chore: Move channel priming out of InstantiatingGrpcTransportProvider

Change-Id: I7214aa3016bd7e7f7f167c64cbaa04134b54a352

Author: nicholsl

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java

@@ -47,7 +47,7 @@
  * channel by sending a ReadRow request for a hardcoded, non-existent row key.
  */
 @BetaApi("Channel priming is not currently stable and might change in the future")
-class BigtableChannelPrimer implements ChannelPrimer {
+public class BigtableChannelPrimer implements ChannelPrimer {
   private static Logger LOG = Logger.getLogger(BigtableChannelPrimer.class.toString());
 
   static final Metadata.Key<String> REQUEST_PARAMS =

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java

@@ -121,21 +121,38 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
         setupCookieHolder(transportProvider);
       }
 
+      BigtableChannelPrimer channelPrimer =
+          BigtableChannelPrimer.create(
+              builder.getProjectId(),
+              builder.getInstanceId(),
+              builder.getAppProfileId(),
+              credentials,
+              builder.getHeaderProvider().getHeaders());
+
+      // // Inject channel priming if enabled
+      // if (builder.isRefreshingChannel()) {
+      //   transportProvider.setChannelPrimer(
+      //       BigtableChannelPrimer.create(
+      //           builder.getProjectId(),
+      //           builder.getInstanceId(),
+      //           builder.getAppProfileId(),
+      //           credentials,
+      //           builder.getHeaderProvider().getHeaders()));
+      // }
+
+      BigtableTransportChannelProvider btTransportProvider;
+
       // Inject channel priming if enabled
       if (builder.isRefreshingChannel()) {
-        transportProvider.setChannelPrimer(
-            BigtableChannelPrimer.create(
-                builder.getProjectId(),
-                builder.getInstanceId(),
-                builder.getAppProfileId(),
-                credentials,
-                builder.getHeaderProvider().getHeaders()));
+        btTransportProvider =
+            BigtableTransportChannelProvider.create(
+                (InstantiatingGrpcChannelProvider) transportProvider.build(), channelPrimer);
+      } else {
+        btTransportProvider =
+            BigtableTransportChannelProvider.createWithoutChannelPriming(
+                (InstantiatingGrpcChannelProvider) transportProvider.build());
       }
 
-      BigtableTransportChannelProvider btTransportProvider =
-          BigtableTransportChannelProvider.create(
-              (InstantiatingGrpcChannelProvider) transportProvider.build());
-
       builder.setTransportChannelProvider(btTransportProvider);
     }
 

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java

