add ability to disable grpc-gcp if needed

Author: rahul2393

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java

@@ -151,6 +151,7 @@ public class SpannerOptions extends ServiceOptions<Spanner, SpannerOptions> {
   private final InstanceAdminStubSettings instanceAdminStubSettings;
   private final DatabaseAdminStubSettings databaseAdminStubSettings;
   private final Duration partitionedDmlTimeout;
+  private final boolean grpcGcpExtensionEnabled;
   private final GcpManagedChannelOptions grpcGcpOptions;
   private final boolean autoThrottleAdministrativeRequests;
   private final RetrySettings retryAdministrativeRequestsSettings;
@@ -797,6 +798,7 @@ protected SpannerOptions(Builder builder) {
       throw SpannerExceptionFactory.newSpannerException(e);
     }
     partitionedDmlTimeout = builder.partitionedDmlTimeout;
+    grpcGcpExtensionEnabled = builder.grpcGcpExtensionEnabled;
     grpcGcpOptions = builder.grpcGcpOptions;
     autoThrottleAdministrativeRequests = builder.autoThrottleAdministrativeRequests;
     retryAdministrativeRequestsSettings = builder.retryAdministrativeRequestsSettings;
@@ -1023,6 +1025,7 @@ public static class Builder
     private DatabaseAdminStubSettings.Builder databaseAdminStubSettingsBuilder =
         DatabaseAdminStubSettings.newBuilder();
     private Duration partitionedDmlTimeout = Duration.ofHours(2L);
+    private boolean grpcGcpExtensionEnabled = true;
     private GcpManagedChannelOptions grpcGcpOptions;
     private RetrySettings retryAdministrativeRequestsSettings =
         DEFAULT_ADMIN_REQUESTS_LIMIT_EXCEEDED_RETRY_SETTINGS;
@@ -1094,6 +1097,7 @@ protected Builder() {
       this.instanceAdminStubSettingsBuilder = options.instanceAdminStubSettings.toBuilder();
       this.databaseAdminStubSettingsBuilder = options.databaseAdminStubSettings.toBuilder();
       this.partitionedDmlTimeout = options.partitionedDmlTimeout;
+      this.grpcGcpExtensionEnabled = options.grpcGcpExtensionEnabled;
       this.grpcGcpOptions = options.grpcGcpOptions;
       this.autoThrottleAdministrativeRequests = options.autoThrottleAdministrativeRequests;
       this.retryAdministrativeRequestsSettings = options.retryAdministrativeRequestsSettings;
@@ -1569,12 +1573,14 @@ public Builder enableGrpcGcpExtension() {
      * Multiplexed sessions are not supported for gRPC-GCP.
      */
     public Builder enableGrpcGcpExtension(GcpManagedChannelOptions options) {
+      this.grpcGcpExtensionEnabled = true;
       this.grpcGcpOptions = options;
       return this;
     }
 
-    /** Disables gRPC-GCP extension. */
+    /** Disables gRPC-GCP extension and uses GAX channel pool instead. */
     public Builder disableGrpcGcpExtension() {
+      this.grpcGcpExtensionEnabled = false;
       return this;
     }
 
@@ -1787,7 +1793,8 @@ public SpannerOptions build() {
         credentials = environment.getDefaultExperimentalHostCredentials();
       }
       if (this.numChannels == null) {
-        this.numChannels = GRPC_GCP_ENABLED_DEFAULT_CHANNELS;
+        this.numChannels =
+            this.grpcGcpExtensionEnabled ? GRPC_GCP_ENABLED_DEFAULT_CHANNELS : DEFAULT_CHANNELS;
       }
 
       synchronized (lock) {
@@ -1982,7 +1989,7 @@ public Duration getPartitionedDmlTimeoutDuration() {
   }
 
   public boolean isGrpcGcpExtensionEnabled() {
-    return true;
+    return grpcGcpExtensionEnabled;
   }
 
   public GcpManagedChannelOptions getGrpcGcpOptions() {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -279,6 +279,7 @@ public class GapicSpannerRpc implements SpannerRpc {
   private final boolean leaderAwareRoutingEnabled;
   private final boolean endToEndTracingEnabled;
   private final int numChannels;
+  private final boolean isGrpcGcpExtensionEnabled;
 
   private final GrpcCallContext baseGrpcCallContext;
 
@@ -334,6 +335,7 @@ public GapicSpannerRpc(final SpannerOptions options) {
     this.leaderAwareRoutingEnabled = options.isLeaderAwareRoutingEnabled();
     this.endToEndTracingEnabled = options.isEndToEndTracingEnabled();
     this.numChannels = options.getNumChannels();
+    this.isGrpcGcpExtensionEnabled = options.isGrpcGcpExtensionEnabled();
     this.baseGrpcCallContext = createBaseCallContext();
 
     if (initializeStubs) {
@@ -590,6 +592,10 @@ private static GcpManagedChannelOptions grpcGcpOptionsWithMetrics(SpannerOptions
   private static void maybeEnableGrpcGcpExtension(
       InstantiatingGrpcChannelProvider.Builder defaultChannelProviderBuilder,
       final SpannerOptions options) {
+    if (!options.isGrpcGcpExtensionEnabled()) {
+      return;
+    }
+
     final String jsonApiConfig = parseGrpcGcpApiConfig();
     final GcpManagedChannelOptions grpcGcpOptions = grpcGcpOptionsWithMetrics(options);
 
@@ -2031,14 +2037,20 @@ <ReqT, RespT> GrpcCallContext newCallContext(
     GrpcCallContext context = this.baseGrpcCallContext;
     Long affinity = options == null ? null : Option.CHANNEL_HINT.getLong(options);
     if (affinity != null) {
-      // Set channel affinity in gRPC-GCP.
-      // Compute bounded channel hint to prevent gRPC-GCP affinity map from getting unbounded.
-      int boundedChannelHint = affinity.intValue() % this.numChannels;
-      context =
-          context.withCallOptions(
-              context
-                  .getCallOptions()
-                  .withOption(GcpManagedChannel.AFFINITY_KEY, String.valueOf(boundedChannelHint)));
+      if (this.isGrpcGcpExtensionEnabled) {
+        // Set channel affinity in gRPC-GCP.
+        // Compute bounded channel hint to prevent gRPC-GCP affinity map from getting unbounded.
+        int boundedChannelHint = affinity.intValue() % this.numChannels;
+        context =
+            context.withCallOptions(
+                context
+                    .getCallOptions()
+                    .withOption(
+                        GcpManagedChannel.AFFINITY_KEY, String.valueOf(boundedChannelHint)));
+      } else {
+        // Set channel affinity in GAX.
+        context = context.withChannelAffinity(affinity.intValue());
+      }
     }
     if (options != null) {
       // TODO(@odeke-em): Infer the affinity if it doesn't match up with in the request-id.