feat(grpc-gcp): move scale-up to background worker

Author: rahul2393

grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java

@@ -28,6 +28,7 @@
 import com.google.cloud.grpc.proto.MethodConfig;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
+import com.google.common.base.Preconditions;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.MessageOrBuilder;
@@ -65,13 +66,15 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.IntUnaryOperator;
@@ -162,7 +165,9 @@ private static int stateFromChannelId(int channelId) {
   private int minRpcPerChannel = 0;
   private int maxRpcPerChannel = 0;
   private Duration scaleDownInterval = Duration.ZERO;
+  private Duration scaleUpCooldown = Duration.ofSeconds(10);
   private int scaleDownConsecutiveLowLoadChecks = 3;
+  private int maxScaleUpPercent = 30;
   private int consecutiveLowLoadChecks;
   private int maxScaleDownChannels = 2;
   private boolean isDynamicScalingEnabled = false;
@@ -273,6 +278,12 @@ private static int stateFromChannelId(int channelId) {
   private AtomicLong maxUnresponsiveDrops = new AtomicLong();
   private AtomicLong scaleUpCount = new AtomicLong();
   private AtomicLong scaleDownCount = new AtomicLong();
+  private final AtomicBoolean scaleUpSignalPending = new AtomicBoolean();
+  private final AtomicBoolean scaleUpWorkerRunning = new AtomicBoolean();
+  private final AtomicLong scaleUpSignalVersion = new AtomicLong();
+  private Executor scaleUpExecutor = SHARED_BACKGROUND_SERVICE;
+  private volatile long lastScaleUpNanos = Long.MIN_VALUE;
+  private volatile boolean shuttingDown;
 
   // Clock supplier for nanoTime, injectable for testing.
   private Supplier<Long> nanoClock = System::nanoTime;
@@ -294,6 +305,11 @@ int readyChannelCountForTest() {
     return readyChannels.get();
   }
 
+  @VisibleForTesting
+  void setScaleUpExecutorForTest(Executor scaleUpExecutor) {
+    this.scaleUpExecutor = Preconditions.checkNotNull(scaleUpExecutor);
+  }
+
   @VisibleForTesting
   void setNanoClock(Supplier<Long> nanoClock) {
     this.nanoClock = nanoClock;
@@ -489,7 +505,9 @@ private void initOptions() {
       minRpcPerChannel = poolOptions.getMinRpcPerChannel();
       maxRpcPerChannel = poolOptions.getMaxRpcPerChannel();
       scaleDownInterval = poolOptions.getScaleDownInterval();
+      scaleUpCooldown = poolOptions.getScaleUpCooldown();
       scaleDownConsecutiveLowLoadChecks = poolOptions.getScaleDownConsecutiveLowLoadChecks();
+      maxScaleUpPercent = poolOptions.getMaxScaleUpPercent();
       maxScaleDownChannels = poolOptions.getMaxScaleDownChannels();
       isDynamicScalingEnabled =
           minRpcPerChannel > 0 && maxRpcPerChannel > 0 && !scaleDownInterval.isZero();
@@ -517,7 +535,13 @@ private synchronized void initScaleDownChecker(Duration scaleDownInterval) {
 
     scaleDownTask =
         SHARED_BACKGROUND_SERVICE.scheduleAtFixedRate(
-            this::checkScaleDown,
+            () -> {
+              try {
+                checkScaleDown();
+              } catch (Throwable failure) {
+                logger.log(Level.WARNING, log("Scale-down check failed"), failure);
+              }
+            },
             scaleDownInterval.toMillis(),
             scaleDownInterval.toMillis(),
             MILLISECONDS);
@@ -1650,11 +1674,19 @@ public int getStreamsLowWatermark() {
   }
 
   public int getMinActiveStreams() {
-    return channelRefs.stream().mapToInt(ChannelRef::getActiveStreamsCount).min().orElse(0);
+    return channelRefs.stream()
+        .filter(ChannelRef::isActive)
+        .mapToInt(ChannelRef::getActiveStreamsCount)
+        .min()
+        .orElse(0);
   }
 
   public int getMaxActiveStreams() {
-    return channelRefs.stream().mapToInt(ChannelRef::getActiveStreamsCount).max().orElse(0);
+    return channelRefs.stream()
+        .filter(ChannelRef::isActive)
+        .mapToInt(ChannelRef::getActiveStreamsCount)
+        .max()
+        .orElse(0);
   }
 
   /**
@@ -1691,12 +1723,15 @@ protected synchronized ChannelRef getChannelRefRoundRobin() {
     if (!isDynamicScalingEnabled && channelRefs.size() < maxSize) {
       return createNewChannel();
     }
-    maybeDynamicUpscale();
-    bindingIndex++;
-    if (bindingIndex >= channelRefs.size()) {
-      bindingIndex = 0;
+    Object[] snapshot = channelRefs.toArray();
+    for (int attempts = 0; attempts < snapshot.length; attempts++) {
+      bindingIndex = (bindingIndex + 1) % snapshot.length;
+      ChannelRef candidate = (ChannelRef) snapshot[bindingIndex];
+      if (candidate.isActive()) {
+        return candidate;
+      }
     }
-    return channelRefs.get(bindingIndex);
+    return pickFromCandidates(channelRefs);
   }
 
   /**
@@ -1710,7 +1745,6 @@ protected synchronized ChannelRef getChannelRefRoundRobin() {
    *     Otherwise pick the one with the smallest number of streams.
    */
   protected ChannelRef getChannelRef(@Nullable String key) {
-    maybeDynamicUpscale();
     if (key == null || key.isEmpty()) {
       return pickLeastBusyChannel(/* forFallback= */ false);
     }
@@ -1777,7 +1811,6 @@ protected ChannelRef getChannelRef(@Nullable String key) {
    * Pick a {@link ChannelRef} using a caller-owned reference instead of grpc-gcp's affinity map.
    */
   protected ChannelRef getChannelRefByAffinityRef(ChannelAffinityRef affinityRef) {
-    maybeDynamicUpscale();
     // Retry if another thread updates the caller-owned affinity ref while we are picking a channel.
     while (true) {
       int state = affinityRef.state.get();
@@ -1894,24 +1927,158 @@ private ChannelRef tryCreateNewChannel() {
     return null;
   }
 
-  private void maybeDynamicUpscale() {
-    if (!isDynamicScalingEnabled || channelRefs.size() >= maxSize) {
+  private void maybeSignalScaleUp(ChannelRef selectedChannel) {
+    int activeChannels = channelRefs.size();
+    if (!selectedChannel.isActive()
+        || !isDynamicScalingEnabled
+        || shuttingDown
+        || activeChannels == 0
+        || activeChannels >= maxSize) {
       return;
     }
-
-    if ((totalActiveStreams.get() / channelRefs.size()) >= maxRpcPerChannel) {
-      dynamicUpscale();
+    if (selectedChannel.getActiveStreamsCount() <= maxRpcPerChannel
+        && ((double) totalActiveStreams.get() / activeChannels) <= maxRpcPerChannel) {
+      return;
     }
+    signalScaleUp();
   }
 
-  private synchronized void dynamicUpscale() {
-    if (!isDynamicScalingEnabled || channelRefs.size() >= maxSize) {
+  private void signalScaleUp() {
+    long signalVersion = scaleUpSignalVersion.incrementAndGet();
+    scaleUpSignalPending.set(true);
+    if (!scaleUpWorkerRunning.compareAndSet(false, true)) {
       return;
     }
+    submitScaleUpWorker(signalVersion);
+  }
 
-    if ((totalActiveStreams.get() / channelRefs.size()) >= maxRpcPerChannel) {
-      createNewChannel();
-      scaleUpCount.incrementAndGet();
+  private void submitScaleUpWorker(long signalVersion) {
+    try {
+      scaleUpExecutor.execute(this::runScaleUpWorker);
+    } catch (Throwable failure) {
+      scaleUpWorkerRunning.set(false);
+      if (scaleUpSignalVersion.get() == signalVersion) {
+        scaleUpSignalPending.set(false);
+      } else if (!shuttingDown && scaleUpWorkerRunning.compareAndSet(false, true)) {
+        submitScaleUpWorker(scaleUpSignalVersion.get());
+      }
+      logger.log(Level.FINE, log("Scale-up task rejected"), failure);
+    }
+  }
+
+  private void runScaleUpWorker() {
+    try {
+      do {
+        scaleUpSignalPending.set(false);
+        try {
+          dynamicUpscale();
+        } catch (Throwable failure) {
+          logger.log(Level.WARNING, log("Scale-up failed"), failure);
+        }
+      } while (scaleUpSignalPending.get() && !shuttingDown);
+    } finally {
+      scaleUpWorkerRunning.set(false);
+      if (scaleUpSignalPending.get() && !shuttingDown) {
+        signalScaleUp();
+      }
+    }
+  }
+
+  private void dynamicUpscale() {
+    final int channelsToBuild;
+    final long scaleUpNanos;
+    int reused = 0;
+    synchronized (this) {
+      if (!isDynamicScalingEnabled || shuttingDown || channelRefs.size() >= maxSize) {
+        return;
+      }
+      long now = nanoClock.get();
+      if (lastScaleUpNanos != Long.MIN_VALUE
+          && now - lastScaleUpNanos < scaleUpCooldown.toNanos()) {
+        return;
+      }
+      int active = channelRefs.size();
+      int targetRpcPerChannel = Math.max(1, (minRpcPerChannel + maxRpcPerChannel) / 2);
+      long load = totalActiveStreams.get();
+      int desired =
+          load == 0
+              ? active
+              : (int) Math.min(Integer.MAX_VALUE, 1 + ((load - 1) / targetRpcPerChannel));
+      int percentCap = Math.max(2, (int) (1 + (((long) active * maxScaleUpPercent - 1) / 100)));
+      int add = Math.min(Math.max(0, desired - active), percentCap);
+      add = Math.min(add, maxSize - active);
+      while (reused < add) {
+        Optional<ChannelRef> reusable = pickChannelForReuse();
+        if (!reusable.isPresent()) {
+          break;
+        }
+        ChannelRef channelRef = reusable.get();
+        removedChannelRefs.remove(channelRef);
+        channelRefs.add(channelRef);
+        channelIdToChannelRef.put(channelRef.getId(), channelRef);
+        channelRef.activateAndAccountReadiness();
+        reused++;
+        scaleUpCount.incrementAndGet();
+        lastScaleUpNanos = now;
+        maxChannels.accumulateAndGet(getNumberOfChannels(), Math::max);
+      }
+      channelsToBuild = add - reused;
+      scaleUpNanos = now;
+    }
+
+    List<ManagedChannel> builtChannels = new ArrayList<>(channelsToBuild);
+    try {
+      for (int i = 0; i < channelsToBuild; i++) {
+        builtChannels.add(
+            Preconditions.checkNotNull(
+                delegateChannelBuilder.build(), "Delegate channel builder returned null."));
+      }
+    } catch (Throwable failure) {
+      builtChannels.forEach(this::shutdownUnpublishedChannel);
+      throw failure;
+    }
+
+    for (int i = 0; i < builtChannels.size(); i++) {
+      ManagedChannel channel = builtChannels.get(i);
+      try {
+        boolean published = false;
+        synchronized (this) {
+          if (!shuttingDown && channelRefs.size() < maxSize) {
+            ChannelRef channelRef = new ChannelRef(channel);
+            channelRefs.add(channelRef);
+            try {
+              channelIdToChannelRef.put(channelRef.getId(), channelRef);
+              channelRef.activateAndAccountReadiness();
+              published = true;
+            } catch (Throwable failure) {
+              channelRefs.remove(channelRef);
+              channelIdToChannelRef.remove(channelRef.getId(), channelRef);
+              throw failure;
+            }
+          }
+        }
+        if (published) {
+          scaleUpCount.incrementAndGet();
+          lastScaleUpNanos = scaleUpNanos;
+          maxChannels.accumulateAndGet(getNumberOfChannels(), Math::max);
+        } else {
+          shutdownUnpublishedChannel(channel);
+        }
+      } catch (Throwable failure) {
+        shutdownUnpublishedChannel(channel);
+        for (int j = i + 1; j < builtChannels.size(); j++) {
+          shutdownUnpublishedChannel(builtChannels.get(j));
+        }
+        throw failure;
+      }
+    }
+  }
+
+  private void shutdownUnpublishedChannel(ManagedChannel channel) {
+    try {
+      channel.shutdownNow();
+    } catch (Throwable failure) {
+      logger.log(Level.WARNING, log("Failed to close unpublished scale-up channel"), failure);
     }
   }
 
@@ -1982,13 +2149,16 @@ private ChannelRef pickLeastBusyNoFallback() {
   private ChannelRef pickLeastBusyWithFallback(boolean forFallback) {
     // Full scan to collect eligible ("ready") channels not in fallbackMap and under max streams.
     List<ChannelRef> readyCandidates = new ArrayList<>();
-    ChannelRef overallCandidate = channelRefs.get(0);
-    int overallMinStreams = overallCandidate.getActiveStreamsCount();
+    ChannelRef overallCandidate = null;
+    int overallMinStreams = Integer.MAX_VALUE;
     int readyMaxStreams = 0;
 
     for (ChannelRef channelRef : channelRefs) {
+      if (!channelRef.isActive()) {
+        continue;
+      }
       int cnt = channelRef.getActiveStreamsCount();
-      if (cnt < overallMinStreams) {
+      if (overallCandidate == null || cnt < overallMinStreams) {
         overallMinStreams = cnt;
         overallCandidate = channelRef;
       }
@@ -2000,6 +2170,10 @@ private ChannelRef pickLeastBusyWithFallback(boolean forFallback) {
       }
     }
 
+    if (overallCandidate == null) {
+      return pickFromCandidates(channelRefs);
+    }
+
     // For scale-up, use maxStreams among ready channels (consistent with non-fallback path).
     int scaleUpStreams = readyCandidates.isEmpty() ? Integer.MAX_VALUE : readyMaxStreams;
     if (shouldScaleUp(scaleUpStreams)) {
@@ -2084,8 +2258,10 @@ ChannelRef pickFromCandidates(List<ChannelRef> candidates) {
 
   @Override
   public String authority() {
-    if (!channelRefs.isEmpty()) {
-      return channelRefs.get(0).getChannel().authority();
+    for (ChannelRef channelRef : channelRefs) {
+      if (channelRef.isActive()) {
+        return channelRef.getChannel().authority();
+      }
     }
     final ManagedChannel channel = delegateChannelBuilder.build();
     final String authority = channel.authority();
@@ -2163,6 +2339,8 @@ private synchronized void cancelBackgroundTasks() {
   @Override
   public ManagedChannel shutdownNow() {
     logger.finer(log("Shutdown now started."));
+    shuttingDown = true;
+    scaleUpSignalPending.set(false);
     for (ChannelRef channelRef : channelRefs) {
       if (!channelRef.getChannel().isTerminated()) {
         channelRef.getChannel().shutdownNow();
@@ -2183,6 +2361,8 @@ public ManagedChannel shutdownNow() {
   @Override
   public ManagedChannel shutdown() {
     logger.finer(log("Shutdown started."));
+    shuttingDown = true;
+    scaleUpSignalPending.set(false);
     for (ChannelRef channelRef : channelRefs) {
       channelRef.getChannel().shutdown();
     }
@@ -2540,6 +2720,7 @@ protected void activeStreamsCountIncr() {
       maxActiveStreams.accumulateAndGet(actStreams, Math::max);
       int totalActStreams = totalActiveStreams.incrementAndGet();
       maxTotalActiveStreams.accumulateAndGet(totalActStreams, Math::max);
+      maybeSignalScaleUp(this);
     }
 
     protected void activeStreamsCountDecr(long startNanos, Status status, boolean fromClientSide) {
@@ -2572,6 +2753,12 @@ protected int getActiveStreamsCount() {
       return activeStreamsCount.get();
     }
 
+    @VisibleForTesting
+    void setActiveStreamsForTest(int streams) {
+      int previous = activeStreamsCount.getAndSet(streams);
+      totalActiveStreams.addAndGet(streams - previous);
+    }
+
     protected long getAndResetOkCalls() {
       return okCalls.getAndSet(0);
     }