Merge branch 'main' into PR #2663 to update

Author: lqiu96

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/GaxBasicChannelProvider.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigtable.data.v2.internal.compat;
 
+import com.google.api.core.ApiFunction;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.auth.Credentials;
 import com.google.bigtable.v2.FeatureFlags;
@@ -32,11 +33,15 @@
 import javax.annotation.Nullable;
 
 public class GaxBasicChannelProvider implements ChannelProvider {
+  private final ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> channelConfigurator;
   private final InstantiatingGrpcChannelProvider inner;
   private final @Nullable CallCredentials credentials;
 
+  @SuppressWarnings("rawtypes")
   public GaxBasicChannelProvider(
       InstantiatingGrpcChannelProvider inner, @Nullable Credentials credentials) {
+    this.channelConfigurator =
+        Optional.ofNullable(inner.toBuilder().getChannelConfigurator()).orElse(b -> b);
     this.inner = inner.toBuilder().setAttemptDirectPath(false).build();
     this.credentials = Optional.ofNullable(credentials).map(MoreCallCredentials::from).orElse(null);
   }
@@ -56,6 +61,7 @@ public ManagedChannelBuilder<?> newChannelBuilder() {
       if (credentials != null) {
         builder.intercept(new CredInterceptor(credentials));
       }
+      builder = channelConfigurator.apply(builder);
       return builder;
     } catch (IOException e) {
       throw new RuntimeException("Gax channel provider failed to provide a channel builder", e);

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/GaxDirectAccessChannelProvider.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigtable.data.v2.internal.compat;
 
+import com.google.api.core.ApiFunction;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.auth.Credentials;
 import com.google.bigtable.v2.FeatureFlags;
@@ -26,6 +27,10 @@
 
 public class GaxDirectAccessChannelProvider implements ChannelProvider {
   private final InstantiatingGrpcChannelProvider inner;
+
+  @SuppressWarnings("rawtypes")
+  private final ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> channelConfigurator;
+
   private final Optional<ChannelProvider> fallback;
 
   public static ChannelProvider create(
@@ -67,6 +72,9 @@ private GaxDirectAccessChannelProvider(
       @SuppressWarnings("unused") @Nullable Credentials credentials,
       Optional<ChannelProvider> fallback) {
     this.inner = directAccessProvider;
+    this.channelConfigurator =
+        Optional.ofNullable(directAccessProvider.toBuilder().getChannelConfigurator())
+            .orElse(b -> b);
     this.fallback = fallback;
   }
 
@@ -81,7 +89,7 @@ public FeatureFlags updateFeatureFlags(FeatureFlags featureFlags) {
   @Override
   public ManagedChannelBuilder<?> newChannelBuilder() {
     try {
-      return inner.createChannelBuilder();
+      return channelConfigurator.apply(inner.createChannelBuilder());
     } catch (IOException e) {
       throw new RuntimeException("Gax channel provider failed to provide a channel builder", e);
     }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java

@@ -34,6 +34,8 @@
 import io.grpc.StatusRuntimeException;
 import java.time.Duration;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ThreadLocalRandom;
 
 /** A callable to fork traffic between classic and session based operations. */
@@ -110,15 +112,25 @@ private boolean useExperimental(ReqT req) {
   }
 
   ApiException translateException(Throwable e) {
+    Throwable cause = e;
+    while (cause instanceof CompletionException || cause instanceof ExecutionException) {
+      if (cause.getCause() != null) {
+        cause = cause.getCause();
+      } else {
+        break;
+      }
+    }
+
     Status.Code code = Status.Code.UNKNOWN;
 
-    if (e instanceof StatusRuntimeException) {
-      code = ((StatusRuntimeException) e).getStatus().getCode();
+    if (cause instanceof StatusRuntimeException) {
+      code = ((StatusRuntimeException) cause).getStatus().getCode();
     }
-    if (e instanceof StatusException) {
-      code = ((StatusException) e).getStatus().getCode();
+    if (cause instanceof StatusException) {
+      code = ((StatusException) cause).getStatus().getCode();
     }
 
-    return ApiExceptionFactory.createException(e.getMessage(), e, GrpcStatusCode.of(code), false);
+    return ApiExceptionFactory.createException(
+        cause.getMessage(), e, GrpcStatusCode.of(code), false);
   }
 }