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

Author: nareshz

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

@@ -150,6 +150,7 @@ public class SpannerOptions extends ServiceOptions<Spanner, SpannerOptions> {
   private final boolean useVirtualThreads;
   private final OpenTelemetry openTelemetry;
   private final boolean enableExtendedTracing;
+  private final boolean enableEndToEndTracing;
 
   enum TracingFramework {
     OPEN_CENSUS,
@@ -655,6 +656,7 @@ protected SpannerOptions(Builder builder) {
     useVirtualThreads = builder.useVirtualThreads;
     openTelemetry = builder.openTelemetry;
     enableExtendedTracing = builder.enableExtendedTracing;
+    enableEndToEndTracing = builder.enableEndToEndTracing;
   }
 
   /**
@@ -779,6 +781,7 @@ public static class Builder
     private boolean useVirtualThreads = false;
     private OpenTelemetry openTelemetry;
     private boolean enableExtendedTracing = SpannerOptions.environment.isEnableExtendedTracing();
+    private boolean enableEndToEndTracing = false;
 
     private static String createCustomClientLibToken(String token) {
       return token + " " + ServiceOptions.getGoogApiClientLibName();
@@ -843,6 +846,7 @@ protected Builder() {
       this.directedReadOptions = options.directedReadOptions;
       this.useVirtualThreads = options.useVirtualThreads;
       this.enableExtendedTracing = options.enableExtendedTracing;
+      this.enableEndToEndTracing = options.enableEndToEndTracing;
     }
 
     @Override
@@ -1352,6 +1356,15 @@ 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.
+     */
+    public Builder setEnableEndToEndTracing(boolean enableEndToEndTracing) {
+      this.enableEndToEndTracing = enableEndToEndTracing;
+      return this;
+    }
+ 
     @SuppressWarnings("rawtypes")
     @Override
     public SpannerOptions build() {
@@ -1606,6 +1619,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

@@ -266,6 +266,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;
 
   public static GapicSpannerRpc create(SpannerOptions options) {
     return new GapicSpannerRpc(options);
@@ -317,6 +318,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
     this.callCredentialsProvider = options.getCallCredentialsProvider();
     this.compressorName = options.getCompressorName();
     this.leaderAwareRoutingEnabled = options.isLeaderAwareRoutingEnabled();
+    this.endToEndTracingEnabled = options.isEndToEndTracingEnabled();
 
     if (initializeStubs) {
       // First check if SpannerOptions provides a TransportChannelProvider. Create one
@@ -1956,6 +1958,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/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/test/java/com/google/cloud/spanner/SpannerOptionsTest.java

@@ -731,6 +731,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 =

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

@@ -148,6 +148,7 @@ public class GapicSpannerRpcTest {
   private static String defaultUserAgent;
   private static Spanner spanner;
   private static boolean isRouteToLeader;
+  private static boolean isEndToEndTracing;
 
   @Parameter public Dialect dialect;
 
@@ -199,8 +200,15 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
                               Key.of(
                                   "x-goog-spanner-route-to-leader",
                                   Metadata.ASCII_STRING_MARSHALLER));
+                      String endToEndTracingHeader =
+                          headers.get(
+                              Key.of(
+                                  "x-goog-spanner-end-to-end-tracing",
+                                  Metadata.ASCII_STRING_MARSHALLER));
                       isRouteToLeader =
                           (routeToLeaderHeader != null && routeToLeaderHeader.equals("true"));
+                      isEndToEndTracing =
+                          (endToEndTracingHeader != null && endToEndTracingHeader.equals("true"));
                     }
                     return Contexts.interceptCall(Context.current(), call, headers, next);
                   }
@@ -224,6 +232,7 @@ public void reset() throws InterruptedException {
       server.awaitTermination();
     }
     isRouteToLeader = false;
+    isEndToEndTracing = false;
   }
 
   @Test
@@ -464,6 +473,77 @@ public void testNewCallContextWithRouteToLeaderHeaderAndLarDisabled() {
     rpc.shutdown();
   }
 
+  @Test
+  public void testNewCallContextWithEndToEndTracingHeader() {
+    SpannerOptions options =
+        SpannerOptions.newBuilder().setProjectId("some-project").setEnableEndToEndTracing(true).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 testNewCallContextWithoutEndToEndTracingHeader() {
+    SpannerOptions options =
+        SpannerOptions.newBuilder().setProjectId("some-project").setEnableEndToEndTracing(false).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 testEndToEndTracingHeaderWithEnabledTracing() {
+    final SpannerOptions options =
+        createSpannerOptions().toBuilder().setEnableEndToEndTracing(true).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(isEndToEndTracing);
+  }
+
+  @Test
+  public void testEndToEndTracingHeaderWithDisabledTracing() {
+    final SpannerOptions options =
+        createSpannerOptions().toBuilder().setEnableEndToEndTracing(false).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(isEndToEndTracing);
+  }
+
   @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.newEndToEndTracingHeader();
+    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");