Skip to content

Commit ea1ebad

Browse files
surbhigarg92surbhigarg92
authored
fix: directpath enabled attribute (#3477)
* chore: add directpath_enabled attribute * Created new grpc wrapper * Created new grpc wrapper * review comments
1 parent bf350b0 commit ea1ebad

File tree

6 files changed

+93
-20
lines changed

6 files changed

+93
-20
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInOpenTelemetryMetricsProvider.java

-3
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import static com.google.cloud.spanner.BuiltInMetricsConstant.CLIENT_HASH_KEY;
2121
import static com.google.cloud.spanner.BuiltInMetricsConstant.CLIENT_NAME_KEY;
2222
import static com.google.cloud.spanner.BuiltInMetricsConstant.CLIENT_UID_KEY;
23-
import static com.google.cloud.spanner.BuiltInMetricsConstant.DIRECT_PATH_ENABLED_KEY;
2423
import static com.google.cloud.spanner.BuiltInMetricsConstant.INSTANCE_CONFIG_ID_KEY;
2524
import static com.google.cloud.spanner.BuiltInMetricsConstant.LOCATION_ID_KEY;
2625
import static com.google.cloud.spanner.BuiltInMetricsConstant.PROJECT_ID_KEY;
@@ -83,8 +82,6 @@ Map<String, String> createClientAttributes(String projectId, String client_name)
8382
Map<String, String> clientAttributes = new HashMap<>();
8483
clientAttributes.put(LOCATION_ID_KEY.getKey(), detectClientLocation());
8584
clientAttributes.put(PROJECT_ID_KEY.getKey(), projectId);
86-
// TODO: Replace this with real value.
87-
clientAttributes.put(DIRECT_PATH_ENABLED_KEY.getKey(), "false");
8885
clientAttributes.put(INSTANCE_CONFIG_ID_KEY.getKey(), "unknown");
8986
clientAttributes.put(CLIENT_NAME_KEY.getKey(), client_name);
9087
String clientUid = getDefaultTaskValue();

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

+30-14
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import com.google.api.gax.grpc.GrpcCallContext;
3030
import com.google.api.gax.grpc.GrpcCallSettings;
3131
import com.google.api.gax.grpc.GrpcStubCallableFactory;
32+
import com.google.api.gax.grpc.GrpcTransportChannel;
3233
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
3334
import com.google.api.gax.longrunning.OperationFuture;
3435
import com.google.api.gax.retrying.ResultRetryAlgorithm;
@@ -78,13 +79,14 @@
7879
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
7980
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
8081
import com.google.cloud.spanner.encryption.EncryptionConfigProtoMapper;
81-
import com.google.cloud.spanner.v1.stub.GrpcSpannerStub;
8282
import com.google.cloud.spanner.v1.stub.SpannerStub;
8383
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
8484
import com.google.common.annotations.VisibleForTesting;
8585
import com.google.common.base.Function;
8686
import com.google.common.base.MoreObjects;
8787
import com.google.common.base.Preconditions;
88+
import com.google.common.base.Supplier;
89+
import com.google.common.base.Suppliers;
8890
import com.google.common.collect.ImmutableList;
8991
import com.google.common.collect.ImmutableSet;
9092
import com.google.common.io.Resources;
@@ -276,6 +278,8 @@ public class GapicSpannerRpc implements SpannerRpc {
276278
private final int numChannels;
277279
private final boolean isGrpcGcpExtensionEnabled;
278280

281+
private Supplier<Boolean> directPathEnabledSupplier = () -> false;
282+
279283
public static GapicSpannerRpc create(SpannerOptions options) {
280284
return new GapicSpannerRpc(options);
281285
}
@@ -351,7 +355,9 @@ public GapicSpannerRpc(final SpannerOptions options) {
351355
SpannerInterceptorProvider.create(
352356
MoreObjects.firstNonNull(
353357
options.getInterceptorProvider(),
354-
SpannerInterceptorProvider.createDefault(options.getOpenTelemetry())))
358+
SpannerInterceptorProvider.createDefault(
359+
options.getOpenTelemetry(),
360+
(() -> directPathEnabledSupplier.get()))))
355361
// This sets the trace context headers.
356362
.withTraceContext(endToEndTracingEnabled, options.getOpenTelemetry())
357363
// This sets the response compressor (Server -> Client).
@@ -396,18 +402,27 @@ public GapicSpannerRpc(final SpannerOptions options) {
396402
final String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST");
397403

398404
try {
405+
SpannerStubSettings spannerStubSettings =
406+
options
407+
.getSpannerStubSettings()
408+
.toBuilder()
409+
.setTransportChannelProvider(channelProvider)
410+
.setCredentialsProvider(credentialsProvider)
411+
.setStreamWatchdogProvider(watchdogProvider)
412+
.setTracerFactory(
413+
options.getApiTracerFactory(
414+
/* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
415+
.build();
416+
ClientContext clientContext = ClientContext.create(spannerStubSettings);
399417
this.spannerStub =
400-
GrpcSpannerStub.create(
401-
options
402-
.getSpannerStubSettings()
403-
.toBuilder()
404-
.setTransportChannelProvider(channelProvider)
405-
.setCredentialsProvider(credentialsProvider)
406-
.setStreamWatchdogProvider(watchdogProvider)
407-
.setTracerFactory(
408-
options.getApiTracerFactory(
409-
/* isAdminClient = */ false, isEmulatorEnabled(options, emulatorHost)))
410-
.build());
418+
GrpcSpannerStubWithStubSettingsAndClientContext.create(
419+
spannerStubSettings, clientContext);
420+
this.directPathEnabledSupplier =
421+
Suppliers.memoize(
422+
() -> {
423+
return ((GrpcTransportChannel) clientContext.getTransportChannel()).isDirectPath()
424+
&& isAttemptDirectPathXds;
425+
});
411426
this.readRetrySettings =
412427
options.getSpannerStubSettings().streamingReadSettings().getRetrySettings();
413428
this.readRetryableCodes =
@@ -455,7 +470,8 @@ public GapicSpannerRpc(final SpannerOptions options) {
455470
.getStreamWatchdogProvider()
456471
.withCheckInterval(pdmlSettings.getStreamWatchdogCheckInterval()));
457472
}
458-
this.partitionedDmlStub = GrpcSpannerStub.create(pdmlSettings.build());
473+
this.partitionedDmlStub =
474+
GrpcSpannerStubWithStubSettingsAndClientContext.create(pdmlSettings.build());
459475
this.instanceAdminStubSettings =
460476
options
461477
.getInstanceAdminStubSettings()

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

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.spi.v1;
18+
19+
import com.google.api.gax.rpc.ClientContext;
20+
import com.google.cloud.spanner.v1.stub.GrpcSpannerStub;
21+
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
22+
import java.io.IOException;
23+
24+
/**
25+
* Wrapper around {@link GrpcSpannerStub} to make the constructor available inside this package.
26+
* This makes it possible to create a {@link GrpcSpannerStub} with a {@link SpannerStubSettings} and
27+
* a {@link ClientContext}.
28+
*/
29+
class GrpcSpannerStubWithStubSettingsAndClientContext extends GrpcSpannerStub {
30+
31+
static final GrpcSpannerStubWithStubSettingsAndClientContext create(
32+
SpannerStubSettings settings, ClientContext clientContext) throws IOException {
33+
return new GrpcSpannerStubWithStubSettingsAndClientContext(settings, clientContext);
34+
}
35+
36+
protected GrpcSpannerStubWithStubSettingsAndClientContext(
37+
SpannerStubSettings settings, ClientContext clientContext) throws IOException {
38+
super(settings, clientContext);
39+
}
40+
}

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

+9-1
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.google.cloud.spanner.CompositeTracer;
2929
import com.google.cloud.spanner.SpannerExceptionFactory;
3030
import com.google.cloud.spanner.SpannerRpcMetrics;
31+
import com.google.common.base.Supplier;
3132
import com.google.common.cache.Cache;
3233
import com.google.common.cache.CacheBuilder;
3334
import com.google.spanner.admin.database.v1.DatabaseName;
@@ -93,8 +94,12 @@ class HeaderInterceptor implements ClientInterceptor {
9394
private static final Level LEVEL = Level.INFO;
9495
private final SpannerRpcMetrics spannerRpcMetrics;
9596

96-
HeaderInterceptor(SpannerRpcMetrics spannerRpcMetrics) {
97+
private final Supplier<Boolean> directPathEnabledSupplier;
98+
99+
HeaderInterceptor(
100+
SpannerRpcMetrics spannerRpcMetrics, Supplier<Boolean> directPathEnabledSupplier) {
97101
this.spannerRpcMetrics = spannerRpcMetrics;
102+
this.directPathEnabledSupplier = directPathEnabledSupplier;
98103
}
99104

100105
@Override
@@ -228,6 +233,9 @@ private Map<String, String> getBuiltInMetricAttributes(String key, DatabaseName
228233
attributes.put(BuiltInMetricsConstant.DATABASE_KEY.getKey(), databaseName.getDatabase());
229234
attributes.put(
230235
BuiltInMetricsConstant.INSTANCE_ID_KEY.getKey(), databaseName.getInstance());
236+
attributes.put(
237+
BuiltInMetricsConstant.DIRECT_PATH_ENABLED_KEY.getKey(),
238+
String.valueOf(this.directPathEnabledSupplier.get()));
231239
return attributes;
232240
});
233241
}

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

+14-1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import com.google.api.core.ObsoleteApi;
2020
import com.google.api.gax.grpc.GrpcInterceptorProvider;
2121
import com.google.cloud.spanner.SpannerRpcMetrics;
22+
import com.google.common.base.Supplier;
23+
import com.google.common.base.Suppliers;
2224
import com.google.common.collect.ImmutableList;
2325
import io.grpc.ClientInterceptor;
2426
import io.opentelemetry.api.GlobalOpenTelemetry;
@@ -46,11 +48,22 @@ public static SpannerInterceptorProvider createDefault() {
4648
}
4749

4850
public static SpannerInterceptorProvider createDefault(OpenTelemetry openTelemetry) {
51+
return createDefault(
52+
openTelemetry,
53+
Suppliers.memoize(
54+
() -> {
55+
return false;
56+
}));
57+
}
58+
59+
public static SpannerInterceptorProvider createDefault(
60+
OpenTelemetry openTelemetry, Supplier<Boolean> directPathEnabledSupplier) {
4961
List<ClientInterceptor> defaultInterceptorList = new ArrayList<>();
5062
defaultInterceptorList.add(new SpannerErrorInterceptor());
5163
defaultInterceptorList.add(
5264
new LoggingInterceptor(Logger.getLogger(GapicSpannerRpc.class.getName()), Level.FINER));
53-
defaultInterceptorList.add(new HeaderInterceptor(new SpannerRpcMetrics(openTelemetry)));
65+
defaultInterceptorList.add(
66+
new HeaderInterceptor(new SpannerRpcMetrics(openTelemetry), directPathEnabledSupplier));
5467
return new SpannerInterceptorProvider(ImmutableList.copyOf(defaultInterceptorList));
5568
}
5669

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetryBuiltInMetricsTracerTest.java

-1
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,6 @@ public static void setup() {
9494
Attributes.builder()
9595
.put(BuiltInMetricsConstant.PROJECT_ID_KEY, "test-project")
9696
.put(BuiltInMetricsConstant.INSTANCE_CONFIG_ID_KEY, "unknown")
97-
.put(BuiltInMetricsConstant.DIRECT_PATH_ENABLED_KEY, "false")
9897
.put(
9998
BuiltInMetricsConstant.LOCATION_ID_KEY,
10099
BuiltInOpenTelemetryMetricsProvider.detectClientLocation())

0 commit comments

Comments
 (0)