resolve comments

Author: nareshz

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

@@ -1480,7 +1480,7 @@ public static void enableOpenCensusTraces() {
    */
   @ObsoleteApi(
       "The OpenCensus project is deprecated. Use enableOpenTelemetryTraces to switch to OpenTelemetry traces")
-  static void resetActiveTracingFramework() {
+  public static void resetActiveTracingFramework() {
     activeTracingFramework = null;
   }
 

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

@@ -23,7 +23,6 @@
 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;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 Google LLC
+ * 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.
@@ -23,52 +23,51 @@
 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.TextMapPropagator;
 import io.opentelemetry.context.propagation.TextMapSetter;
 
-public class TraceContextInterceptor implements ClientInterceptor  {
-
-    private final ContextPropagators propagators;
+/**
+ * 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 {
 
-    TraceContextInterceptor(OpenTelemetry openTelemetry) {
-        this.propagators = openTelemetry.getPropagators();
-    }
+  private final TextMapPropagator textMapPropagator;
 
-    enum MetadataSetter implements TextMapSetter<Metadata> {
-        INSTANCE;
+  TraceContextInterceptor(OpenTelemetry openTelemetry) {
+    this.textMapPropagator = openTelemetry.getPropagators().getTextMapPropagator();
+  }
 
-        @SuppressWarnings("null")
-        @Override
-        public void set(Metadata carrier, String key, String value) {
-            carrier.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
-        }
-    }
+  enum MetadataSetter implements TextMapSetter<Metadata> {
+    INSTANCE;
 
+    @SuppressWarnings("null")
     @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);
-                    }
+    public void set(Metadata carrier, String key, String value) {
+      carrier.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+    }
+  }
 
-                    @Override
-                    public void onClose(Status status, Metadata trailers) {
-                        super.onClose(status, trailers);
-                    }
-                }, headers);
-            }
-        };
+  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/spi/v1/GapicSpannerRpcTest.java

@@ -82,7 +82,6 @@
 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;
@@ -167,6 +166,7 @@ public static Object[] data() {
   @Before
   public void startServer() throws IOException {
     // Enable OpenTelemetry tracing.
+    SpannerOptions.resetActiveTracingFramework();
     SpannerOptions.enableOpenTelemetryTraces();
 
     assumeTrue(
@@ -552,36 +552,38 @@ public void testCustomUserAgent() {
 
   @Test
   public void testTraceContextHeaderWithOpenTelemetry() {
-    OpenTelemetry openTelemetry = OpenTelemetrySdk.builder()
-      .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
-      .setTracerProvider(SdkTracerProvider.builder().setSampler(Sampler.alwaysOn()).build())
-      .build();
+    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();
+    final SpannerOptions options =
+        createSpannerOptions().toBuilder().setOpenTelemetry(openTelemetry).build();
     try (Spanner spanner = options.getService()) {
       final DatabaseClient databaseClient =
-        spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
 
-        try (final ResultSet rs = databaseClient.singleUse().executeQuery(SELECT1AND2)) {
-          rs.next();
-        }
+      try (final ResultSet rs = databaseClient.singleUse().executeQuery(SELECT1AND2)) {
+        rs.next();
+      }
 
-        assertTrue(isTraceContextPresent);
-    } 
+      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]"));
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
 
-        try (final ResultSet rs = databaseClient.singleUse().executeQuery(SELECT1AND2)) {
-          rs.next();
-        }
+      try (final ResultSet rs = databaseClient.singleUse().executeQuery(SELECT1AND2)) {
+        rs.next();
+      }
 
-        assertFalse(isTraceContextPresent);
+      assertFalse(isTraceContextPresent);
     }
   }