int commit

Author: sushanb

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

@@ -26,6 +26,7 @@
 import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer;
 import io.grpc.CallCredentials;
 import io.grpc.CallOptions;
+import io.grpc.Channel;
 import io.grpc.ClientCall;
 import io.grpc.Deadline;
 import io.grpc.ManagedChannel;
@@ -88,21 +89,21 @@ static BigtableChannelPrimer create(
   }
 
   @Override
-  public void primeChannel(ManagedChannel managedChannel) {
+  public void primeChannel(Channel channel) {
     try {
-      primeChannelUnsafe(managedChannel);
+      primeChannelUnsafe(channel);
     } catch (IOException | RuntimeException e) {
       LOG.log(Level.WARNING, "Unexpected error while trying to prime a channel", e);
     }
   }
 
-  private void primeChannelUnsafe(ManagedChannel managedChannel) throws IOException {
-    sendPrimeRequestsBlocking(managedChannel);
+  private void primeChannelUnsafe(Channel channel) throws IOException {
+    sendPrimeRequestsBlocking(channel);
   }
 
-  private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) {
+  private void sendPrimeRequestsBlocking(Channel channel) {
     try {
-      sendPrimeRequestsAsync(managedChannel).get(1, TimeUnit.MINUTES);
+      sendPrimeRequestsAsync(channel).get(1, TimeUnit.MINUTES);
     } catch (Throwable e) {
       // TODO: Not sure if we should swallow the error here. We are pre-emptively swapping
       // channels if the new
@@ -112,7 +113,7 @@ private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) {
   }
 
   @Override
-  public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel managedChannel) {
+  public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(Channel managedChannel) {
     ClientCall<PingAndWarmRequest, PingAndWarmResponse> clientCall =
         managedChannel.newCall(
             BigtableGrpc.getPingAndWarmMethod(),

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

@@ -242,6 +242,25 @@ public ClientOperationSettings getPerOpSettings() {
     return perOpSettings;
   }
 
+  /** Applies common pool, message size, and keep-alive settings to the provided builder. */
+  private static InstantiatingGrpcChannelProvider.Builder commonTraits(
+          InstantiatingGrpcChannelProvider.Builder builder) {
+    return builder
+            .setChannelPoolSettings(
+                    ChannelPoolSettings.builder()
+                            .setInitialChannelCount(10)
+                            .setMinRpcsPerChannel(1)
+                            // Keep it conservative as we scale the channel size every 1min
+                            // and delta is 2 channels.
+                            .setMaxRpcsPerChannel(25)
+                            .setPreemptiveRefreshEnabled(true)
+                            .build())
+            .setMaxInboundMessageSize(MAX_MESSAGE_SIZE)
+            .setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval
+            .setKeepAliveTimeout(
+                    Duration.ofSeconds(10)); // wait this long before considering the connection dead
+  }
+
   /** Returns a builder for the default ChannelProvider for this service. */
   public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
     InstantiatingGrpcChannelProvider.Builder grpcTransportProviderBuilder =
@@ -261,22 +280,27 @@ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProvi
             Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS));
       }
     }
-    return grpcTransportProviderBuilder
-        .setChannelPoolSettings(
-            ChannelPoolSettings.builder()
-                .setInitialChannelCount(10)
-                .setMinRpcsPerChannel(1)
-                // Keep it conservative as we scale the channel size every 1min
-                // and delta is 2 channels.
-                .setMaxRpcsPerChannel(25)
-                .setPreemptiveRefreshEnabled(true)
-                .build())
-        .setMaxInboundMessageSize(MAX_MESSAGE_SIZE)
-        .setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval
-        .setKeepAliveTimeout(
-            Duration.ofSeconds(10)); // wait this long before considering the connection dead
+    return commonTraits(grpcTransportProviderBuilder);
   }
 
+  /** Applies Direct Access traits (DirectPath & ALTS) to an existing builder. */
+  public static InstantiatingGrpcChannelProvider.Builder applyDirectAccessTraits(
+          InstantiatingGrpcChannelProvider.Builder builder) {
+
+    builder
+            .setAttemptDirectPathXds()
+            .setAttemptDirectPath(true)
+            .setAllowNonDefaultServiceAccount(true);
+
+    if (!DIRECT_PATH_BOUND_TOKEN_DISABLED) {
+      builder.setAllowHardBoundTokenTypes(
+              Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS));
+    }
+
+    return builder;
+  }
+
+
   @SuppressWarnings("WeakerAccess")
   public static TransportChannelProvider defaultTransportChannelProvider() {
     return defaultGrpcTransportProviderBuilder().build();

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

@@ -16,11 +16,11 @@
 package com.google.cloud.bigtable.data.v2.stub;
 
 import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
 import com.google.api.core.InternalApi;
-import com.google.api.core.SettableApiFuture;
 import com.google.bigtable.v2.PingAndWarmResponse;
 import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer;
-import io.grpc.ManagedChannel;
+import io.grpc.Channel;
 
 @InternalApi
 public class NoOpChannelPrimer implements ChannelPrimer {
@@ -31,14 +31,12 @@ static NoOpChannelPrimer create() {
   private NoOpChannelPrimer() {}
 
   @Override
-  public void primeChannel(ManagedChannel channel) {
+  public void primeChannel(Channel channel) {
     // No op
   }
 
   @Override
-  public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel channel) {
-    SettableApiFuture future = SettableApiFuture.create();
-    future.set(PingAndWarmResponse.getDefaultInstance());
-    return future;
+  public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(Channel channel) {
+    return ApiFutures.immediateFuture(PingAndWarmResponse.getDefaultInstance());
   }
 }