fix: channel primer respect other client settings

googleapis · Dec 2, 2024 · 990330b · 990330b
1 parent bac7005
commit 990330b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 26 deletions.
diff --git a/...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java
@@ -16,12 +16,11 @@
 package com.google.cloud.bigtable.data.v2.stub;
 
 import com.google.api.core.BetaApi;
-import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.core.InstantiatingExecutorProvider;
 import com.google.api.gax.grpc.ChannelPrimer;
 import com.google.api.gax.grpc.GrpcTransportChannel;
 import com.google.api.gax.rpc.FixedTransportChannelProvider;
-import com.google.auth.Credentials;
+import com.google.api.gax.rpc.InstantiatingWatchdogProvider;
 import com.google.bigtable.v2.PingAndWarmRequest;
 import com.google.cloud.bigtable.data.v2.internal.NameUtil;
 import com.google.common.base.Preconditions;
@@ -43,18 +42,13 @@ class BigtableChannelPrimer implements ChannelPrimer {
 
   private final EnhancedBigtableStubSettings settingsTemplate;
 
-  static BigtableChannelPrimer create(
-      Credentials credentials, String projectId, String instanceId, String appProfileId) {
-    EnhancedBigtableStubSettings.Builder builder =
-        EnhancedBigtableStubSettings.newBuilder()
-            .setProjectId(projectId)
-            .setInstanceId(instanceId)
-            .setAppProfileId(appProfileId)
-            .setCredentialsProvider(FixedCredentialsProvider.create(credentials))
-            // Disable refreshing channel here to avoid creating settings in a loop
-            .setRefreshingChannel(false)
-            .setExecutorProvider(
-                InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(1).build());
+  static BigtableChannelPrimer create(EnhancedBigtableStubSettings settings) {
+    EnhancedBigtableStubSettings.Builder builder = settings.toBuilder();
+    builder
+        .setRefreshingChannel(false)
+        .setBackgroundExecutorProvider(
+            InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(1).build())
+        .setStreamWatchdogProvider(InstantiatingWatchdogProvider.create());
 
     return new BigtableChannelPrimer(builder.build());
   }

diff --git a/...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
@@ -103,12 +103,7 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
       }
       // Inject channel priming if enabled
       if (builder.isRefreshingChannel()) {
-        transportProvider.setChannelPrimer(
-            BigtableChannelPrimer.create(
-                credentials,
-                settings.getProjectId(),
-                settings.getInstanceId(),
-                settings.getAppProfileId()));
+        transportProvider.setChannelPrimer(BigtableChannelPrimer.create(builder.build()));
       }
 
       builder.setTransportChannelProvider(transportProvider.build());

diff --git a/...table/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java b/...table/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java
@@ -18,6 +18,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.api.core.ApiFunction;
+import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.bigtable.v2.BigtableGrpc.BigtableImplBase;
@@ -67,12 +68,16 @@ public void setup() throws IOException {
 
     server = FakeServiceBuilder.create(fakeService).intercept(metadataInterceptor).start();
 
-    primer =
-        BigtableChannelPrimer.create(
-            OAuth2Credentials.create(new AccessToken(TOKEN_VALUE, null)),
-            "fake-project",
-            "fake-instance",
-            "fake-app-profile");
+    EnhancedBigtableStubSettings settings =
+        EnhancedBigtableStubSettings.newBuilder()
+            .setProjectId("fake-project")
+            .setInstanceId("fake-instance")
+            .setAppProfileId("fake-app-profile")
+            .setCredentialsProvider(
+                FixedCredentialsProvider.create(
+                    OAuth2Credentials.create(new AccessToken(TOKEN_VALUE, null))))
+            .build();
+    primer = BigtableChannelPrimer.create(settings);
 
     channel =
         ManagedChannelBuilder.forAddress("localhost", server.getPort()).usePlaintext().build();