fix: provide API to share the same background executor for channel po… (#4030)

…ol resizing and refresh

Before this change ChannelPool creates a new
SingleThreadScheduledExecutor every time getTransportChannel is called.
We can reuse the same background executor from the clientSettings and
can reduce the number of threads when the user creates many instances of
the client.

---------

Co-authored-by: Blake Li <[email protected]>

Author: mutianf

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

@@ -30,6 +30,7 @@
 package com.google.api.gax.grpc;
 
 import com.google.api.core.InternalApi;
+import com.google.api.gax.core.FixedExecutorProvider;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
@@ -48,6 +49,7 @@
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -72,33 +74,45 @@ class ChannelPool extends ManagedChannel {
 
   private final ChannelPoolSettings settings;
   private final ChannelFactory channelFactory;
-  private final ScheduledExecutorService executor;
+  private final FixedExecutorProvider backgroundExecutorProvider;
+
+  private ScheduledFuture<?> refreshFuture = null;
+  private ScheduledFuture<?> resizeFuture = null;
 
   private final Object entryWriteLock = new Object();
   @VisibleForTesting final AtomicReference<ImmutableList<Entry>> entries = new AtomicReference<>();
   private final AtomicInteger indexTicker = new AtomicInteger();
   private final String authority;
 
-  static ChannelPool create(ChannelPoolSettings settings, ChannelFactory channelFactory)
+  static ChannelPool create(
+      ChannelPoolSettings settings,
+      ChannelFactory channelFactory,
+      @Nullable ScheduledExecutorService backgroundExecutor)
       throws IOException {
-    return new ChannelPool(settings, channelFactory, Executors.newSingleThreadScheduledExecutor());
+
+    FixedExecutorProvider executorProvider =
+        backgroundExecutor == null
+            ? FixedExecutorProvider.create(Executors.newSingleThreadScheduledExecutor(), true)
+            : FixedExecutorProvider.create(backgroundExecutor, false);
+    return new ChannelPool(settings, channelFactory, executorProvider);
   }
 
   /**
    * Initializes the channel pool. Assumes that all channels have the same authority.
    *
    * @param settings options for controling the ChannelPool sizing behavior
    * @param channelFactory method to create the channels
-   * @param executor periodically refreshes the channels
+   * @param executorProvider provides the executor that periodically refreshes the channels
    */
   @VisibleForTesting
   ChannelPool(
       ChannelPoolSettings settings,
       ChannelFactory channelFactory,
-      ScheduledExecutorService executor)
+      FixedExecutorProvider executorProvider)
       throws IOException {
     this.settings = settings;
     this.channelFactory = channelFactory;
+    this.backgroundExecutorProvider = executorProvider;
 
     ImmutableList.Builder<Entry> initialListBuilder = ImmutableList.builder();
 
@@ -108,21 +122,26 @@ static ChannelPool create(ChannelPoolSettings settings, ChannelFactory channelFa
 
     entries.set(initialListBuilder.build());
     authority = entries.get().get(0).channel.authority();
-    this.executor = executor;
 
     if (!settings.isStaticSize()) {
-      executor.scheduleAtFixedRate(
-          this::resizeSafely,
-          ChannelPoolSettings.RESIZE_INTERVAL.getSeconds(),
-          ChannelPoolSettings.RESIZE_INTERVAL.getSeconds(),
-          TimeUnit.SECONDS);
+      resizeFuture =
+          backgroundExecutorProvider
+              .getExecutor()
+              .scheduleAtFixedRate(
+                  this::resizeSafely,
+                  ChannelPoolSettings.RESIZE_INTERVAL.getSeconds(),
+                  ChannelPoolSettings.RESIZE_INTERVAL.getSeconds(),
+                  TimeUnit.SECONDS);
     }
     if (settings.isPreemptiveRefreshEnabled()) {
-      executor.scheduleAtFixedRate(
-          this::refreshSafely,
-          REFRESH_PERIOD.getSeconds(),
-          REFRESH_PERIOD.getSeconds(),
-          TimeUnit.SECONDS);
+      refreshFuture =
+          backgroundExecutorProvider
+              .getExecutor()
+              .scheduleAtFixedRate(
+                  this::refreshSafely,
+                  REFRESH_PERIOD.getSeconds(),
+                  REFRESH_PERIOD.getSeconds(),
+                  TimeUnit.SECONDS);
     }
   }
 
@@ -153,14 +172,25 @@ Channel getChannel(int affinity) {
   public ManagedChannel shutdown() {
     LOG.fine("Initiating graceful shutdown due to explicit request");
 
+    // Resize and refresh tasks can block on channel priming. We don't need
+    // to wait for the channels to be ready since we're shutting down the
+    // pool. Allowing interrupt to speed it up.
+    if (resizeFuture != null) {
+      resizeFuture.cancel(true);
+    }
+    if (refreshFuture != null) {
+      refreshFuture.cancel(true);
+    }
+
     List<Entry> localEntries = entries.get();
     for (Entry entry : localEntries) {
       entry.channel.shutdown();
     }
-    if (executor != null) {
-      // shutdownNow will cancel scheduled tasks
-      executor.shutdownNow();
+
+    if (backgroundExecutorProvider.shouldAutoClose()) {
+      backgroundExecutorProvider.getExecutor().shutdown();
     }
+
     return this;
   }
 
@@ -173,7 +203,7 @@ public boolean isShutdown() {
         return false;
       }
     }
-    return executor == null || executor.isShutdown();
+    return true;
   }
 
   /** {@inheritDoc} */
@@ -185,21 +215,28 @@ public boolean isTerminated() {
         return false;
       }
     }
-
-    return executor == null || executor.isTerminated();
+    return true;
   }
 
   /** {@inheritDoc} */
   @Override
   public ManagedChannel shutdownNow() {
     LOG.fine("Initiating immediate shutdown due to explicit request");
 
+    if (resizeFuture != null) {
+      resizeFuture.cancel(true);
+    }
+    if (refreshFuture != null) {
+      refreshFuture.cancel(true);
+    }
+
     List<Entry> localEntries = entries.get();
     for (Entry entry : localEntries) {
       entry.channel.shutdownNow();
     }
-    if (executor != null) {
-      executor.shutdownNow();
+
+    if (backgroundExecutorProvider.shouldAutoClose()) {
+      backgroundExecutorProvider.getExecutor().shutdownNow();
     }
     return this;
   }
@@ -216,10 +253,6 @@ public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedE
       }
       entry.channel.awaitTermination(awaitTimeNanos, TimeUnit.NANOSECONDS);
     }
