feat(gax): implement dynamic channel refreshing on 401 retries

Author: blakeli0

sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java

@@ -82,6 +82,7 @@ class ChannelPool extends ManagedChannel {
   private ScheduledFuture<?> resizeFuture = null;
 
   private final Object entryWriteLock = new Object();
+  private long lastRefreshTimeNanos = 0;
   @VisibleForTesting final AtomicReference<ImmutableList<Entry>> entries = new AtomicReference<>();
   private final AtomicInteger indexTicker = new AtomicInteger();
   private final String authority;
@@ -441,6 +442,13 @@ void refresh() {
     // - then thread2 will shut down channel that thread1 will put back into circulation (after it
     //   replaces the list)
     synchronized (entryWriteLock) {
+      long now = System.nanoTime();
+      if (now - lastRefreshTimeNanos < TimeUnit.SECONDS.toNanos(5)) {
+        LOG.fine("Channel pool was refreshed recently, skipping duplicate refresh");
+        return;
+      }
+      lastRefreshTimeNanos = now;
+
       LOG.fine("Refreshing all channels");
       ArrayList<Entry> newEntries = new ArrayList<>(entries.get());
 

sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java

@@ -211,6 +211,14 @@ public GrpcCallContext withCredentials(Credentials newCredentials) {
         isDirectPath);
   }
 
+  @Override
+  public TransportChannel getTransportChannel() {
+    if (channel instanceof io.grpc.ManagedChannel) {
+      return GrpcTransportChannel.create((io.grpc.ManagedChannel) channel);
+    }
+    return null;
+  }
+
   @Override
   public GrpcCallContext withTransportChannel(TransportChannel inputChannel) {
     Preconditions.checkNotNull(inputChannel);

sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcTransportChannel.java

@@ -66,6 +66,14 @@ public Channel getChannel() {
     return getManagedChannel();
   }
 
+  @Override
+  public void refresh() {
+    Channel channel = getChannel();
+    if (channel instanceof ChannelPool) {
+      ((ChannelPool) channel).refresh();
+    }
+  }
+
   @Override
   public void shutdown() {
     getManagedChannel().shutdown();

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/BasicRetryingFuture.java

@@ -116,6 +116,11 @@ public TimedAttemptSettings getAttemptSettings() {
     }
   }
 
+  @Override
+  public RetryingContext getRetryingContext() {
+    return retryingContext;
+  }
+
   @Override
   public ApiFuture<ResponseT> peekAttemptResult() {
     synchronized (lock) {

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/RetryingFuture.java

@@ -128,4 +128,12 @@ public interface RetryingFuture<ResponseT> extends ApiFuture<ResponseT> {
    * </ul>
    */
   ApiFuture<ResponseT> getAttemptResult();
+
+  /**
+   * Returns the retrying context associated with this future, or {@code null} if none is set.
+   */
+  @com.google.api.core.BetaApi("The surface for passing per operation state is not yet stable")
+  default RetryingContext getRetryingContext() {
+    return null;
+  }
 }

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/retrying/ScheduledRetryingExecutor.java

@@ -111,6 +111,10 @@ public RetryingFuture<ResponseT> createFuture(
    */
   @Override
   public ApiFuture<ResponseT> submit(RetryingFuture<ResponseT> retryingFuture) {
+    if ("true".equalsIgnoreCase(System.getenv("isMwlidEnvironment"))) {
+      checkForFailedChannelRefresh(retryingFuture);
+    }
+
     try {
       ListenableFuture<ResponseT> attemptFuture =
           scheduler.schedule(
@@ -122,4 +126,27 @@ public ApiFuture<ResponseT> submit(RetryingFuture<ResponseT> retryingFuture) {
       return ApiFutures.immediateFailedFuture(e);
     }
   }
+
+  private void checkForFailedChannelRefresh(RetryingFuture<ResponseT> retryingFuture) {
+    ApiFuture<ResponseT> lastAttemptResult = retryingFuture.peekAttemptResult();
+    if (lastAttemptResult != null && lastAttemptResult.isDone()) {
+      try {
+        lastAttemptResult.get();
+      } catch (java.util.concurrent.ExecutionException e) {
+        Throwable cause = e.getCause();
+        if (cause instanceof com.google.api.gax.rpc.UnauthenticatedException) {
+          RetryingContext context = retryingFuture.getRetryingContext();
+          if (context instanceof com.google.api.gax.rpc.ApiCallContext) {
+            com.google.api.gax.rpc.TransportChannel transportChannel =
+                ((com.google.api.gax.rpc.ApiCallContext) context).getTransportChannel();
+            if (transportChannel != null) {
+              transportChannel.refresh();
+            }
+          }
+        }
+      } catch (Exception ignored) {
+        // Ignore cancellations or interruptions
+      }
+    }
+  }
 }

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiCallContext.java

@@ -63,6 +63,14 @@ public interface ApiCallContext extends RetryingContext {
   /** Returns a new ApiCallContext with the given channel set. */
   ApiCallContext withTransportChannel(TransportChannel channel);
 
+  /**
+   * Returns the {@link TransportChannel} associated with this call context, or {@code null} if none
+   * is set.
+   */
+  default TransportChannel getTransportChannel() {
+    return null;
+  }
+
   /** Returns a new ApiCallContext with the given Endpoint Context. */
   ApiCallContext withEndpointContext(EndpointContext endpointContext);
 

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiResultRetryAlgorithm.java

@@ -38,6 +38,10 @@ class ApiResultRetryAlgorithm<ResponseT> extends BasicResultRetryAlgorithm<Respo
   /** Returns true if previousThrowable is an {@link ApiException} that is retryable. */
   @Override
   public boolean shouldRetry(Throwable previousThrowable, ResponseT previousResponse) {
+    if ("true".equalsIgnoreCase(System.getenv("isMwlidEnvironment"))
+        && previousThrowable instanceof UnauthenticatedException) {
+      return true;
+    }
     return (previousThrowable instanceof ApiException)
         && ((ApiException) previousThrowable).isRetryable();
   }
@@ -51,6 +55,10 @@ public boolean shouldRetry(Throwable previousThrowable, ResponseT previousRespon
   @Override
   public boolean shouldRetry(
       RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) {
+    if ("true".equalsIgnoreCase(System.getenv("isMwlidEnvironment"))
+        && previousThrowable instanceof UnauthenticatedException) {
+      return true;
+    }
     if (context.getRetryableCodes() != null) {
       // Ignore the isRetryable() value of the throwable if the RetryingContext has a specific list
       // of codes that should be retried.

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/TransportChannel.java

@@ -47,4 +47,12 @@ public interface TransportChannel extends BackgroundResource {
    * Returns an empty {@link ApiCallContext} that is compatible with this {@code TransportChannel}.
    */
   ApiCallContext getEmptyCallContext();
+
+  /**
+   * Refreshes or recreates the underlying network connections of this transport channel.
+   *
+   * <p>By default, this is a no-op for transports that do not require stateful connection lifecycle
+   * management.
+   */
+  default void refresh() {}
 }