Merge branch 'e2e-tracing-header' into grpc-telemetry-client-interceptor

Author: nareshz

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;
   }
 
   /**
@@ -799,6 +801,7 @@ public static class Builder
     private OpenTelemetry openTelemetry;
     private boolean enableApiTracing = SpannerOptions.environment.isEnableApiTracing();
     private boolean enableExtendedTracing = SpannerOptions.environment.isEnableExtendedTracing();
+    private boolean enableServerSideTracing = false;
 
     private static String createCustomClientLibToken(String token) {
       return token + " " + ServiceOptions.getGoogApiClientLibName();
@@ -864,6 +867,7 @@ protected Builder() {
       this.useVirtualThreads = options.useVirtualThreads;
       this.enableApiTracing = options.enableApiTracing;
       this.enableExtendedTracing = options.enableExtendedTracing;
+      this.enableServerSideTracing = options.enableServerSideTracing;
     }
 
     @Override
@@ -1391,6 +1395,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() {
@@ -1681,6 +1703,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();
 
@@ -1988,6 +1990,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/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/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 =

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

@@ -154,6 +154,7 @@ public class GapicSpannerRpcTest {
   private static String defaultUserAgent;
   private static Spanner spanner;
   private static boolean isRouteToLeader;
+  private static boolean isServerSideTracing;
   private static boolean isTraceContextPresent;
 
   @Parameter public Dialect dialect;
@@ -213,8 +214,15 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
                               Key.of(
                                   "x-goog-spanner-route-to-leader",
                                   Metadata.ASCII_STRING_MARSHALLER));
+                      String serverSideTracingHeader =
+                          headers.get(
+                              Key.of(
+                                  "x-goog-spanner-end-to-end-tracing",
+                                  Metadata.ASCII_STRING_MARSHALLER));
                       isRouteToLeader =
                           (routeToLeaderHeader != null && routeToLeaderHeader.equals("true"));
+                      isServerSideTracing =
+                          (serverSideTracingHeader != null && serverSideTracingHeader.equals("true"));
                     }
                     return Contexts.interceptCall(Context.current(), call, headers, next);
                   }
@@ -238,6 +246,7 @@ public void reset() throws InterruptedException {
       server.awaitTermination();
     }
     isRouteToLeader = false;
+    isServerSideTracing = false;
     isTraceContextPresent = false;
   }
 
@@ -479,6 +488,77 @@ public void testNewCallContextWithRouteToLeaderHeaderAndLarDisabled() {
     rpc.shutdown();
   }
 
+  @Test
+  public void testNewCallContextWithServerSideTracingHeader() {
+    SpannerOptions options =
+        SpannerOptions.newBuilder().setProjectId("some-project").enableServerSideTracing().build();
+    GapicSpannerRpc rpc = new GapicSpannerRpc(options, false);
+    GrpcCallContext callContext =
+        rpc.newCallContext(
+            optionsMap,
+            "/some/resource",
+            ExecuteSqlRequest.getDefaultInstance(),
+            SpannerGrpc.getExecuteSqlMethod());
+    assertNotNull(callContext);
+    assertEquals(
+        ImmutableList.of("true"),
+        callContext.getExtraHeaders().get("x-goog-spanner-end-to-end-tracing"));
+    assertEquals(
+        ImmutableList.of("projects/some-project"),
+        callContext.getExtraHeaders().get(ApiClientHeaderProvider.getDefaultResourceHeaderKey()));
+    rpc.shutdown();
+  }
+
+  @Test
+  public void testNewCallContextWithoutServerSideTracingHeader() {
+    SpannerOptions options =
+        SpannerOptions.newBuilder().setProjectId("some-project").disableServerSideTracing().build();
+    GapicSpannerRpc rpc = new GapicSpannerRpc(options, false);
+    GrpcCallContext callContext =
+        rpc.newCallContext(
+            optionsMap,
+            "/some/resource",
+            ExecuteSqlRequest.getDefaultInstance(),
+            SpannerGrpc.getExecuteSqlMethod());
+    assertNotNull(callContext);
+    assertNull(callContext.getExtraHeaders().get("x-goog-spanner-end-to-end-tracing"));
+    rpc.shutdown();
+  }
+
+  @Test
+  public void testServerSideTracingHeaderWithEnabledTracing() {
+    final SpannerOptions options =
+        createSpannerOptions().toBuilder().enableServerSideTracing().build();
+    try (Spanner spanner = options.getService()) {
+      final DatabaseClient databaseClient =
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+      TransactionRunner runner = databaseClient.readWriteTransaction();
+      runner.run(
+          transaction -> {
+            transaction.executeUpdate(UPDATE_FOO_STATEMENT);
+            return null;
+          });
+    }
+    assertTrue(isServerSideTracing);
+  }
+
+  @Test
+  public void testServerSideTracingHeaderWithDisabledTracing() {
+    final SpannerOptions options =
+        createSpannerOptions().toBuilder().disableServerSideTracing().build();
+    try (Spanner spanner = options.getService()) {
+      final DatabaseClient databaseClient =
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+      TransactionRunner runner = databaseClient.readWriteTransaction();
+      runner.run(
+          transaction -> {
+            transaction.executeUpdate(UPDATE_FOO_STATEMENT);
+            return null;
+          });
+    }
+    assertFalse(isServerSideTracing);
+  }
+
   @Test
   public void testAdminRequestsLimitExceededRetryAlgorithm() {
     AdminRequestsLimitExceededRetryAlgorithm<Long> alg =

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

@@ -94,6 +94,17 @@ public void testNewRouteToLeaderHeader() {
     assertTrue(Maps.difference(extraHeaders, expectedHeaders).areEqual());
   }
 
+  @Test
+  public void testNewEndToEndTracingHeader() {
+    SpannerMetadataProvider metadataProvider =
+        SpannerMetadataProvider.create(ImmutableMap.of(), "header1");
+    Map<String, List<String>> extraHeaders = metadataProvider.newServerSideTracingHeader();
+    Map<String, List<String>> expectedHeaders =
+        ImmutableMap.<String, List<String>>of(
+            "x-goog-spanner-end-to-end-tracing", ImmutableList.of("true"));
+    assertTrue(Maps.difference(extraHeaders, expectedHeaders).areEqual());
+  }
+
   private String getResourceHeaderValue(
       SpannerMetadataProvider headerProvider, String resourceTokenTemplate) {
     Metadata metadata = headerProvider.newMetadata(resourceTokenTemplate, "projects/p");