fix: remove unused DmlBatch span (#3147)

olavloite · web-flow · commit f7891c1ca427 · 2024-06-21T09:08:12.000+05:30
The DmlBatch span that was created for DML batches was never used,
as all traces were registered on the underlying transaction of the
batch. The DmlBatch span was therefore always just an empty
sibling span of the Transaction span, which did contain the traces
of the statements that were executed.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java
@@ -165,6 +165,11 @@ B setSpan(@Nullable Span span) {
     this.span = Preconditions.checkNotNull(builder.span);
   }
 
+  @Override
+  public Span getSpan() {
+    return this.span;
+  }
+
   ApiFuture<Void> asyncEndUnitOfWorkSpan() {
     return this.statementExecutor.submit(this::endUnitOfWorkSpan);
   }
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java
@@ -96,7 +96,6 @@ class ConnectionImpl implements Connection {
   private static final String SINGLE_USE_TRANSACTION = "SingleUseTransaction";
   private static final String READ_ONLY_TRANSACTION = "ReadOnlyTransaction";
   private static final String READ_WRITE_TRANSACTION = "ReadWriteTransaction";
-  private static final String DML_BATCH = "DmlBatch";
   private static final String DDL_BATCH = "DdlBatch";
   private static final String DDL_STATEMENT = "DdlStatement";
 
@@ -1932,7 +1931,8 @@ UnitOfWork createNewUnitOfWork(
               .setStatementTag(statementTag)
               .setExcludeTxnFromChangeStreams(excludeTxnFromChangeStreams)
               .setRpcPriority(rpcPriority)
-              .setSpan(createSpanForUnitOfWork(DML_BATCH))
+              // Use the transaction Span for the DML batch.
+              .setSpan(transactionStack.peek().getSpan())
               .build();
         case DDL_BATCH:
           return DdlBatch.newBuilder()
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java
@@ -31,6 +31,7 @@
 import com.google.cloud.spanner.TransactionContext;
 import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
 import com.google.spanner.v1.ResultSetStats;
+import io.opentelemetry.api.trace.Span;
 import java.util.concurrent.ExecutionException;
 import javax.annotation.Nonnull;
 
@@ -77,6 +78,9 @@ public boolean isActive() {
   /** @return <code>true</code> if this unit of work is still active. */
   boolean isActive();
 
+  /** @return the {@link Span} that is used by this {@link UnitOfWork}. */
+  Span getSpan();
+
   /** Returns true if this transaction can only be used for a single statement. */
   boolean isSingleUse();
 

Original file line number	Diff line number	Diff line change
`@@ -165,6 +165,11 @@ B setSpan(@Nullable Span span) {`
`165`	`165`	`this.span = Preconditions.checkNotNull(builder.span);`
`166`	`166`	`}`
`167`	`167`
	`168`	`+ @Override`
	`169`	`+ public Span getSpan() {`
	`170`	`+ return this.span;`
	`171`	`+ }`
	`172`	`+`
`168`	`173`	`ApiFuture<Void> asyncEndUnitOfWorkSpan() {`
`169`	`174`	`return this.statementExecutor.submit(this::endUnitOfWorkSpan);`
`170`	`175`	`}`