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

* feat(spanner): Add x-goog-spanner-end-to-end-tracing header for requests to SpanFE

* Add Grpc Telemetry client interceptor for trace context propagation

* Remove print statement

* copy grpc telemetry client interceptor code

* rename spanner option for end to end tracing

* add test for custom client interceptor code

* resolve comments

* Make trace context interceptor conditional based on client opt-in

* resolve comments

* reformat code

* resolve comments

* fix clirr build failure

* fix lint errors

* combine enable/disable fn into single fn

* Rename server side tracing to spanner tracing

* Rename spanner tracing to end to end tracing

* fix comments

* Fix lint error

Author: nareshz

google-cloud-spanner/clirr-ignored-differences.xml

@@ -695,6 +695,13 @@
     <method>boolean isEnableApiTracing()</method>
   </difference>
 
+  <!-- Added end to end tracing option -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/SpannerOptions$SpannerEnvironment</className>
+    <method>boolean isEnableEndToEndTracing()</method>
+  </difference>
+
   <!-- Added Built In Metrics option -->
   <difference>
     <differenceType>7012</differenceType>

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

@@ -163,6 +163,7 @@ public class SpannerOptions extends ServiceOptions<Spanner, SpannerOptions> {
   private final boolean enableApiTracing;
   private final boolean enableBuiltInMetrics;
   private final boolean enableExtendedTracing;
+  private final boolean enableEndToEndTracing;
 
   enum TracingFramework {
     OPEN_CENSUS,
@@ -670,6 +671,7 @@ protected SpannerOptions(Builder builder) {
     enableApiTracing = builder.enableApiTracing;
     enableExtendedTracing = builder.enableExtendedTracing;
     enableBuiltInMetrics = builder.enableBuiltInMetrics;
+    enableEndToEndTracing = builder.enableEndToEndTracing;
   }
 
   /**
@@ -706,6 +708,10 @@ default boolean isEnableApiTracing() {
     default boolean isEnableBuiltInMetrics() {
       return false;
     }
+
+    default boolean isEnableEndToEndTracing() {
+      return false;
+    }
   }
 
   /**
@@ -720,6 +726,8 @@ private static class SpannerEnvironmentImpl implements SpannerEnvironment {
     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_BUILTIN_METRICS = "SPANNER_ENABLE_BUILTIN_METRICS";
+    private static final String SPANNER_ENABLE_END_TO_END_TRACING =
+        "SPANNER_ENABLE_END_TO_END_TRACING";
 
     private SpannerEnvironmentImpl() {}
 
@@ -752,6 +760,11 @@ public boolean isEnableBuiltInMetrics() {
       // removed in the future.
       return Boolean.parseBoolean(System.getenv(SPANNER_ENABLE_BUILTIN_METRICS));
     }
+
+    @Override
+    public boolean isEnableEndToEndTracing() {
+      return Boolean.parseBoolean(System.getenv(SPANNER_ENABLE_END_TO_END_TRACING));
+    }
   }
 
   /** Builder for {@link SpannerOptions} instances. */
@@ -816,6 +829,7 @@ public static class Builder
     private boolean enableApiTracing = SpannerOptions.environment.isEnableApiTracing();
     private boolean enableExtendedTracing = SpannerOptions.environment.isEnableExtendedTracing();
     private boolean enableBuiltInMetrics = SpannerOptions.environment.isEnableBuiltInMetrics();
+    private boolean enableEndToEndTracing = SpannerOptions.environment.isEnableEndToEndTracing();
 
     private static String createCustomClientLibToken(String token) {
       return token + " " + ServiceOptions.getGoogApiClientLibName();
@@ -882,6 +896,7 @@ protected Builder() {
       this.enableApiTracing = options.enableApiTracing;
       this.enableExtendedTracing = options.enableExtendedTracing;
       this.enableBuiltInMetrics = options.enableBuiltInMetrics;
+      this.enableEndToEndTracing = options.enableEndToEndTracing;
     }
 
     @Override
@@ -1415,6 +1430,17 @@ public Builder setEnableExtendedTracing(boolean enableExtendedTracing) {
       return this;
     }
 
+    /**
+     * Sets whether to enable end to end tracing. Enabling this option will create the trace spans
+     * at the Spanner layer. By default, end to end tracing is disabled. Enabling end to end tracing
+     * requires OpenTelemetry to be set up. Simply enabling this option won't generate traces at
+     * Spanner layer.
+     */
+    public Builder setEnableEndToEndTracing(boolean enableEndToEndTracing) {
+      this.enableEndToEndTracing = enableEndToEndTracing;
+      return this;
+    }
+
     @SuppressWarnings("rawtypes")
     @Override
     public SpannerOptions build() {
@@ -1504,6 +1530,7 @@ public static void enableOpenCensusTraces() {
    */
   @ObsoleteApi(
       "The OpenCensus project is deprecated. Use enableOpenTelemetryTraces to switch to OpenTelemetry traces")
+  @VisibleForTesting
   static void resetActiveTracingFramework() {
     activeTracingFramework = null;
   }
@@ -1745,6 +1772,14 @@ public boolean isEnableExtendedTracing() {
     return enableExtendedTracing;
   }
 
+  /**
+   * Returns whether end to end tracing is enabled. If this option is enabled then trace spans will
+   * be created at the Spanner layer.
+   */
+  public boolean isEndToEndTracingEnabled() {
+    return enableEndToEndTracing;
+  }
+
   /** 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

@@ -272,6 +272,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 endToEndTracingEnabled;
   private final int numChannels;
   private final boolean isGrpcGcpExtensionEnabled;
 
@@ -325,6 +326,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
     this.callCredentialsProvider = options.getCallCredentialsProvider();
     this.compressorName = options.getCompressorName();
     this.leaderAwareRoutingEnabled = options.isLeaderAwareRoutingEnabled();
+    this.endToEndTracingEnabled = options.isEndToEndTracingEnabled();
     this.numChannels = options.getNumChannels();
     this.isGrpcGcpExtensionEnabled = options.isGrpcGcpExtensionEnabled();
 
@@ -350,6 +352,8 @@ public GapicSpannerRpc(final SpannerOptions options) {
                           MoreObjects.firstNonNull(
                               options.getInterceptorProvider(),
                               SpannerInterceptorProvider.createDefault(options.getOpenTelemetry())))
+                      // This sets the trace context headers.
+                      .withTraceContext(endToEndTracingEnabled, options.getOpenTelemetry())
                       // This sets the response compressor (Server -> Client).
                       .withEncoding(compressorName))
               .setHeaderProvider(headerProviderWithUserAgent)
@@ -2007,6 +2011,9 @@ <ReqT, RespT> GrpcCallContext newCallContext(
     if (routeToLeader && leaderAwareRoutingEnabled) {
       context = context.withExtraHeaders(metadataProvider.newRouteToLeaderHeader());
     }
+    if (endToEndTracingEnabled) {
+      context = context.withExtraHeaders(metadataProvider.newEndToEndTracingHeader());
+    }
     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,14 @@ SpannerInterceptorProvider withEncoding(String encoding) {
     return this;
   }
 
+  SpannerInterceptorProvider withTraceContext(
+      boolean endToEndTracingEnabled, OpenTelemetry openTelemetry) {
+    if (endToEndTracingEnabled) {
+      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 END_TO_END_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>> END_TO_END_TRACING_HEADER_MAP =
+      ImmutableMap.of(END_TO_END_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>> newEndToEndTracingHeader() {
+    return END_TO_END_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/SpannerOptionsHelper.java

@@ -0,0 +1,29 @@
+/*
+ * 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;
+
+/** Helper to configure SpannerOptions for tests. */
+public class SpannerOptionsHelper {
+
+  /**
+   * Resets the activeTracingFramework. This variable is used for internal testing, and is not a
+   * valid production scenario.
+   */
+  public static void resetActiveTracingFramework() {
+    SpannerOptions.resetActiveTracingFramework();
+  }
+}

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

@@ -736,6 +736,24 @@ public void testLeaderAwareRoutingEnablement() {
             .isLeaderAwareRoutingEnabled());
   }
 
+  @Test
+  public void testEndToEndTracingEnablement() {
+    // Test that end to end tracing is disabled by default.
+    assertFalse(SpannerOptions.newBuilder().setProjectId("p").build().isEndToEndTracingEnabled());
+    assertTrue(
+        SpannerOptions.newBuilder()
+            .setProjectId("p")
+            .setEnableEndToEndTracing(true)
+            .build()
+            .isEndToEndTracingEnabled());
+    assertFalse(
+        SpannerOptions.newBuilder()
+            .setProjectId("p")
+            .setEnableEndToEndTracing(false)
+            .build()
+            .isEndToEndTracingEnabled());
+  }
+
   @Test
   public void testSetDirectedReadOptions() {
     final DirectedReadOptions directedReadOptions =