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

@@ -97,6 +97,7 @@ public final class GrpcCallContext implements ApiCallContext {
   private final ApiCallContextOptions options;
   private final EndpointContext endpointContext;
   private final boolean isDirectPath;
+  @Nullable private final TransportChannel transportChannel;
 
   /** Returns an empty instance with a null channel and default {@link CallOptions}. */
   public static GrpcCallContext createDefault() {
@@ -113,7 +114,8 @@ public static GrpcCallContext createDefault() {
         null,
         null,
         null,
-        false);
+        false,
+        null);
   }
 
   /** Returns an instance with the given channel and {@link CallOptions}. */
@@ -131,7 +133,8 @@ public static GrpcCallContext of(Channel channel, CallOptions callOptions) {
         null,
         null,
         null,
-        false);
+        false,
+        null);
   }
 
   private GrpcCallContext(
@@ -147,7 +150,8 @@ private GrpcCallContext(
       @Nullable RetrySettings retrySettings,
       @Nullable Set<StatusCode.Code> retryableCodes,
       @Nullable EndpointContext endpointContext,
-      boolean isDirectPath) {
+      boolean isDirectPath,
+      @Nullable TransportChannel transportChannel) {
     this.channel = channel;
     this.credentials = credentials;
     Preconditions.checkNotNull(callOptions);
@@ -167,6 +171,7 @@ private GrpcCallContext(
     this.endpointContext =
         endpointContext == null ? EndpointContext.getDefaultInstance() : endpointContext;
     this.isDirectPath = isDirectPath;
+    this.transportChannel = transportChannel;
   }
 
   /**
@@ -208,7 +213,13 @@ public GrpcCallContext withCredentials(Credentials newCredentials) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
+  }
+
+  @Override
+  public TransportChannel getTransportChannel() {
+    return transportChannel;
   }
 
   @Override
@@ -232,7 +243,8 @@ public GrpcCallContext withTransportChannel(TransportChannel inputChannel) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        transportChannel.isDirectPath());
+        transportChannel.isDirectPath(),
+        inputChannel);
   }
 
   @Override
@@ -251,7 +263,8 @@ public GrpcCallContext withEndpointContext(EndpointContext endpointContext) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   /** This method is obsolete. Use {@link #withTimeoutDuration(java.time.Duration)} instead. */
@@ -286,7 +299,8 @@ public GrpcCallContext withTimeoutDuration(@Nullable java.time.Duration timeout)
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   /** This method is obsolete. Use {@link #getTimeoutDuration()} instead. */
@@ -335,7 +349,8 @@ public GrpcCallContext withStreamWaitTimeoutDuration(
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   /**
@@ -370,7 +385,8 @@ public GrpcCallContext withStreamIdleTimeoutDuration(
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   @BetaApi("The surface for channel affinity is not stable yet and may change in the future.")
@@ -388,7 +404,8 @@ public GrpcCallContext withChannelAffinity(@Nullable Integer affinity) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   @BetaApi("The surface for extra headers is not stable yet and may change in the future.")
@@ -410,7 +427,8 @@ public GrpcCallContext withExtraHeaders(Map<String, List<String>> extraHeaders)
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   @Override
@@ -433,7 +451,8 @@ public GrpcCallContext withRetrySettings(RetrySettings retrySettings) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   @Override
@@ -456,7 +475,8 @@ public GrpcCallContext withRetryableCodes(Set<StatusCode.Code> retryableCodes) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   @Override
@@ -558,7 +578,8 @@ public ApiCallContext merge(ApiCallContext inputCallContext) {
         newRetrySettings,
         newRetryableCodes,
         endpointContext,
-        newIsDirectPath);
+        newIsDirectPath,
+        transportChannel);
   }
 
   /** The {@link Channel} set on this context. */
@@ -641,7 +662,8 @@ public GrpcCallContext withChannel(Channel newChannel) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   /** Returns a new instance with the call options set to the given call options. */
@@ -659,7 +681,8 @@ public GrpcCallContext withCallOptions(CallOptions newCallOptions) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   public GrpcCallContext withRequestParamsDynamicHeaderOption(String requestParams) {
@@ -704,7 +727,8 @@ public <T> GrpcCallContext withOption(Key<T> key, T value) {
         retrySettings,
         retryableCodes,
         endpointContext,
-        isDirectPath);
+        isDirectPath,
+        transportChannel);
   }
 
   /** {@inheritDoc} */

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/AttemptListener.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.retrying;
+
+import com.google.api.core.BetaApi;
+
+/**
+ * A listener that is notified in between retry attempts, allowing custom transport-level
+ * side effects to be executed out-of-band.
+ *
+ * @param <ResponseT> response type
+ */
+@BetaApi("The surface for attempt listeners is not stable yet and may change in the future.")
+public interface AttemptListener<ResponseT> {
+  /** Called in between retry attempts before the next attempt task is scheduled or fired. */
+  void onAttemptComplete(RetryingFuture<ResponseT> retryingFuture);
+}

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

@@ -59,6 +59,7 @@ public class ScheduledRetryingExecutor<ResponseT>
 
   private final RetryAlgorithm<ResponseT> retryAlgorithm;
   private final ListeningScheduledExecutorService scheduler;
+  private final AttemptListener<ResponseT> attemptListener;
 
   /**
    * Creates a new scheduled retry executor, which will be using {@code scheduler} for actual
@@ -69,8 +70,24 @@ public class ScheduledRetryingExecutor<ResponseT>
    */
   public ScheduledRetryingExecutor(
       RetryAlgorithm<ResponseT> retryAlgorithm, ScheduledExecutorService scheduler) {
+    this(retryAlgorithm, scheduler, null);
+  }
+
+  /**
+   * Creates a new scheduled retry executor with a custom {@link AttemptListener}.
+   *
+   * @param retryAlgorithm retry algorithm to use
+   * @param scheduler scheduler
+   * @param attemptListener a listener to be notified in between retry attempts
+   */
+  @BetaApi("The surface for attempt listeners is not stable yet and may change in the future.")
+  public ScheduledRetryingExecutor(
+      RetryAlgorithm<ResponseT> retryAlgorithm,
+      ScheduledExecutorService scheduler,
+      AttemptListener<ResponseT> attemptListener) {
     this.retryAlgorithm = retryAlgorithm;
     this.scheduler = MoreExecutors.listeningDecorator(scheduler);
+    this.attemptListener = attemptListener;
   }
 
   /**
@@ -111,6 +128,10 @@ public RetryingFuture<ResponseT> createFuture(
    */
   @Override
   public ApiFuture<ResponseT> submit(RetryingFuture<ResponseT> retryingFuture) {
+    if (attemptListener != null) {
+      attemptListener.onAttemptComplete(retryingFuture);
+    }
+
     try {
       ListenableFuture<ResponseT> attemptFuture =
           scheduler.schedule(

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.