diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index c5599a749f6..ebba9fa9e17 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/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 rollbackAsync() { ignored -> ApiFutures.immediateFuture(null), MoreExecutors.directExecutor()); } finally { + if (span != null) { + span.addAnnotation("Transaction rolled back"); + span.end(); + } txnState = TransactionState.ROLLED_BACK; } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryApiTracerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryApiTracerTest.java index 65bb5f5f0d7..67012ed9622 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryApiTracerTest.java +++ b/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 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 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 spans) { assertTrue( "Expected " + spansToString(spans) + " to contain " + expected,