feat: Create new environment variable to toggle directpath scoped to …

…cloud bigtable. (#2261) * mend * mend * address code comments * formatting fix * fixed pom to set environment variable and not a system property * tagged myself in todo,added environment variable to directpath ipv4 only profile
googleapis · Jul 26, 2024 · 9a7ac88 · 9a7ac88
1 parent 53db8ed
commit 9a7ac88
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/google-cloud-bigtable/pom.xml b/google-cloud-bigtable/pom.xml
@@ -499,6 +499,9 @@
                     <bigtable.grpc-log-dir>${project.build.directory}/test-grpc-logs/directpath-it</bigtable.grpc-log-dir>
                     <bigtable.connection-mode>REQUIRE_DIRECT_PATH</bigtable.connection-mode>
                   </systemPropertyVariables>
+                  <environmentVariables>
+                    <CBT_ENABLE_DIRECTPATH>true</CBT_ENABLE_DIRECTPATH>
+                  </environmentVariables>
                   <includes>
                     <!-- TODO(igorbernstein): Once the control plane is accessible via directpath, add admin tests -->
                     <include>com.google.cloud.bigtable.data.v2.it.*IT</include>
@@ -575,6 +578,9 @@
                     <bigtable.grpc-log-dir>${project.build.directory}/test-grpc-logs/directpath-ipv4only-it</bigtable.grpc-log-dir>
                     <bigtable.connection-mode>REQUIRE_DIRECT_PATH_IPV4</bigtable.connection-mode>
                   </systemPropertyVariables>
+                  <environmentVariables>
+                    <CBT_ENABLE_DIRECTPATH>true</CBT_ENABLE_DIRECTPATH>
+                  </environmentVariables>
                   <includes>
                     <!-- TODO(igorbernstein): Once the control plane is accessible via directpath, add admin tests -->
                     <include>com.google.cloud.bigtable.data.v2.it.*IT</include>

diff --git a/...le/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/...le/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
@@ -100,6 +100,8 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
   private static final int MAX_MESSAGE_SIZE = 256 * 1024 * 1024;
   private static final String SERVER_DEFAULT_APP_PROFILE_ID = "";
 
+  // TODO(meeral-k): add documentation
+  private static final String CBT_ENABLE_DIRECTPATH = "CBT_ENABLE_DIRECTPATH";
   private static final Set<Code> IDEMPOTENT_RETRY_CODES =
       ImmutableSet.of(Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE);
 
@@ -345,7 +347,15 @@ public boolean getEnableRetryInfo() {
 
   /** Returns a builder for the default ChannelProvider for this service. */
   public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
-    return BigtableStubSettings.defaultGrpcTransportProviderBuilder()
+    Boolean isDirectpathEnabled = Boolean.parseBoolean(System.getenv(CBT_ENABLE_DIRECTPATH));
+    InstantiatingGrpcChannelProvider.Builder grpcTransportProviderBuilder =
+        BigtableStubSettings.defaultGrpcTransportProviderBuilder();
+    if (isDirectpathEnabled) {
+      // Attempts direct access to CBT service over gRPC to improve throughput,
+      // whether the attempt is allowed is totally controlled by service owner.
+      grpcTransportProviderBuilder.setAttemptDirectPathXds().setAttemptDirectPath(true);
+    }
+    return grpcTransportProviderBuilder
         .setChannelPoolSettings(
             ChannelPoolSettings.builder()
                 .setInitialChannelCount(10)

diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/test_helpers/env/CloudEnv.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/test_helpers/env/CloudEnv.java
@@ -168,21 +168,18 @@ private static void injectTracingCookie(
   private void configureConnection(StubSettings.Builder stubSettings) {
     // Build an remote address restricting interceptor
     final ClientInterceptor interceptor;
-    boolean enableDirectPath = false;
 
     switch (getConnectionMode()) {
       case DEFAULT:
         // nothing special
         return;
       case REQUIRE_DIRECT_PATH:
-        enableDirectPath = true;
         interceptor =
             buildRemoteAddrInterceptor(
                 "DirectPath IPv4 or IPv6",
                 Predicates.or(DIRECT_PATH_IPV4_MATCHER, DIRECT_PATH_IPV6_MATCHER));
         break;
       case REQUIRE_DIRECT_PATH_IPV4:
-        enableDirectPath = true;
         interceptor =
             buildRemoteAddrInterceptor("DirectPath IPv4", Predicates.or(DIRECT_PATH_IPV4_MATCHER));
         break;
@@ -205,10 +202,6 @@ private void configureConnection(StubSettings.Builder stubSettings) {
     final ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> oldConfigurator =
         channelProvider.getChannelConfigurator();
 
-    if (enableDirectPath) {
-      channelProvider.setAttemptDirectPath(true).setAttemptDirectPathXds();
-    }
-
     @SuppressWarnings("rawtypes")
     final ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> newConfigurator =
         new ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder>() {