feat: Create new environment variable to toggle directpath scoped to cloud bigtable. (#2261)

meeral-k · web-flow · commit 906294461027 · 2024-07-08T12:00:58.000-04:00
* 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
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/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/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>() {
