From 588006664721c4775c000aab16fa81ba8c25d44d Mon Sep 17 00:00:00 2001
From: meeral-k <108340829+meeral-k@users.noreply.github.com>
Date: Mon, 8 Jul 2024 12:00:58 -0400
Subject: [PATCH] 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
---
 google-cloud-bigtable/pom.xml                        |  6 ++++++
 .../data/v2/stub/EnhancedBigtableStubSettings.java   | 12 +++++++++++-
 .../cloud/bigtable/test_helpers/env/CloudEnv.java    |  7 -------
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/google-cloud-bigtable/pom.xml b/google-cloud-bigtable/pom.xml
index 2f41833576..862eede954 100644
--- 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
index 9535f1d1ae..5a9e03cf10 100644
--- 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
index d10fd5ea5d..d7b9523b83 100644
--- 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>() {