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

@@ -84,6 +84,28 @@ public ResponseT call() {
           .attemptStarted(request, externalFuture.getAttemptSettings().getOverallAttemptCount());
 
       ApiFuture<ResponseT> internalFuture = callable.futureCall(request, callContext);
+
+      if ("true".equalsIgnoreCase(System.getenv("isMwlidEnvironment"))) {
+        final ApiCallContext finalContext = callContext;
+        ApiFutures.addCallback(
+            internalFuture,
+            new com.google.api.core.ApiFutureCallback<ResponseT>() {
+              @Override
+              public void onFailure(Throwable t) {
+                if (t instanceof UnauthenticatedException) {
+                  TransportChannel transportChannel = finalContext.getTransportChannel();
+                  if (transportChannel != null) {
+                    transportChannel.refresh();
+                  }
+                }
+              }
+
+              @Override
+              public void onSuccess(ResponseT result) {}
+            },
+            com.google.common.util.concurrent.MoreExecutors.directExecutor());
+      }
+
       externalFuture.setAttemptFuture(internalFuture);
     } catch (Throwable e) {
       externalFuture.setAttemptFuture(ApiFutures.<ResponseT>immediateFailedFuture(e));

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

@@ -219,6 +219,7 @@ public Void call() {
         .getTracer()
         .attemptStarted(request, outerRetryingFuture.getAttemptSettings().getOverallAttemptCount());
 
+    final ApiCallContext finalContext = attemptContext;
     innerCallable.call(
         request,
         new StateCheckingResponseObserver<ResponseT>() {
@@ -234,6 +235,18 @@ public void onResponseImpl(ResponseT response) {
 
           @Override
           public void onErrorImpl(Throwable t) {
+            if ("true".equalsIgnoreCase(System.getenv("isMwlidEnvironment"))) {
+              Throwable cause = t;
+              if (cause instanceof com.google.api.gax.retrying.ServerStreamingAttemptException) {
+                cause = cause.getCause();
+              }
+              if (cause instanceof UnauthenticatedException) {
+                TransportChannel transportChannel = finalContext.getTransportChannel();
+                if (transportChannel != null) {
+                  transportChannel.refresh();
+                }
+              }
+            }
             onAttemptError(t);
           }
 

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() {}
 }