chore: add directpath_enabled attribute

googleapis · Oct 4, 2024 · 2bf89b8 · 2bf89b8
1 parent f910f70
commit 2bf89b8
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 13 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java
@@ -101,6 +101,7 @@ public class BuiltInMetricsConstant {
           DIRECT_PATH_ENABLED_KEY,
           DIRECT_PATH_USED_KEY);
 
+  public static boolean DIRECT_PATH_ENABLED;
   static Aggregation AGGREGATION_WITH_MILLIS_HISTOGRAM =
       Aggregation.explicitBucketHistogram(
           ImmutableList.of(

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/CompositeTracer.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/CompositeTracer.java
@@ -184,7 +184,7 @@ public void addAttributes(Map<String, String> attributes) {
     for (ApiTracer child : children) {
       if (child instanceof MetricsTracer) {
         MetricsTracer metricsTracer = (MetricsTracer) child;
-        attributes.forEach((key, value) -> metricsTracer.addAttributes(key, value));
+        metricsTracer.addAttributes(attributes);
       }
     }
   }

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
@@ -29,6 +29,7 @@
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.grpc.GrpcCallSettings;
 import com.google.api.gax.grpc.GrpcStubCallableFactory;
+import com.google.api.gax.grpc.GrpcTransportChannel;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
@@ -63,6 +64,7 @@
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.spanner.AdminRequestsPerMinuteExceededException;
 import com.google.cloud.spanner.BackupId;
+import com.google.cloud.spanner.BuiltInMetricsConstant;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Restore;
 import com.google.cloud.spanner.SpannerException;
@@ -396,18 +398,22 @@ public GapicSpannerRpc(final SpannerOptions options) {
       final String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
 
       try {
-        this.spannerStub =
-            GrpcSpannerStub.create(
-                options
-                    .getSpannerStubSettings()
-                    .toBuilder()
-                    .setTransportChannelProvider(channelProvider)
-                    .setCredentialsProvider(credentialsProvider)
-                    .setStreamWatchdogProvider(watchdogProvider)
-                    .setTracerFactory(
-                        options.getApiTracerFactory(
-                            /* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
-                    .build());
+        SpannerStubSettings spannerStubSettings =
+            options
+                .getSpannerStubSettings()
+                .toBuilder()
+                .setTransportChannelProvider(channelProvider)
+                .setCredentialsProvider(credentialsProvider)
+                .setStreamWatchdogProvider(watchdogProvider)
+                .setTracerFactory(
+                    options.getApiTracerFactory(
+                        /* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
+                .build();
+        ClientContext clientContext = ClientContext.create(spannerStubSettings);
+        this.spannerStub = GrpcSpannerStub.create(spannerStubSettings, clientContext);
+        BuiltInMetricsConstant.DIRECT_PATH_ENABLED =
+            ((GrpcTransportChannel) clientContext.getTransportChannel()).isDirectPath()
+                && isAttemptDirectPathXds;
         this.readRetrySettings =
             options.getSpannerStubSettings().streamingReadSettings().getRetrySettings();
         this.readRetryableCodes =

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java
@@ -228,6 +228,9 @@ private Map<String, String> getBuiltInMetricAttributes(String key, DatabaseName
           attributes.put(BuiltInMetricsConstant.DATABASE_KEY.getKey(), databaseName.getDatabase());
           attributes.put(
               BuiltInMetricsConstant.INSTANCE_ID_KEY.getKey(), databaseName.getInstance());
+          attributes.put(
+              BuiltInMetricsConstant.DIRECT_PATH_ENABLED_KEY.getKey(),
+              String.valueOf(BuiltInMetricsConstant.DIRECT_PATH_ENABLED));
           return attributes;
         });
   }

diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/v1/stub/GrpcSpannerStub.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/v1/stub/GrpcSpannerStub.java
@@ -247,6 +247,10 @@ public static final GrpcSpannerStub create(ClientContext clientContext) throws I
     return new GrpcSpannerStub(SpannerStubSettings.newBuilder().build(), clientContext);
   }
 
+  public static final GrpcSpannerStub create(SpannerStubSettings settings, ClientContext clientContext) throws IOException {
+    return new GrpcSpannerStub(settings, clientContext);
+  }
+
   public static final GrpcSpannerStub create(
       ClientContext clientContext, GrpcStubCallableFactory callableFactory) throws IOException {
     return new GrpcSpannerStub(