fix: add missing span.end calls for AsyncTransactionManager (#4012)

The AsyncTransactionManager did not end the span when the transaction was
committed or rolled back. This caused the spans not to be collected and
exported.

Author: olavloite

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java

@@ -165,12 +165,19 @@ public void onFailure(Throwable t) {
               txnState = TransactionState.ABORTED;
             } else {
               txnState = TransactionState.COMMIT_FAILED;
+              if (span != null) {
+                span.setStatus(t);
+                span.end();
+              }
               commitResponse.setException(t);
             }
           }
 
           @Override
           public void onSuccess(CommitResponse result) {
+            if (span != null) {
+              span.end();
+            }
             commitResponse.set(result);
           }
         },
@@ -190,6 +197,10 @@ public ApiFuture<Void> rollbackAsync() {
           ignored -> ApiFutures.immediateFuture(null),
           MoreExecutors.directExecutor());
     } finally {
+      if (span != null) {
+        span.addAnnotation("Transaction rolled back");
+        span.end();
+      }
       txnState = TransactionState.ROLLED_BACK;
     }
   }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryApiTracerTest.java

@@ -27,11 +27,14 @@
 import com.google.api.gax.longrunning.OperationTimedPollAlgorithm;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.NoCredentials;
+import com.google.cloud.spanner.AsyncTransactionManager.CommitTimestampFuture;
+import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture;
 import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
 import com.google.cloud.spanner.SpannerOptions.SpannerEnvironment;
 import com.google.cloud.spanner.connection.RandomResultSetGenerator;
 import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
 import io.grpc.Status;
 import io.opentelemetry.api.GlobalOpenTelemetry;
@@ -451,6 +454,58 @@ public boolean isEnableApiTracing() {
         "CloudSpannerOperation.ExecuteStreamingQuery", "Spanner.ExecuteStreamingSql", spans);
   }
 
+  @Test
+  public void testAsyncTransactionManagerCommit() throws Exception {
+    try (AsyncTransactionManager manager = client.transactionManagerAsync()) {
+      TransactionContextFuture transactionFuture = manager.beginAsync();
+      CommitTimestampFuture commitTimestamp =
+          transactionFuture
+              .then(
+                  (transaction, __) -> transaction.executeUpdateAsync(UPDATE_RANDOM),
+                  MoreExecutors.directExecutor())
+              .commitAsync();
+      commitTimestamp.get();
+    }
+
+    assertEquals(CompletableResultCode.ofSuccess(), spanExporter.flush());
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertContains("CloudSpanner.ReadWriteTransaction", spans);
+    assertContains("CloudSpannerOperation.ExecuteUpdate", spans);
+    assertContains("CloudSpannerOperation.Commit", spans);
+    assertContains("Spanner.ExecuteSql", spans);
+    assertContains("Spanner.Commit", spans);
+
+    assertParent("CloudSpanner.ReadWriteTransaction", "CloudSpannerOperation.ExecuteUpdate", spans);
+    assertParent("CloudSpanner.ReadWriteTransaction", "CloudSpannerOperation.Commit", spans);
+    assertParent("CloudSpannerOperation.ExecuteUpdate", "Spanner.ExecuteSql", spans);
+  }
+
+  @Test
+  public void testAsyncTransactionManagerRollback() throws Exception {
+    try (AsyncTransactionManager manager = client.transactionManagerAsync()) {
+      TransactionContextFuture transactionFuture = manager.beginAsync();
+      transactionFuture
+          .then(
+              (transaction, __) -> transaction.executeUpdateAsync(UPDATE_RANDOM),
+              MoreExecutors.directExecutor())
+          .get();
+      manager.rollbackAsync().get();
+    }
+
+    assertEquals(CompletableResultCode.ofSuccess(), spanExporter.flush());
+    List<SpanData> spans = spanExporter.getFinishedSpanItems();
+    assertContains("CloudSpanner.ReadWriteTransaction", spans);
+    assertContains("CloudSpannerOperation.ExecuteUpdate", spans);
+    assertContains("Spanner.ExecuteSql", spans);
+    assertContains("Spanner.Rollback", spans);
+
+    assertParent("CloudSpanner.ReadWriteTransaction", "CloudSpannerOperation.ExecuteUpdate", spans);
+    assertParent("CloudSpannerOperation.ExecuteUpdate", "Spanner.ExecuteSql", spans);
+    SpanData transactionSpan = getSpan("CloudSpanner.ReadWriteTransaction", spans);
+    assertNotNull(transactionSpan);
+    assertContainsEvent("Transaction rolled back", transactionSpan.getEvents());
+  }
+
   void assertContains(String expected, List<SpanData> spans) {
     assertTrue(
         "Expected " + spansToString(spans) + " to contain " + expected,