feat: include thread name in traces (#3173)

Include the name of the tread that is executing a statement in the traces that are generated, so it is easier to debug concurrency issues. Sometimes, it is not clear whether a statement is being executed by a background worker thread or the main thread.

.readme-partials.yaml

@@ -153,8 +153,9 @@ custom_content: |
   
   #### OpenTelemetry SQL Statement Tracing
   The OpenTelemetry traces that are generated by the Java client include any request and transaction
-  tags that have been set. The traces can also include the SQL statements that are executed. Enable
-  this with the `enableExtendedTracing` option:
+  tags that have been set. The traces can also include the SQL statements that are executed and the
+  name of the thread that executes the statement. Enable this with the `enableExtendedTracing`
+  option:
   
   ```
   SpannerOptions options = SpannerOptions.newBuilder()

README.md

@@ -259,8 +259,9 @@ Spanner spanner = options.getService();
 
 #### OpenTelemetry SQL Statement Tracing
 The OpenTelemetry traces that are generated by the Java client include any request and transaction
-tags that have been set. The traces can also include the SQL statements that are executed. Enable
-this with the `enableExtendedTracing` option:
+tags that have been set. The traces can also include the SQL statements that are executed and the
+name of the thread that executes the statement. Enable this with the `enableExtendedTracing`
+option:
 
 ```
 SpannerOptions options = SpannerOptions.newBuilder()

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner;
+
+import com.google.api.core.InternalApi;
+import io.opentelemetry.context.ContextKey;
+
+/**
+ * Keys for OpenTelemetry context variables that are used by the Spanner client library. Only
+ * intended for internal use.
+ */
+@InternalApi
+public class OpenTelemetryContextKeys {
+  @InternalApi
+  public static final ContextKey<String> THREAD_NAME_KEY = ContextKey.named("thread.name");
+}

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

@@ -1383,6 +1383,7 @@ public Builder setEnableApiTracing(boolean enableApiTracing) {
      *
      * <ul>
      *   <li>db.statement: Contains the SQL statement that is being executed.
+     *   <li>thread.name: The name of the thread that executes the statement.
      * </ul>
      */
     public Builder setEnableExtendedTracing(boolean enableExtendedTracing) {
@@ -1673,6 +1674,7 @@ public boolean isUseVirtualThreads() {
    *
    * <ul>
    *   <li>db.statement: Contains the SQL statement that is being executed.
+   *   <li>thread.name: The name of the thread that executes the statement.
    * </ul>
    */
   public boolean isEnableExtendedTracing() {

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

@@ -19,6 +19,7 @@
 import com.google.cloud.spanner.Options.TagOption;
 import com.google.cloud.spanner.Options.TransactionOption;
 import com.google.cloud.spanner.SpannerOptions.TracingFramework;
+import com.google.common.base.MoreObjects;
 import io.opencensus.trace.BlankSpan;
 import io.opencensus.trace.Span;
 import io.opencensus.trace.Tracer;
@@ -41,6 +42,7 @@ class TraceWrapper {
       AttributeKey.stringKey("db.statement");
   private static final AttributeKey<List<String>> DB_STATEMENT_ARRAY_KEY =
       AttributeKey.stringArrayKey("db.statement");
+  private static final AttributeKey<String> THREAD_NAME_KEY = AttributeKey.stringKey("thread.name");
 
   private final Tracer openCensusTracer;
   private final io.opentelemetry.api.trace.Tracer openTelemetryTracer;
@@ -154,6 +156,7 @@ Attributes createStatementAttributes(Statement statement, Options options) {
       AttributesBuilder builder = Attributes.builder();
       if (this.enableExtendedTracing) {
         builder.put(DB_STATEMENT_KEY, statement.getSql());
+        builder.put(THREAD_NAME_KEY, getTraceThreadName());
       }
       if (options != null && options.hasTag()) {
         builder.put(STATEMENT_TAG_KEY, options.tag());
@@ -172,6 +175,7 @@ Attributes createStatementBatchAttributes(Iterable<Statement> statements, Option
             StreamSupport.stream(statements.spliterator(), false)
                 .map(Statement::getSql)
                 .collect(Collectors.toList()));
+        builder.put(THREAD_NAME_KEY, getTraceThreadName());
       }
       if (options != null && options.hasTag()) {
         builder.put(STATEMENT_TAG_KEY, options.tag());
@@ -180,4 +184,10 @@ Attributes createStatementBatchAttributes(Iterable<Statement> statements, Option
     }
     return Attributes.empty();
   }
+
+  private static String getTraceThreadName() {
+    return MoreObjects.firstNonNull(
+        Context.current().get(OpenTelemetryContextKeys.THREAD_NAME_KEY),
+        Thread.currentThread().getName());
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java

@@ -25,6 +25,7 @@
 import com.google.cloud.spanner.BatchTransactionId;
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.OpenTelemetryContextKeys;
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.Partition;
@@ -47,6 +48,7 @@
 import io.grpc.Status;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Scope;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
@@ -358,40 +360,47 @@ public <ReqT, RespT> ApiCallContext configure(
                 }
               });
     }
-    ApiFuture<T> f = statementExecutor.submit(context.wrap(callable));
-    final SpannerAsyncExecutionException caller =
-        callType == CallType.ASYNC
-            ? new SpannerAsyncExecutionException(statement.getStatement())
-            : null;
-    final ApiFuture<T> future =
-        ApiFutures.catching(
-            f,
-            Throwable.class,
-            input -> {
-              if (caller != null) {
-                input.addSuppressed(caller);
+    // Register the name of the thread that called this method as the thread name that should be
+    // traced.
+    try (Scope ignore =
+        io.opentelemetry.context.Context.current()
+            .with(OpenTelemetryContextKeys.THREAD_NAME_KEY, Thread.currentThread().getName())
+            .makeCurrent()) {
+      ApiFuture<T> f = statementExecutor.submit(context.wrap(callable));
+      final SpannerAsyncExecutionException caller =
+          callType == CallType.ASYNC
+              ? new SpannerAsyncExecutionException(statement.getStatement())
+              : null;
+      final ApiFuture<T> future =
+          ApiFutures.catching(
+              f,
+              Throwable.class,
+              input -> {
+                if (caller != null) {
+                  input.addSuppressed(caller);
+                }
+                throw SpannerExceptionFactory.asSpannerException(input);
+              },
+              MoreExecutors.directExecutor());
+      synchronized (this) {
+        this.currentlyRunningStatementFuture = future;
+      }
+      future.addListener(
+          new Runnable() {
+            @Override
+            public void run() {
+              synchronized (this) {
+                if (currentlyRunningStatementFuture == future) {
+                  currentlyRunningStatementFuture = null;
+                }
               }
-              throw SpannerExceptionFactory.asSpannerException(input);
-            },
-            MoreExecutors.directExecutor());
-    synchronized (this) {
-      this.currentlyRunningStatementFuture = future;
-    }
-    future.addListener(
-        new Runnable() {
-          @Override
-          public void run() {
-            synchronized (this) {
-              if (currentlyRunningStatementFuture == future) {
-                currentlyRunningStatementFuture = null;
+              if (isSingleUse()) {
+                endUnitOfWorkSpan();
               }
             }
-            if (isSingleUse()) {
-              endUnitOfWorkSpan();
-            }
-          }
-        },
-        MoreExecutors.directExecutor());
-    return future;
+          },
+          MoreExecutors.directExecutor());
+      return future;
+    }
   }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/OpenTelemetryTracingTest.java

@@ -164,6 +164,14 @@ public boolean isEnableExtendedTracing() {
         "CloudSpannerOperation.ExecuteStreamingQuery",
         Attributes.of(AttributeKey.stringKey("db.statement"), SELECT1_STATEMENT.getSql()),
         spans);
+    SpanData executeQuerySpan =
+        getSpan(
+            "CloudSpannerOperation.ExecuteStreamingQuery",
+            Attributes.of(
+                AttributeKey.stringKey("db.statement"), SELECT1_STATEMENT.getSql(),
+                AttributeKey.stringKey("thread.name"), Thread.currentThread().getName()),
+            spans);
+
     assertParent(
         "CloudSpannerJdbc.SingleUseTransaction", "CloudSpanner.ReadOnlyTransaction", spans);
     assertParent(
@@ -190,6 +198,14 @@ public void testSingleUseQuery() {
         "CloudSpannerOperation.ExecuteStreamingQuery",
         Attributes.of(AttributeKey.stringKey("db.statement"), SELECT1_STATEMENT.getSql()),
         spans);
+    SpanData executeQuerySpan =
+        getSpan(
+            "CloudSpannerOperation.ExecuteStreamingQuery",
+            Attributes.of(
+                AttributeKey.stringKey("db.statement"), SELECT1_STATEMENT.getSql(),
+                AttributeKey.stringKey("thread.name"), Thread.currentThread().getName()),
+            spans);
+
     assertParent(
         "CloudSpannerJdbc.SingleUseTransaction", "CloudSpanner.ReadOnlyTransaction", spans);
     assertParent(
@@ -222,6 +238,13 @@ public void testSingleUseUpdate() {
         "CloudSpannerOperation.ExecuteUpdate",
         Attributes.of(AttributeKey.stringKey("db.statement"), INSERT_STATEMENT.getSql()),
         spans);
+    SpanData executeQuerySpan =
+        getSpan(
+            "CloudSpannerOperation.ExecuteUpdate",
+            Attributes.of(
+                AttributeKey.stringKey("db.statement"), INSERT_STATEMENT.getSql(),
+                AttributeKey.stringKey("thread.name"), Thread.currentThread().getName()),
+            spans);
     assertParent("CloudSpanner.ReadWriteTransaction", "CloudSpannerOperation.Commit", spans);
   }
 
@@ -244,6 +267,15 @@ public void testSingleUseBatchUpdate() {
             AttributeKey.stringArrayKey("db.statement"),
             ImmutableList.of(INSERT_STATEMENT.getSql(), INSERT_STATEMENT.getSql())),
         spans);
+    SpanData executeQuerySpan =
+        getSpan(
+            "CloudSpannerOperation.BatchUpdate",
+            Attributes.of(
+                AttributeKey.stringArrayKey("db.statement"),
+                ImmutableList.of(INSERT_STATEMENT.getSql(), INSERT_STATEMENT.getSql())),
+            spans);
+    String threadName = executeQuerySpan.getAttributes().get(AttributeKey.stringKey("thread.name"));
+    assertEquals(Thread.currentThread().getName(), threadName);
     assertContains("CloudSpannerOperation.Commit", spans);
 
     assertParent(