feat: Add opt-in flag and ClientInterceptor to propagate trace context for Spanner end to end tracing

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

@@ -158,6 +158,7 @@ public class SpannerOptions extends ServiceOptions<Spanner, SpannerOptions> {
   private final OpenTelemetry openTelemetry;
   private final boolean enableApiTracing;
   private final boolean enableExtendedTracing;
+  private final boolean enableServerSideTracing;
 
   enum TracingFramework {
     OPEN_CENSUS,
@@ -664,6 +665,7 @@ protected SpannerOptions(Builder builder) {
     openTelemetry = builder.openTelemetry;
     enableApiTracing = builder.enableApiTracing;
     enableExtendedTracing = builder.enableExtendedTracing;
+    enableServerSideTracing = builder.enableServerSideTracing;
   }
 
   /**
@@ -696,6 +698,10 @@ default boolean isEnableExtendedTracing() {
     default boolean isEnableApiTracing() {
       return false;
     }
+
+    default boolean isEnableServerSideTracing() {
+      return false;
+    }
   }
 
   /**
@@ -709,6 +715,7 @@ private static class SpannerEnvironmentImpl implements SpannerEnvironment {
         "SPANNER_OPTIMIZER_STATISTICS_PACKAGE";
     private static final String SPANNER_ENABLE_EXTENDED_TRACING = "SPANNER_ENABLE_EXTENDED_TRACING";
     private static final String SPANNER_ENABLE_API_TRACING = "SPANNER_ENABLE_API_TRACING";
+    private static final String SPANNER_ENABLE_SERVER_SIDE_TRACING = "SPANNER_ENABLE_SERVER_SIDE_TRACING";
 
     private SpannerEnvironmentImpl() {}
 
@@ -734,6 +741,11 @@ public boolean isEnableExtendedTracing() {
     public boolean isEnableApiTracing() {
       return Boolean.parseBoolean(System.getenv(SPANNER_ENABLE_API_TRACING));
     }
+
+    @Override
+    public boolean isEnableServerSideTracing() {
+      return Boolean.parseBoolean(System.getenv(SPANNER_ENABLE_SERVER_SIDE_TRACING));
+    }
   }
 
   /** Builder for {@link SpannerOptions} instances. */
@@ -799,6 +811,7 @@ public static class Builder
     private OpenTelemetry openTelemetry;
     private boolean enableApiTracing = SpannerOptions.environment.isEnableApiTracing();
     private boolean enableExtendedTracing = SpannerOptions.environment.isEnableExtendedTracing();
+    private boolean enableServerSideTracing = SpannerOptions.environment.isEnableServerSideTracing();
 
     private static String createCustomClientLibToken(String token) {
       return token + " " + ServiceOptions.getGoogApiClientLibName();
@@ -864,6 +877,7 @@ protected Builder() {
       this.useVirtualThreads = options.useVirtualThreads;
       this.enableApiTracing = options.enableApiTracing;
       this.enableExtendedTracing = options.enableExtendedTracing;
+      this.enableServerSideTracing = options.enableServerSideTracing;
     }
 
     @Override
@@ -1391,6 +1405,24 @@ public Builder setEnableExtendedTracing(boolean enableExtendedTracing) {
       return this;
     }
 
+    /**
+     * Enable spanner server side tracing. Enabling this option will create the trace
+     * spans at the Spanner layer. By default, server side tracing is disabled.
+     */
+    public Builder enableServerSideTracing() {
+      this.enableServerSideTracing = true;
+      return this;
+    }
+
+    /**
+     * Disable spanner server side tracing. If server side is disabled, trace
+     * spans won't be created at the Spanner layer.
+     */
+    public Builder disableServerSideTracing() {
+      this.enableServerSideTracing = false;
+      return this;
+    }
+
     @SuppressWarnings("rawtypes")
     @Override
     public SpannerOptions build() {
@@ -1480,7 +1512,8 @@ public static void enableOpenCensusTraces() {
    */
   @ObsoleteApi(
       "The OpenCensus project is deprecated. Use enableOpenTelemetryTraces to switch to OpenTelemetry traces")
-  static void resetActiveTracingFramework() {
+  @VisibleForTesting
+  public static void resetActiveTracingFramework() {
     activeTracingFramework = null;
   }
 
@@ -1681,6 +1714,14 @@ public boolean isEnableExtendedTracing() {
     return enableExtendedTracing;
   }
 
+  /**
+   * Returns whether Spanner server side tracing is enabled. If this option is enabled then trace spans
+   * will be created at the Spanner layer.
+   */
+  public boolean isServerSideTracingEnabled() {
+    return enableServerSideTracing;
+  }
+
   /** Returns the default query options to use for the specific database. */
   public QueryOptions getDefaultQueryOptions(DatabaseId databaseId) {
     // Use the specific query options for the database if any have been specified. These have

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -268,6 +268,7 @@ public class GapicSpannerRpc implements SpannerRpc {
   private static final ConcurrentMap<String, RateLimiter> ADMINISTRATIVE_REQUESTS_RATE_LIMITERS =
       new ConcurrentHashMap<>();
   private final boolean leaderAwareRoutingEnabled;
+  private final boolean serverSideTracingEnabled;
   private final int numChannels;
   private final boolean isGrpcGcpExtensionEnabled;
 
@@ -321,6 +322,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
     this.callCredentialsProvider = options.getCallCredentialsProvider();
     this.compressorName = options.getCompressorName();
     this.leaderAwareRoutingEnabled = options.isLeaderAwareRoutingEnabled();
+    this.serverSideTracingEnabled = options.isServerSideTracingEnabled();
     this.numChannels = options.getNumChannels();
     this.isGrpcGcpExtensionEnabled = options.isGrpcGcpExtensionEnabled();
 
@@ -346,6 +348,8 @@ public GapicSpannerRpc(final SpannerOptions options) {
                           MoreObjects.firstNonNull(
                               options.getInterceptorProvider(),
                               SpannerInterceptorProvider.createDefault(options.getOpenTelemetry())))
+                      // This sets the trace context headers.
+                      .withTraceContext(serverSideTracingEnabled, options.getOpenTelemetry())
                       // This sets the response compressor (Server -> Client).
                       .withEncoding(compressorName))
               .setHeaderProvider(headerProviderWithUserAgent)
@@ -1988,6 +1992,9 @@ <ReqT, RespT> GrpcCallContext newCallContext(
     if (routeToLeader && leaderAwareRoutingEnabled) {
       context = context.withExtraHeaders(metadataProvider.newRouteToLeaderHeader());
     }
+    if (serverSideTracingEnabled) {
+      context = context.withExtraHeaders(metadataProvider.newServerSideTracingHeader());
+    }
     if (callCredentialsProvider != null) {
       CallCredentials callCredentials = callCredentialsProvider.getCallCredentials();
       if (callCredentials != null) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java

@@ -74,6 +74,13 @@ SpannerInterceptorProvider withEncoding(String encoding) {
     return this;
   }
 
+  SpannerInterceptorProvider withTraceContext(boolean serverSideTracingEnabled, OpenTelemetry openTelemetry) {
+    if (serverSideTracingEnabled) {
+      return with(new TraceContextInterceptor(openTelemetry));
+    }
+    return this;
+  }
+
   @Override
   public List<ClientInterceptor> getInterceptors() {
     return clientInterceptors;

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java

@@ -37,13 +37,16 @@ class SpannerMetadataProvider {
   private final Map<Metadata.Key<String>, String> headers;
   private final String resourceHeaderKey;
   private static final String ROUTE_TO_LEADER_HEADER_KEY = "x-goog-spanner-route-to-leader";
+  private static final String SERVER_SIDE_TRACING_HEADER_KEY = "x-goog-spanner-end-to-end-tracing";
   private static final Pattern[] RESOURCE_TOKEN_PATTERNS = {
     Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*/databases/[^/]*)(.*)?"),
     Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*)(.*)?")
   };
 
   private static final Map<String, List<String>> ROUTE_TO_LEADER_HEADER_MAP =
       ImmutableMap.of(ROUTE_TO_LEADER_HEADER_KEY, Collections.singletonList("true"));
+  private static final Map<String, List<String>> SERVER_SIDE_TRACING_HEADER_MAP =
+      ImmutableMap.of(SERVER_SIDE_TRACING_HEADER_KEY, Collections.singletonList("true"));
 
   private SpannerMetadataProvider(Map<String, String> headers, String resourceHeaderKey) {
     this.resourceHeaderKey = resourceHeaderKey;
@@ -89,6 +92,10 @@ Map<String, List<String>> newRouteToLeaderHeader() {
     return ROUTE_TO_LEADER_HEADER_MAP;
   }
 
+  Map<String, List<String>> newServerSideTracingHeader() {
+    return SERVER_SIDE_TRACING_HEADER_MAP;
+  }
+
   private Map<Metadata.Key<String>, String> constructHeadersAsMetadata(
       Map<String, String> headers) {
     ImmutableMap.Builder<Metadata.Key<String>, String> headersAsMetadataBuilder =

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/TraceContextInterceptor.java

@@ -0,0 +1,73 @@
+/*
+ * 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
+ *
+ *      https://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.spi.v1;
+
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.ClientInterceptor;
+import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
+import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener;
+import io.grpc.Metadata;
+import io.grpc.MethodDescriptor;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapPropagator;
+import io.opentelemetry.context.propagation.TextMapSetter;
+
+/**
+ * Intercepts all gRPC calls and injects trace context related headers to propagate trace context to
+ * Spanner. This class takes reference from OpenTelemetry's JAVA instrumentation library for gRPC.
+ * https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/9ecf7965aa455d41ea8cc0761b6c6b6eeb106324/instrumentation/grpc-1.6/library/src/main/java/io/opentelemetry/instrumentation/grpc/v1_6/TracingClientInterceptor.java#L27
+ */
+class TraceContextInterceptor implements ClientInterceptor {
+
+  private final TextMapPropagator textMapPropagator;
+
+  TraceContextInterceptor(OpenTelemetry openTelemetry) {
+    this.textMapPropagator = openTelemetry.getPropagators().getTextMapPropagator();
+  }
+
+  enum MetadataSetter implements TextMapSetter<Metadata> {
+    INSTANCE;
+
+    @SuppressWarnings("null")
+    @Override
+    public void set(Metadata carrier, String key, String value) {
+      carrier.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+    }
+  }
+
+  private static final class NoopSimpleForwardingClientCallListener<RespT>
+      extends SimpleForwardingClientCallListener<RespT> {
+    public NoopSimpleForwardingClientCallListener(ClientCall.Listener<RespT> responseListener) {
+      super(responseListener);
+    }
+  }
+
+  @Override
+  public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
+      MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
+    return new SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
+      @Override
+      public void start(Listener<RespT> responseListener, Metadata headers) {
+        Context parentContext = Context.current();
+        textMapPropagator.inject(parentContext, headers, MetadataSetter.INSTANCE);
+        super.start(new NoopSimpleForwardingClientCallListener<RespT>(responseListener), headers);
+      }
+    };
+  }
+}

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

@@ -731,6 +731,24 @@ public void testLeaderAwareRoutingEnablement() {
             .isLeaderAwareRoutingEnabled());
   }
 
+  @Test
+  public void testServerSideTracingEnablement() {
+    // Test that end to end tracing is disabled by default.
+    assertFalse(SpannerOptions.newBuilder().setProjectId("p").build().isServerSideTracingEnabled());
+    assertTrue(
+        SpannerOptions.newBuilder()
+            .setProjectId("p")
+            .enableServerSideTracing()
+            .build()
+            .isServerSideTracingEnabled());
+    assertFalse(
+        SpannerOptions.newBuilder()
+            .setProjectId("p")
+            .disableServerSideTracing()
+            .build()
+            .isServerSideTracingEnabled());
+  }
+
   @Test
   public void testSetDirectedReadOptions() {
     final DirectedReadOptions directedReadOptions =