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/spi/v1/SpannerInterceptorProvider.java

@@ -23,6 +23,7 @@
 import io.grpc.ClientInterceptor;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.logging.Level;
@@ -51,6 +52,7 @@ public static SpannerInterceptorProvider createDefault(OpenTelemetry openTelemet
     defaultInterceptorList.add(
         new LoggingInterceptor(Logger.getLogger(GapicSpannerRpc.class.getName()), Level.FINER));
     defaultInterceptorList.add(new HeaderInterceptor(new SpannerRpcMetrics(openTelemetry)));
+    defaultInterceptorList.add(new TraceContextInterceptor(openTelemetry));
     return new SpannerInterceptorProvider(ImmutableList.copyOf(defaultInterceptorList));
   }
 

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

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2021 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.grpc.Status;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.context.propagation.TextMapSetter;
+
+public class TraceContextInterceptor implements ClientInterceptor  {
+
+    private final ContextPropagators propagators;
+
+    TraceContextInterceptor(OpenTelemetry openTelemetry) {
+        this.propagators = openTelemetry.getPropagators();
+    }
+
+    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);
+        }
+    }
+
+    @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();
+
+                propagators.getTextMapPropagator().inject(parentContext, headers, MetadataSetter.INSTANCE);
+
+                super.start(new SimpleForwardingClientCallListener<RespT>(responseListener) {
+                    @Override
+                    public void onHeaders(Metadata headers) {
+                        super.onHeaders(headers);
+                    }
+
+                    @Override
+                    public void onClose(Status status, Metadata trailers) {
+                        super.onClose(status, trailers);
+                    }
+                }, headers);
+            }
+        };
+    }
+}

google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java

@@ -76,6 +76,13 @@
 import io.grpc.auth.MoreCallCredentials;
 import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
 import io.grpc.protobuf.lite.ProtoLiteUtils;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.util.HashMap;
@@ -148,6 +155,7 @@ public class GapicSpannerRpcTest {
   private static String defaultUserAgent;
   private static Spanner spanner;
   private static boolean isRouteToLeader;
+  private static boolean isTraceContextPresent;
 
   @Parameter public Dialect dialect;
 
@@ -158,6 +166,9 @@ public static Object[] data() {
 
   @Before
   public void startServer() throws IOException {
+    // Enable OpenTelemetry tracing.
+    SpannerOptions.enableOpenTelemetryTraces();
+
     assumeTrue(
         "Skip tests when emulator is enabled as this test interferes with the check whether the emulator is running",
         System.getenv("SPANNER_EMULATOR_HOST") == null);
@@ -194,6 +205,9 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
                     if (call.getMethodDescriptor()
                             .equals(SpannerGrpc.getExecuteStreamingSqlMethod())
                         || call.getMethodDescriptor().equals(SpannerGrpc.getExecuteSqlMethod())) {
+                      String traceParentHeader =
+                          headers.get(Key.of("traceparent", Metadata.ASCII_STRING_MARSHALLER));
+                      isTraceContextPresent = (traceParentHeader != null);
                       String routeToLeaderHeader =
                           headers.get(
                               Key.of(
@@ -224,6 +238,7 @@ public void reset() throws InterruptedException {
       server.awaitTermination();
     }
     isRouteToLeader = false;
+    isTraceContextPresent = false;
   }
 
   @Test
@@ -535,6 +550,41 @@ public void testCustomUserAgent() {
     }
   }
 
+  @Test
+  public void testTraceContextHeaderWithOpenTelemetry() {
+    OpenTelemetry openTelemetry = OpenTelemetrySdk.builder()
+      .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+      .setTracerProvider(SdkTracerProvider.builder().setSampler(Sampler.alwaysOn()).build())
+      .build();
+
+    final SpannerOptions options = createSpannerOptions().toBuilder().setOpenTelemetry(openTelemetry).build();
+    try (Spanner spanner = options.getService()) {
+      final DatabaseClient databaseClient =
+        spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+
+        try (final ResultSet rs = databaseClient.singleUse().executeQuery(SELECT1AND2)) {
+          rs.next();
+        }
+
+        assertTrue(isTraceContextPresent);
+    } 
+  }
+
+  @Test
+  public void testTraceContextHeaderWithoutOpenTelemetry() {
+    final SpannerOptions options = createSpannerOptions();
+    try (Spanner spanner = options.getService()) {
+      final DatabaseClient databaseClient =
+        spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+
+        try (final ResultSet rs = databaseClient.singleUse().executeQuery(SELECT1AND2)) {
+          rs.next();
+        }
+
+        assertFalse(isTraceContextPresent);
+    }
+  }
+
   @Test
   public void testRouteToLeaderHeaderForReadOnly() {
     final SpannerOptions options =