From 990330b917d3bacf2fa2cf94db61c5aa3054dfa9 Mon Sep 17 00:00:00 2001
From: Mattie Fu <mattiefu@google.com>
Date: Mon, 2 Dec 2024 09:57:10 -0500
Subject: [PATCH] fix: channel primer respect other client settings

---
 .../data/v2/stub/BigtableChannelPrimer.java   | 22 +++++++------------
 .../data/v2/stub/BigtableClientContext.java   |  7 +-----
 .../v2/stub/BigtableChannelPrimerTest.java    | 17 +++++++++-----
 3 files changed, 20 insertions(+), 26 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java
index ecbef85be5..3fd6614a3d 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java
+++ b/google-cloud-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/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
index d23b487caf..ebd4df0bcd 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java
+++ b/google-cloud-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/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java
index e1f22bebbd..bfacd1d769 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java
+++ b/google-cloud-bigtable/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();