fix: do not end transaction span when rolling back to savepoint (#3167)

olavloite · gcf-owl-bot[bot] · web-flow · commit 8ec0cf2032de · 2024-06-21T12:53:31.000+02:00
* fix: do not end transaction span when rolling back to savepoint When a transaction is rolled back to a savepoint instead of rolled back entirely, the transaction span should not be ended, as the transaction continues to live. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/README.md b/README.md
@@ -50,7 +50,7 @@ If you are using Maven without the BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies:
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:26.41.0')
+implementation platform('com.google.cloud:libraries-bom:26.42.0')
 
 implementation 'com.google.cloud:google-cloud-spanner'
 ```
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java
@@ -1142,22 +1142,29 @@ public ApiFuture<Void> rollbackAsync(CallType callType) {
     }
   }
 
-  private ApiFuture<Void> rollbackAsync(CallType callType, boolean updateStatus) {
+  private ApiFuture<Void> rollbackAsync(CallType callType, boolean updateStatusAndEndSpan) {
     ConnectionPreconditions.checkState(
         state == UnitOfWorkState.STARTED || state == UnitOfWorkState.ABORTED,
         "This transaction has status " + state.name());
-    if (updateStatus) {
+    if (updateStatusAndEndSpan) {
       state = UnitOfWorkState.ROLLED_BACK;
-      asyncEndUnitOfWorkSpan();
     }
     if (txContextFuture != null && state != UnitOfWorkState.ABORTED) {
       ApiFuture<Void> result =
           executeStatementAsync(
               callType, ROLLBACK_STATEMENT, rollbackCallable, SpannerGrpc.getRollbackMethod());
-      asyncEndUnitOfWorkSpan();
+      if (updateStatusAndEndSpan) {
+        // Note: We end the transaction span after executing the rollback to include the rollback in
+        // the transaction span. Even though both methods are executed asynchronously, they are both
+        // executed using the same single-threaded executor, meaning that the span will only be
+        // ended after the rollback has finished.
+        asyncEndUnitOfWorkSpan();
+      }
       return result;
-    } else {
+    } else if (updateStatusAndEndSpan) {
       return asyncEndUnitOfWorkSpan();
+    } else {
+      return ApiFutures.immediateFuture(null);
     }
   }
 
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/OpenTelemetryTracingTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/OpenTelemetryTracingTest.java
@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import com.google.cloud.spanner.MockSpannerServiceImpl;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.SpannerOptions.SpannerEnvironment;
@@ -472,6 +473,59 @@ public void testMultiUseReadWriteAborted() {
     assertParent("CloudSpanner.ReadWriteTransaction", "CloudSpannerOperation.Commit", spans);
   }
 
+  @Test
+  public void testSavepoint() {
+    Statement statement1 = Statement.of("insert into foo (id) values (1)");
+    Statement statement2 = Statement.of("insert into foo (id) values (2)");
+    mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement1, 1));
+    mockSpanner.putStatementResult(MockSpannerServiceImpl.StatementResult.update(statement2, 1));
+
+    try (Connection connection = createTestConnection()) {
+      connection.setAutocommit(false);
+      connection.setReadOnly(false);
+      connection.setSavepointSupport(SavepointSupport.ENABLED);
+      assertEquals(1L, connection.executeUpdate(statement1));
+      connection.savepoint("test");
+      assertEquals(1L, connection.executeUpdate(statement2));
+      connection.rollbackToSavepoint("test");
+      connection.commit();
+    }
+    assertEquals(CompletableResultCode.ofSuccess(), spanExporter.flush());
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertContains("CloudSpannerJdbc.ReadWriteTransaction", spans);
+    assertContains("CloudSpanner.ReadWriteTransaction", spans);
+    // Statement 1 is executed 2 times, because the original transaction needs to be
+    // retried after the transaction was rolled back to the savepoint.
+    assertContains(
+        "CloudSpannerOperation.ExecuteUpdate",
+        2,
+        Attributes.of(AttributeKey.stringKey("db.statement"), statement1.getSql()),
+        spans);
+    assertContains(
+        "CloudSpannerOperation.ExecuteUpdate",
+        1,
+        Attributes.of(AttributeKey.stringKey("db.statement"), statement2.getSql()),
+        spans);
+    assertContains("CloudSpannerOperation.Commit", spans);
+
+    // Verify that we have two Cloud Spanner transactions, and that these are both children of one
+    // JDBC transaction.
+    List<SpanData> transactionSpans =
+        getSpans("CloudSpanner.ReadWriteTransaction", Attributes.empty(), spans);
+    assertEquals(2, transactionSpans.size());
+    assertEquals(
+        transactionSpans.get(0).getParentSpanId(), transactionSpans.get(1).getParentSpanId());
+    List<SpanData> jdbcTransactionSpans =
+        getSpans("CloudSpannerJdbc.ReadWriteTransaction", Attributes.empty(), spans);
+    assertEquals(1, jdbcTransactionSpans.size());
+    assertEquals(
+        jdbcTransactionSpans.get(0).getSpanId(), transactionSpans.get(0).getParentSpanId());
+    List<SpanData> commitSpans =
+        getSpans("CloudSpannerOperation.Commit", Attributes.empty(), spans);
+    assertEquals(1, commitSpans.size());
+    assertEquals(transactionSpans.get(1).getSpanId(), commitSpans.get(0).getParentSpanId());
+  }
+
   @Test
   public void testTransactionTag() {
     try (Connection connection = createTestConnection()) {