-    if (executor != null) {
-      long awaitTimeNanos = endTimeNanos - System.nanoTime();
-      executor.awaitTermination(awaitTimeNanos, TimeUnit.NANOSECONDS);
-    }
     return isTerminated();
   }
 

gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java

@@ -129,6 +129,7 @@ public final class InstantiatingGrpcChannelProvider implements TransportChannelP
 
   private final int processorCount;
   private final Executor executor;
+  @Nullable private final ScheduledExecutorService backgroundExecutor;
   private final HeaderProvider headerProvider;
   private final boolean useS2A;
   private final String endpoint;
@@ -181,6 +182,7 @@ public enum HardBoundTokenTypes {
   private InstantiatingGrpcChannelProvider(Builder builder) {
     this.processorCount = builder.processorCount;
     this.executor = builder.executor;
+    this.backgroundExecutor = builder.backgroundExecutor;
     this.headerProvider = builder.headerProvider;
     this.useS2A = builder.useS2A;
     this.endpoint = builder.endpoint;
@@ -245,6 +247,16 @@ public TransportChannelProvider withExecutor(Executor executor) {
     return toBuilder().setExecutor(executor).build();
   }
 
+  @Override
+  public boolean needsBackgroundExecutor() {
+    return backgroundExecutor == null;
+  }
+
+  @Override
+  public TransportChannelProvider withBackgroundExecutor(ScheduledExecutorService executor) {
+    return toBuilder().setBackgroundExecutor(executor).build();
+  }
+
   @Override
   public boolean needsHeaders() {
     return headerProvider == null;
@@ -356,7 +368,9 @@ private TransportChannel createChannel() throws IOException {
     return GrpcTransportChannel.newBuilder()
         .setManagedChannel(
             ChannelPool.create(
-                channelPoolSettings, InstantiatingGrpcChannelProvider.this::createSingleChannel))
+                channelPoolSettings,
+                InstantiatingGrpcChannelProvider.this::createSingleChannel,
+                backgroundExecutor))
         .setDirectPath(this.canUseDirectPath())
         .build();
   }
@@ -839,6 +853,11 @@ public ChannelPoolSettings getChannelPoolSettings() {
     return channelPoolSettings;
   }
 
+  /** Gets the background executor for channel refresh and resize. */
+  ScheduledExecutorService getBackgroundExecutor() {
+    return backgroundExecutor;
+  }
+
   @Override
   public boolean shouldAutoClose() {
     return true;
@@ -855,6 +874,7 @@ public static Builder newBuilder() {
   public static final class Builder {
     @Deprecated private int processorCount;
     private Executor executor;
+    private ScheduledExecutorService backgroundExecutor;
     private HeaderProvider headerProvider;
     private String endpoint;
     private String mtlsEndpoint;
@@ -891,6 +911,7 @@ private Builder() {
     private Builder(InstantiatingGrpcChannelProvider provider) {
       this.processorCount = provider.processorCount;
       this.executor = provider.executor;
+      this.backgroundExecutor = provider.backgroundExecutor;
       this.headerProvider = provider.headerProvider;
       this.endpoint = provider.endpoint;
       this.useS2A = provider.useS2A;
@@ -950,6 +971,15 @@ public Builder setExecutorProvider(ExecutorProvider executorProvider) {
       return setExecutor((Executor) executorProvider.getExecutor());
     }
 
+    /**
+     * Sets the background executor for this TransportChannelProvider. The life cycle of the
+     * executor should be managed by the caller.
+     */
+    public Builder setBackgroundExecutor(ScheduledExecutorService executor) {
+      this.backgroundExecutor = executor;
+      return this;
+    }
+
     /**
      * Sets the HeaderProvider for this TransportChannelProvider.
      *