@@ -17,6 +17,7 @@
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.grpc.ChannelFactory;
+import com.google.cloud.bigtable.data.v2.stub.BigtableChannelPrimer;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
@@ -60,6 +61,8 @@ public class BigtableChannelPool extends ManagedChannel {
 
   private final BigtableChannelPoolSettings settings;
   private final ChannelFactory channelFactory;
+
+  @Nullable private final BigtableChannelPrimer channelPrimer;
   private final ScheduledExecutorService executor;
 
   private final Object entryWriteLock = new Object();
@@ -68,9 +71,12 @@ public class BigtableChannelPool extends ManagedChannel {
   private final String authority;
 
   public static BigtableChannelPool create(
-      BigtableChannelPoolSettings settings, ChannelFactory channelFactory) throws IOException {
+      BigtableChannelPoolSettings settings,
+      ChannelFactory channelFactory,
+      BigtableChannelPrimer channelPrimer)
+      throws IOException {
     return new BigtableChannelPool(
-        settings, channelFactory, Executors.newSingleThreadScheduledExecutor());
+        settings, channelFactory, channelPrimer, Executors.newSingleThreadScheduledExecutor());
   }
 
   /**
@@ -84,15 +90,21 @@ public static BigtableChannelPool create(
   BigtableChannelPool(
       BigtableChannelPoolSettings settings,
       ChannelFactory channelFactory,
+      BigtableChannelPrimer channelPrimer,
       ScheduledExecutorService executor)
       throws IOException {
     this.settings = settings;
     this.channelFactory = channelFactory;
+    this.channelPrimer = channelPrimer;
 
     ImmutableList.Builder<Entry> initialListBuilder = ImmutableList.builder();
 
     for (int i = 0; i < settings.getInitialChannelCount(); i++) {
-      initialListBuilder.add(new Entry(channelFactory.createSingleChannel()));
+      ManagedChannel newChannel = channelFactory.createSingleChannel();
+      if (channelPrimer != null) {
+        channelPrimer.primeChannel(newChannel);
+      }
+      initialListBuilder.add(new Entry(newChannel));
     }
 
     entries.set(initialListBuilder.build());
@@ -316,7 +328,11 @@ private void expand(int desiredSize) {
 
     for (int i = 0; i < desiredSize - localEntries.size(); i++) {
       try {
-        newEntries.add(new Entry(channelFactory.createSingleChannel()));
+        ManagedChannel newChannel = channelFactory.createSingleChannel();
+        if (this.channelPrimer != null) {
+          this.channelPrimer.primeChannel(newChannel);
+        }
+        newEntries.add(new Entry(newChannel));
       } catch (IOException e) {
         LOG.log(Level.WARNING, "Failed to add channel", e);
       }
@@ -354,7 +370,11 @@ void refresh() {
 
       for (int i = 0; i < newEntries.size(); i++) {
         try {
-          newEntries.set(i, new Entry(channelFactory.createSingleChannel()));
+          ManagedChannel newChannel = channelFactory.createSingleChannel();
+          if (this.channelPrimer != null) {
+            this.channelPrimer.primeChannel(newChannel);
+          }
+          newEntries.set(i, new Entry(newChannel));
         } catch (IOException e) {
           LOG.log(Level.WARNING, "Failed to refresh channel, leaving old channel", e);
         }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java

@@ -23,12 +23,14 @@
 import com.google.api.gax.rpc.TransportChannel;
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.auth.Credentials;
+import com.google.cloud.bigtable.data.v2.stub.BigtableChannelPrimer;
 import com.google.common.base.Preconditions;
 import io.grpc.ManagedChannel;
 import java.io.IOException;
 import java.util.Map;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
+import javax.annotation.Nullable;
 
 /**
  * An instance of TransportChannelProvider that provides a TransportChannel through a supplied
@@ -38,10 +40,13 @@
 public final class BigtableTransportChannelProvider implements TransportChannelProvider {
 
   private final InstantiatingGrpcChannelProvider delegate;
+  @Nullable private final BigtableChannelPrimer channelPrimer;
 
   private BigtableTransportChannelProvider(
-      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider) {
+      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider,
+      BigtableChannelPrimer channelPrimer) {
     delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider);
+    this.channelPrimer = channelPrimer;
   }
 
   @Override
@@ -63,7 +68,7 @@ public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService ex
   public BigtableTransportChannelProvider withExecutor(Executor executor) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   @Override
@@ -75,7 +80,7 @@ public boolean needsHeaders() {
   public BigtableTransportChannelProvider withHeaders(Map<String, String> headers) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   @Override
@@ -87,7 +92,7 @@ public boolean needsEndpoint() {
   public TransportChannelProvider withEndpoint(String endpoint) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   @Deprecated
@@ -101,7 +106,7 @@ public boolean acceptsPoolSize() {
   public TransportChannelProvider withPoolSize(int size) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   /** Expected to only be called once when BigtableClientContext is created */
@@ -130,7 +135,8 @@ public TransportChannel getTransportChannel() throws IOException {
     BigtableChannelPoolSettings btPoolSettings =
         BigtableChannelPoolSettings.copyFrom(delegate.getChannelPoolSettings());
 
-    BigtableChannelPool btChannelPool = BigtableChannelPool.create(btPoolSettings, channelFactory);
+    BigtableChannelPool btChannelPool =
+        BigtableChannelPool.create(btPoolSettings, channelFactory, channelPrimer);
 
     return GrpcTransportChannel.create(btChannelPool);
   }
@@ -149,12 +155,19 @@ public boolean needsCredentials() {
   public TransportChannelProvider withCredentials(Credentials credentials) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withCredentials(credentials);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   /** Creates a BigtableTransportChannelProvider. */
   public static BigtableTransportChannelProvider create(
+      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider,
+      BigtableChannelPrimer channelPrimer) {
+    Preconditions.checkNotNull(channelPrimer);
+    return new BigtableTransportChannelProvider(instantiatingGrpcChannelProvider, channelPrimer);
+  }
+
+  public static BigtableTransportChannelProvider createWithoutChannelPriming(
       InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider) {
-    return new BigtableTransportChannelProvider(instantiatingGrpcChannelProvider);
+    return new BigtableTransportChannelProvider(instantiatingGrpcChannelProvider, null);
   }
 }