From 4cc4447e01d797c7ee0b045391943734aa4f352f Mon Sep 17 00:00:00 2001 From: Rahul Yadav Date: Thu, 27 Aug 2026 19:10:53 +0530 Subject: [PATCH] feat(grpc-gcp): drain scaled-down channels --- .../google/cloud/grpc/GcpManagedChannel.java | 577 ++++++++++++------ .../cloud/grpc/GcpManagedChannelOptions.java | 31 +- .../grpc/GcpManagedChannelDrainingTest.java | 347 +++++++++++ ...anagedChannelHotChannelReproducerTest.java | 11 +- .../grpc/GcpManagedChannelOptionsTest.java | 1 + .../GcpManagedChannelScaleUpWorkerTest.java | 12 + .../cloud/grpc/GcpManagedChannelTest.java | 511 ++++++++++------ 7 files changed, 1135 insertions(+), 355 deletions(-) create mode 100644 grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelDrainingTest.java diff --git a/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java b/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java index 7c1ab00e5325..54cf48920cfc 100644 --- a/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java +++ b/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannel.java @@ -60,13 +60,13 @@ import java.util.List; import java.util.LongSummaryStatistics; import java.util.Map; -import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.ThreadLocalRandom; @@ -74,6 +74,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; import java.util.function.IntUnaryOperator; import java.util.function.Supplier; import java.util.logging.Level; @@ -112,7 +113,12 @@ public class GcpManagedChannel extends ManagedChannel { public static final CallOptions.Key CHANNEL_AFFINITY_REF_KEY = CallOptions.Key.create("GcpChannelAffinityRef"); - /** Opaque sticky channel reference for callers that should not depend on {@link ChannelRef}. */ + /** + * Opaque caller-owned channel reference for transaction-lifetime stickiness. + * + *

The reference remains on a draining channel until its delegate shuts down. Call {@link + * #useDifferentChannelOnNextCall()} to move the next RPC to another active channel. + */ public static final class ChannelAffinityRef { private static final int USE_DIFFERENT_CHANNEL_ON_NEXT_CALL_MASK = 1 << 31; private static final int CHANNEL_ID_MASK = ~USE_DIFFERENT_CHANNEL_ON_NEXT_CALL_MASK; @@ -166,6 +172,7 @@ private static int stateFromChannelId(int channelId) { private int scaleDownConsecutiveLowLoadChecks = 3; private int maxScaleUpPercent = 30; private int maxScaleDownChannels = 2; + private Duration drainIdleGrace = Duration.ofMinutes(1); private boolean isDynamicScalingEnabled = false; private int maxConcurrentStreamsLowWatermark = DEFAULT_MAX_STREAM; private GcpManagedChannelOptions.ChannelPickStrategy channelPickStrategy = @@ -189,6 +196,9 @@ private static int stateFromChannelId(int channelId) { // we can shut them down. final Set removedChannelRefs = ConcurrentHashMap.newKeySet(); + @GuardedBy("this") + private final Map> drainTasks = new HashMap<>(); + // One-slot scale-up signal. At most one worker mutates pool size at a time. private final AtomicBoolean scaleUpSignalPending = new AtomicBoolean(); private final AtomicBoolean scaleUpWorkerRunning = new AtomicBoolean(); @@ -242,6 +252,7 @@ private static int stateFromChannelId(int channelId) { private ScheduledFuture cleanupTask; private ScheduledFuture scaleDownTask; private ScheduledFuture logMetricsTask; + private ScheduledExecutorService drainScheduler = SHARED_BACKGROUND_SERVICE; // Metrics counters. private final AtomicInteger readyChannels = new AtomicInteger(); @@ -287,17 +298,43 @@ private static int stateFromChannelId(int channelId) { private Supplier nanoClock = System::nanoTime; private IntUnaryOperator candidateIndexPicker = bound -> ThreadLocalRandom.current().nextInt(bound); + @Nullable private volatile Consumer pickerValidationHookForTest; + @Nullable private volatile Runnable inactiveMappingRemovalHookForTest; @VisibleForTesting void setNanoClock(Supplier nanoClock) { this.nanoClock = nanoClock; } + @VisibleForTesting + void setPickerValidationHookForTest(Consumer hook) { + pickerValidationHookForTest = hook; + } + @VisibleForTesting void setCandidateIndexPickerForTest(IntUnaryOperator candidateIndexPicker) { this.candidateIndexPicker = candidateIndexPicker; } + @VisibleForTesting + void setInactiveMappingRemovalHookForTest(Runnable hook) { + inactiveMappingRemovalHookForTest = hook; + } + + @VisibleForTesting + void setDrainSchedulerForTest(ScheduledExecutorService drainScheduler) { + this.drainScheduler = drainScheduler; + } + + private boolean validatePickedChannel(ChannelRef channelRef) { + Consumer hook = pickerValidationHookForTest; + if (hook != null) { + pickerValidationHookForTest = null; + hook.accept(channelRef); + } + return channelRef.isActive(); + } + @VisibleForTesting Map> fallbackMapForTest() { return fallbackMap; @@ -308,6 +345,11 @@ int readyChannelCountForTest() { return readyChannels.get(); } + @VisibleForTesting + synchronized int drainTaskCountForTest() { + return drainTasks.size(); + } + private static ScheduledThreadPoolExecutor createSharedBackgroundService() { ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor( @@ -383,8 +425,9 @@ public GcpManagedChannel( } } - private void cleanupAffinityKeys() { - final long cutoff = System.nanoTime() - affinityKeyLifetime.toNanos(); + @VisibleForTesting + void cleanupAffinityKeys() { + final long cutoff = nanoClock.get() - affinityKeyLifetime.toNanos(); affinityKeyLastUsed.forEach( (String key, Long time) -> { if (time < cutoff) { @@ -393,81 +436,224 @@ private void cleanupAffinityKeys() { }); } + /** + * Evaluates instantaneous active load; consecutive checks provide the low-load history rather + * than retaining a maximum observed between checks. + */ @VisibleForTesting - synchronized void checkScaleDown() { - if (!isDynamicScalingEnabled) { - return; - } + void checkScaleDown() { + List removedChannels; + synchronized (this) { + if (!isDynamicScalingEnabled || shuttingDown) { + return; + } - int channelCount = channelRefs.size(); - if (channelCount <= minSize) { - consecutiveLowLoadChecks = 0; - } else { - long activeLoad = totalActiveStreams.get(); - if (activeLoad > (long) minRpcPerChannel * channelCount) { + int channelCount = channelRefs.size(); + if (channelCount <= minSize) { consecutiveLowLoadChecks = 0; - } else if (++consecutiveLowLoadChecks >= scaleDownConsecutiveLowLoadChecks) { + return; + } + long activeLoad = activeLoad(channelRefs); + if (activeLoad > (long) minRpcPerChannel * channelCount) { consecutiveLowLoadChecks = 0; - int targetRpcPerChannel = Math.max(1, (minRpcPerChannel + maxRpcPerChannel) / 2); - int desiredSize = - activeLoad == 0 - ? minSize - : (int) Math.min(Integer.MAX_VALUE, 1 + ((activeLoad - 1) / targetRpcPerChannel)); - int removeCount = - Math.min( - maxScaleDownChannels, Math.max(0, channelCount - Math.max(minSize, desiredSize))); - removeOldestChannels(removeCount); - } - } - - // Shutdown removed channels where all RPCs are completed. - List completedChRefs = - removedChannelRefs.stream() - .filter(chRef -> (chRef.getActiveStreamsCount() == 0)) + return; + } + if (++consecutiveLowLoadChecks < scaleDownConsecutiveLowLoadChecks) { + return; + } + consecutiveLowLoadChecks = 0; + + int desiredSize = Math.max(minSize, ceilDiv(activeLoad, targetRpcPerChannel())); + int removeCount = Math.min(maxScaleDownChannels, Math.max(0, channelCount - desiredSize)); + removedChannels = removeChannels(removeCount); + } + List keysToUnbind = + affinityKeyToChannelRef.entrySet().stream() + .filter(entry -> removedChannels.contains(entry.getValue())) + .map(Map.Entry::getKey) .collect(Collectors.toList()); - removedChannelRefs.removeAll(completedChRefs); - for (ChannelRef channelRef : completedChRefs) { - channelRef.getChannel().shutdown(); - // Remove channel from broken channels map. - fallbackMap.remove(channelRef.getId()); - channelIdToChannelRef.remove(channelRef.getId()); + for (String key : keysToUnbind) { + ChannelRef mappedChannel = affinityKeyToChannelRef.get(key); + if (removedChannels.contains(mappedChannel)) { + unbindInactiveMapping(key, mappedChannel); + } + } + for (ChannelRef channelRef : removedChannels) { + scheduleDrain(channelRef); } } - private void removeOldestChannels(int num) { + @GuardedBy("this") + private List removeChannels(int num) { if (num <= 0) { - return; + return Collections.emptyList(); } - // Select longest connected channels (or disconnected channels). + // Drain least-loaded channels first, preferring fewer affinity bindings, then older channels. final List channelsToRemove = channelRefs.stream() - .sorted(Comparator.comparing(ChannelRef::getConnectedSinceNanos)) + .sorted( + Comparator.comparingInt(ChannelRef::getActiveStreamsCount) + .thenComparingInt(ChannelRef::getAffinityCount) + .thenComparingLong(ChannelRef::getCreatedNanos)) .limit(num) .collect(Collectors.toList()); - // Remove from active channels. + for (ChannelRef channelRef : channelsToRemove) { + // Stop new picks before publishing the shorter active list. + channelRef.deactivateAndAccountReadiness(); + } channelRefs.removeAll(channelsToRemove); for (ChannelRef channelRef : channelsToRemove) { - channelRef.resetAffinityCount(); - channelRef.deactivateAndAccountReadiness(); + removedChannelRefs.add(channelRef); } + minChannels.accumulateAndGet(getNumberOfChannels(), Math::min); + scaleDownCount.addAndGet(channelsToRemove.size()); + executeStateChangeCallbacks(); + return channelsToRemove; + } - // Remove affinity keys mapping for the channels. - affinityKeyToChannelRef - .keySet() - .removeIf(key -> channelsToRemove.contains(affinityKeyToChannelRef.get(key))); + /** Drain task bookkeeping is guarded by the pool monitor. */ + @VisibleForTesting + void scheduleDrain(ChannelRef channelRef) { + boolean closeChannel = false; + @Nullable DrainTask inlineTask = null; + synchronized (this) { + if (channelRef.isActive() || channelRef.getActiveStreamsCount() != 0 || shuttingDown) { + return; + } + long elapsed = Math.max(0, nanoClock.get() - channelRef.getLastActivityNanos()); + long delay = Math.max(0, drainIdleGrace.toNanos() - elapsed); + DrainTask drainTask = new DrainTask(channelRef); + ScheduledFuture task; + try { + task = drainScheduler.schedule(drainTask, delay, NANOSECONDS); + drainTask.future = task; + } catch (RejectedExecutionException e) { + logger.fine(log("Drain task rejected: %s", e.getMessage())); + closeChannel = removeDrainedChannel(channelRef); + task = null; + } + if (task != null) { + ScheduledFuture previous = drainTasks.put(channelRef, task); + if (previous != null) { + previous.cancel(false); + } + if (drainTask.ranBeforeFutureAssignment) { + inlineTask = drainTask; + } + } + } + if (inlineTask != null) { + finishDrain(channelRef, inlineTask); + } else if (closeChannel) { + channelRef.getChannel().shutdown(); + } + } - // Keep them aside to wait for all RPCs to complete. - removedChannelRefs.addAll(channelsToRemove); + @VisibleForTesting + void finishDrain(ChannelRef channelRef) { + finishDrain(channelRef, null); + } - // Track minimum number of channels for metrics. - minChannels.accumulateAndGet(getNumberOfChannels(), Math::min); - scaleDownCount.addAndGet(channelsToRemove.size()); + private void finishDrain(ChannelRef channelRef, @Nullable DrainTask runningTask) { + boolean closeChannel = false; + boolean reschedule = false; + synchronized (this) { + ScheduledFuture drainTask = drainTasks.get(channelRef); + if (runningTask != null && drainTask != runningTask.future) { + return; + } + drainTasks.remove(channelRef); + if (drainTask != null) { + drainTask.cancel(false); + } + if (channelRef.isActive() + || channelRef.getActiveStreamsCount() != 0 + || !removedChannelRefs.contains(channelRef) + || shuttingDown) { + return; + } + long elapsed = Math.max(0, nanoClock.get() - channelRef.getLastActivityNanos()); + if (elapsed < drainIdleGrace.toNanos()) { + reschedule = true; + } else { + closeChannel = removeDrainedChannel(channelRef); + } + } + if (reschedule) { + scheduleDrain(channelRef); + } else if (closeChannel) { + channelRef.getChannel().shutdown(); + } + } - // Removing a channel may change channel pool state. - executeStateChangeCallbacks(); + private final class DrainTask implements Runnable { + private final ChannelRef channelRef; + + @GuardedBy("GcpManagedChannel.this") + private boolean ranBeforeFutureAssignment; + + @GuardedBy("GcpManagedChannel.this") + @Nullable + private ScheduledFuture future; + + private DrainTask(ChannelRef channelRef) { + this.channelRef = channelRef; + } + + @Override + public void run() { + synchronized (GcpManagedChannel.this) { + if (future == null) { + ranBeforeFutureAssignment = true; + return; + } + } + finishDrain(channelRef, this); + } + } + + @GuardedBy("this") + private boolean removeDrainedChannel(ChannelRef channelRef) { + if (!removedChannelRefs.remove(channelRef)) { + return false; + } + fallbackMap.remove(channelRef.getId()); + channelIdToChannelRef.remove(channelRef.getId(), channelRef); + return true; + } + + private static int ceilDiv(long numerator, int denominator) { + if (numerator <= 0) { + return 0; + } + return (int) Math.min(Integer.MAX_VALUE, 1 + ((numerator - 1) / denominator)); + } + + private int targetRpcPerChannel() { + return Math.max(1, (minRpcPerChannel + maxRpcPerChannel) / 2); + } + + private long activeLoad(List refs) { + long load = 0; + for (ChannelRef channelRef : refs) { + if (channelRef.isActive()) { + load += channelRef.getActiveStreamsCount(); + } + } + return load; + } + + private long pickerLoad(List refs) { + long load = 0; + for (ChannelRef channelRef : refs) { + if (channelRef.isActive()) { + load += channelRef.getPickerLoad(); + } + } + return load; } private Supplier log(Supplier messageSupplier) { @@ -502,6 +688,7 @@ private void initOptions() { scaleDownConsecutiveLowLoadChecks = poolOptions.getScaleDownConsecutiveLowLoadChecks(); maxScaleUpPercent = poolOptions.getMaxScaleUpPercent(); maxScaleDownChannels = poolOptions.getMaxScaleDownChannels(); + drainIdleGrace = poolOptions.getDrainIdleGrace(); isDynamicScalingEnabled = minRpcPerChannel > 0 && maxRpcPerChannel > 0 && !scaleDownInterval.isZero(); channelPickStrategy = poolOptions.getChannelPickStrategy(); @@ -1528,9 +1715,8 @@ public void notifyWhenStateChanged(ConnectivityState source, Runnable callback) private class ChannelStateMonitor implements Runnable { private final ChannelRef channelRef; private final ManagedChannel channel; - private ConnectivityState currentState; + private volatile ConnectivityState currentState; private long connectingStartNanos; - private long connectedSinceNanos; @GuardedBy("channelRef") private boolean readyAccounted; @@ -1541,26 +1727,24 @@ private ChannelStateMonitor(ManagedChannel channel, ChannelRef channelRef) { run(); } - public long getConnectedSinceNanos() { - return connectedSinceNanos; - } - public ConnectivityState getCurrentState() { return currentState; } - private void accountReadyIfNeeded() { + private boolean accountReadyIfNeeded() { if (currentState == ConnectivityState.READY && !readyAccounted) { readyAccounted = true; - incReadyChannels(false); + return true; } + return false; } - private void unaccountReadyIfNeeded() { + private boolean unaccountReadyIfNeeded() { if (readyAccounted) { readyAccounted = false; - decReadyChannels(false); + return true; } + return false; } @Override @@ -1580,6 +1764,9 @@ public void run() { ConnectivityState newState = channel.getState(requestConnection); boolean isActive; + boolean incrementReady = false; + boolean decrementReady = false; + long readinessNanos = 0; synchronized (channelRef) { isActive = channelRef.isActive() && channelRefs.contains(channelRef); if (logger.isLoggable(Level.FINER)) { @@ -1589,30 +1776,35 @@ public void run() { channelRef.getId(), currentState, newState)); } if (newState == ConnectivityState.READY && currentState != ConnectivityState.READY) { - connectedSinceNanos = System.nanoTime(); if (isActive && !readyAccounted) { readyAccounted = true; - incReadyChannels(true); + incrementReady = true; if (connectingStartNanos > 0) { - saveReadinessTime(System.nanoTime() - connectingStartNanos); + readinessNanos = nanoClock.get() - connectingStartNanos; } } connectingStartNanos = 0; } if (newState != ConnectivityState.READY && readyAccounted) { readyAccounted = false; - decReadyChannels(true); + decrementReady = true; } if (newState == ConnectivityState.CONNECTING && currentState != ConnectivityState.CONNECTING) { - connectingStartNanos = System.nanoTime(); - } - if (newState != ConnectivityState.READY) { - connectedSinceNanos = 0; + connectingStartNanos = nanoClock.get(); } currentState = newState; } + if (incrementReady) { + incReadyChannels(true); + if (readinessNanos > 0) { + saveReadinessTime(readinessNanos); + } + } else if (decrementReady) { + decReadyChannels(true); + } + processChannelStateChange(channelRef.getId(), newState); if (isActive) { executeStateChangeCallbacks(); @@ -1641,6 +1833,11 @@ void processChannelStateChange(int channelId, ConnectivityState state) { if (!fallbackEnabled) { return; } + ChannelRef channelRef = channelIdToChannelRef.get(channelId); + if (channelRef == null || !channelRef.isActive()) { + fallbackMap.remove(channelId); + return; + } if (state == ConnectivityState.READY || state == ConnectivityState.IDLE) { // Ready fallbackMap.remove(channelId); @@ -1747,11 +1944,26 @@ protected ChannelRef getChannelRef(@Nullable String key) { if (key == null || key.isEmpty()) { return pickLeastBusyChannel(/* forFallback= */ false); } - ChannelRef mappedChannel = affinityKeyToChannelRef.get(key); - affinityKeyLastUsed.put(key, System.nanoTime()); + ChannelRef mappedChannel; + while (true) { + mappedChannel = affinityKeyToChannelRef.get(key); + if (mappedChannel == null) { + break; + } + long lastUsed = nanoClock.get(); + affinityKeyLastUsed.merge(key, lastUsed, Long::max); + if (affinityKeyToChannelRef.get(key) == mappedChannel) { + break; + } + affinityKeyLastUsed.remove(key, lastUsed); + } + while (mappedChannel != null && !mappedChannel.isActive()) { + mappedChannel = + unbindInactiveMapping(key, mappedChannel) ? null : affinityKeyToChannelRef.get(key); + } if (mappedChannel == null) { ChannelRef channelRef = pickLeastBusyChannel(/* forFallback= */ false); - bind(channelRef, Collections.singletonList(key)); + channelRef = bind(channelRef, Collections.singletonList(key)); return channelRef; } if (!fallbackEnabled) { @@ -1766,12 +1978,11 @@ protected ChannelRef getChannelRef(@Nullable String key) { // Channel is not ready. Look up if the affinity key mapped to another channel. Integer channelId = tempMap.get(key); if (channelId != null && !fallbackMap.containsKey(channelId)) { - // Fallback channel is ready. - if (logger.isLoggable(Level.FINEST)) { - logger.finest(log("Using fallback channel: %d -> %d", mappedChannel.getId(), channelId)); - } ChannelRef fallbackChannel = channelIdToChannelRef.get(channelId); if (fallbackChannel != null && fallbackChannel.isActive()) { + if (logger.isLoggable(Level.FINEST)) { + logger.finest(log("Using fallback channel: %d -> %d", mappedChannel.getId(), channelId)); + } fallbacksSucceeded.incrementAndGet(); return fallbackChannel; } @@ -1780,7 +1991,7 @@ protected ChannelRef getChannelRef(@Nullable String key) { // No temp mapping for this key or fallback channel is also broken. ChannelRef channelRef = pickLeastBusyChannel(/* forFallback= */ true); if (!fallbackMap.containsKey(channelRef.getId()) - && channelRef.getActiveStreamsCount() < DEFAULT_MAX_STREAM) { + && channelRef.getActiveStreamsCount() < maxConcurrentStreamsLowWatermark) { // Got a ready and not an overloaded channel. if (channelRef.getId() != mappedChannel.getId()) { if (logger.isLoggable(Level.FINEST)) { @@ -1807,7 +2018,9 @@ protected ChannelRef getChannelRef(@Nullable String key) { } /** - * Pick a {@link ChannelRef} using a caller-owned reference instead of grpc-gcp's affinity map. + * Picks a {@link ChannelRef} using a caller-owned reference instead of grpc-gcp's affinity map. A + * reference remains sticky while its delegate is open, including while the channel drains, and + * re-resolves after delegate shutdown or an explicit request to use a different channel. */ protected ChannelRef getChannelRefByAffinityRef(ChannelAffinityRef affinityRef) { // Retry if another thread updates the caller-owned affinity ref while we are picking a channel. @@ -1837,7 +2050,7 @@ protected ChannelRef getChannelRefByAffinityRef(ChannelAffinityRef affinityRef) private ChannelRef pickLeastBusyChannelDifferentFrom(@Nullable ChannelRef excludedChannelRef) { ChannelRef channelRef = pickLeastBusyChannel(/* forFallback= */ false); - if (excludedChannelRef == null || channelRefs.size() <= 1) { + if (excludedChannelRef == null) { return channelRef; } if (channelRef != excludedChannelRef && channelRef.isActive()) { @@ -1849,7 +2062,7 @@ private ChannelRef pickLeastBusyChannelDifferentFrom(@Nullable ChannelRef exclud if (candidate == excludedChannelRef || !candidate.isActive()) { continue; } - int streams = candidate.getActiveStreamsCount(); + int streams = candidate.getPickerLoad(); if (leastBusyChannelRef == null || streams < leastBusyStreams) { leastBusyChannelRef = candidate; leastBusyStreams = streams; @@ -1859,22 +2072,8 @@ private ChannelRef pickLeastBusyChannelDifferentFrom(@Nullable ChannelRef exclud } // Create a new channel and add it to channelRefs. - // If we have a ready channel not in the pool that we wait for completing its RPCs, - // then re-use that channel instead. @VisibleForTesting - ChannelRef createNewChannel() { - Optional reusedChannelRef = pickChannelForReuse(); - if (reusedChannelRef.isPresent()) { - ChannelRef chRef = reusedChannelRef.get(); - channelRefs.add(chRef); - removedChannelRefs.remove(chRef); - channelIdToChannelRef.put(chRef.getId(), chRef); - chRef.activateAndAccountReadiness(); - logger.finer(log("Channel %d reused.", chRef.getId())); - maxChannels.accumulateAndGet(getNumberOfChannels(), Math::max); - return chRef; - } - + synchronized ChannelRef createNewChannel() { ChannelRef channelRef = new ChannelRef(delegateChannelBuilder.build()); channelRefs.add(channelRef); channelIdToChannelRef.put(channelRef.getId(), channelRef); @@ -1884,13 +2083,6 @@ ChannelRef createNewChannel() { return channelRef; } - private Optional pickChannelForReuse() { - // Pick the most recently connected ready channel, if any. - return removedChannelRefs.stream() - .filter(channelRef -> channelRef.getState() == ConnectivityState.READY) - .max(Comparator.comparing(ChannelRef::getConnectedSinceNanos)); - } - @GuardedBy("this") private ChannelRef addBuiltChannel(ManagedChannel channel) { ChannelRef channelRef = new ChannelRef(channel); @@ -1908,7 +2100,7 @@ private ChannelRef createFirstChannel() { return null; } synchronized (this) { - if (channelRefs.isEmpty()) { + if (channelRefs.isEmpty() && !shuttingDown) { return createNewChannel(); } } @@ -1980,7 +2172,6 @@ private void runScaleUpWorker() { private void dynamicUpscale() { final int channelsToBuild; - int reused = 0; synchronized (this) { if (!isDynamicScalingEnabled || shuttingDown || channelRefs.size() >= maxSize) { return; @@ -1990,43 +2181,26 @@ private void dynamicUpscale() { && now - lastScaleUpNanos < scaleUpCooldown.toNanos()) { return; } - int active = channelRefs.size(); + List activeChannels = + channelRefs.stream().filter(ChannelRef::isActive).collect(Collectors.toList()); + int active = activeChannels.size(); if (active == 0) { return; } - 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 desired = ceilDiv(pickerLoad(activeChannels), targetRpcPerChannel()); int add = desired - active; // Small pools may add two channels per event before percentage growth dominates. - int percentCap = Math.max(2, (int) (1 + (((long) active * maxScaleUpPercent - 1) / 100))); + int percentCap = Math.max(2, ceilDiv((long) active * maxScaleUpPercent, 100)); add = Math.min(add, percentCap); add = Math.min(add, maxSize - active); if (add <= 0) { return; } - while (reused < add) { - Optional reusable = pickChannelForReuse(); - if (!reusable.isPresent()) { - break; - } - ChannelRef channelRef = reusable.get(); - removedChannelRefs.remove(channelRef); - channelRefs.add(channelRef); - channelIdToChannelRef.put(channelRef.getId(), channelRef); - channelRef.activateAndAccountReadiness(); - maxChannels.accumulateAndGet(getNumberOfChannels(), Math::max); - reused++; - } - channelsToBuild = add - reused; + channelsToBuild = add; // Claim cooldown before delegate construction begins. lastScaleUpNanos = now; } - scaleUpCount.addAndGet(reused); List builtChannels = new ArrayList<>(channelsToBuild); try { for (int i = 0; i < channelsToBuild; i++) { @@ -2078,16 +2252,19 @@ private boolean shouldScaleUp(int minStreams) { * be provided if available. */ private ChannelRef pickLeastBusyChannel(boolean forFallback) { - ChannelRef first = createFirstChannel(); - if (first != null) { - return first; - } - - if (!fallbackEnabled) { - return pickLeastBusyNoFallback(); + // Retries cover post-snapshot deactivation, not draining density. + for (int attempt = 0; attempt < 3; attempt++) { + ChannelRef first = createFirstChannel(); + if (first != null) { + return first; + } + ChannelRef picked = + fallbackEnabled ? pickLeastBusyWithFallback(forFallback) : pickLeastBusyNoFallback(); + if (validatePickedChannel(picked)) { + return picked; + } } - - return pickLeastBusyWithFallback(forFallback); + return leastLoadedActiveChannel(channelRefs); } /** @@ -2105,7 +2282,7 @@ private ChannelRef pickLeastBusyNoFallback() { int streams = channelPickStrategy == GcpManagedChannelOptions.ChannelPickStrategy.POWER_OF_TWO ? getMaxActiveStreams() - : channelCandidate.getActiveStreamsCount(); + : channelCandidate.getPickerLoad(); if (streams >= maxConcurrentStreamsLowWatermark) { ChannelRef newChannel = tryCreateNewChannel(); if (newChannel != null) { @@ -2132,12 +2309,12 @@ private ChannelRef pickLeastBusyWithFallback(boolean forFallback) { if (!channelRef.isActive()) { continue; } - int cnt = channelRef.getActiveStreamsCount(); + int cnt = channelRef.getPickerLoad(); if (overallCandidate == null || cnt < overallMinStreams) { overallMinStreams = cnt; overallCandidate = channelRef; } - if (!fallbackMap.containsKey(channelRef.getId()) && cnt < DEFAULT_MAX_STREAM) { + if (!fallbackMap.containsKey(channelRef.getId()) && cnt < maxConcurrentStreamsLowWatermark) { readyCandidates.add(channelRef); if (cnt > readyMaxStreams) { readyMaxStreams = cnt; @@ -2146,7 +2323,7 @@ private ChannelRef pickLeastBusyWithFallback(boolean forFallback) { } if (overallCandidate == null) { - return pickFromCandidates(channelRefs); + return leastLoadedActiveChannel(channelRefs); } // For scale-up, use maxStreams among ready channels (consistent with non-fallback path). @@ -2192,8 +2369,9 @@ private ChannelRef pickLeastBusyWithFallback(boolean forFallback) { /** * Picks a channel from the given candidate list using the configured strategy. * - *

For {@code POWER_OF_TWO}: samples twice with replacement and picks the less busy candidate. - * The first sample wins ties. Inactive candidates are retried before falling back to a full scan. + *

For {@code POWER_OF_TWO}: samples twice with replacement and picks the less busy sample. The + * first sample wins ties. Draining or inactive samples are retried up to twice the candidate + * count before a full active-channel scan. * *

For {@code LINEAR_SCAN}: deterministic scan picking the first least-busy active channel. */ @@ -2208,26 +2386,32 @@ ChannelRef pickFromCandidates(List candidates) { if (!first.isActive() || !second.isActive()) { continue; } - return first.getActiveStreamsCount() <= second.getActiveStreamsCount() ? first : second; + ChannelRef picked = pickLessBusy(first, second); + if (picked.isActive()) { + return picked; + } } } + return leastLoadedActiveChannel(candidates); + } + + private ChannelRef leastLoadedActiveChannel(List candidates) { ChannelRef best = null; - int bestStreams = Integer.MAX_VALUE; - for (Object element : snapshot) { - ChannelRef candidate = (ChannelRef) element; - if (!candidate.isActive()) { - continue; - } - int cnt = candidate.getActiveStreamsCount(); - if (best == null || cnt < bestStreams) { - bestStreams = cnt; + for (ChannelRef candidate : candidates) { + if (candidate.isActive() + && (best == null || candidate.getPickerLoad() < best.getPickerLoad())) { best = candidate; } } - if (best == null) { - throw new IllegalStateException("No active channel available"); + if (best != null) { + return best; } - return best; + throw Status.UNAVAILABLE.withDescription("No available channels").asRuntimeException(); + } + + @VisibleForTesting + ChannelRef pickLessBusy(ChannelRef first, ChannelRef second) { + return first.getPickerLoad() <= second.getPickerLoad() ? first : second; } @Override @@ -2247,6 +2431,9 @@ public String authority() { *

If method-affinity is specified, we will use the GcpClientCall to fetch the affinitykey and * bind/unbind the channel, otherwise we just need the SimpleGcpClientCall to keep track of the * number of streams in each channel. + * + *

A returned simple call reserves one unit of pool load immediately. If never started, callers + * must invoke {@link ClientCall#cancel(String, Throwable)} to release that reservation. */ @Override public ClientCall newCall( @@ -2308,6 +2495,8 @@ private synchronized void cancelBackgroundTasks() { logMetricsTask.cancel(false); logMetricsTask = null; } + drainTasks.values().forEach(task -> task.cancel(false)); + drainTasks.clear(); } @Override @@ -2455,9 +2644,13 @@ public ConnectivityState getState(boolean requestConnection) { *

One channel can be mapped to more than one keys. But one key can only be mapped to one * channel. */ - protected void bind(ChannelRef channelRef, List affinityKeys) { + protected synchronized ChannelRef bind(ChannelRef channelRef, List affinityKeys) { if (channelRef == null || affinityKeys == null) { - return; + return channelRef; + } + if (!channelRef.isActive()) { + channelRef = pickLeastBusyChannel(/* forFallback= */ false); + // Deactivation also holds the pool monitor, so this pick stays active through binding. } if (logger.isLoggable(Level.FINEST)) { logger.finest( @@ -2469,13 +2662,28 @@ protected void bind(ChannelRef channelRef, List affinityKeys) { while (affinityKeyToChannelRef.putIfAbsent(affinityKey, channelRef) != null) { unbind(Collections.singletonList(affinityKey)); } - affinityKeyLastUsed.put(affinityKey, System.nanoTime()); + affinityKeyLastUsed.merge(affinityKey, nanoClock.get(), Long::max); channelRef.affinityCountIncr(); } + return channelRef; + } + + private synchronized boolean unbindInactiveMapping(String affinityKey, ChannelRef mappedChannel) { + Runnable hook = inactiveMappingRemovalHookForTest; + if (hook != null) { + inactiveMappingRemovalHookForTest = null; + hook.run(); + } + if (affinityKeyToChannelRef.remove(affinityKey, mappedChannel)) { + affinityKeyLastUsed.remove(affinityKey); + mappedChannel.affinityCountDecr(); + return true; + } + return false; } /** Unbind channel with affinity key. */ - protected void unbind(List affinityKeys) { + protected synchronized void unbind(List affinityKeys) { if (affinityKeys == null) { return; } @@ -2611,7 +2819,9 @@ protected class ChannelRef { // activeStreamsCount are mutated from the GcpClientCall concurrently using the // `activeStreamsCountIncr()` and `activeStreamsCountDecr()` methods. private final AtomicInteger activeStreamsCount; - private long lastResponseNanos = nanoClock.get(); + private final long createdNanos = nanoClock.get(); + private volatile long lastActivityNanos = createdNanos; + private long lastResponseNanos = createdNanos; private final AtomicInteger deadlineExceededCount = new AtomicInteger(); private final AtomicLong okCalls = new AtomicLong(); private final AtomicLong errCalls = new AtomicLong(); @@ -2630,8 +2840,12 @@ protected ChannelRef(ManagedChannel channel, int affinityCount, int activeStream channelStateMonitor = new ChannelStateMonitor(channel, this); } - protected long getConnectedSinceNanos() { - return channelStateMonitor.getConnectedSinceNanos(); + protected long getCreatedNanos() { + return createdNanos; + } + + protected long getLastActivityNanos() { + return lastActivityNanos; } protected ConnectivityState getState() { @@ -2651,26 +2865,36 @@ protected boolean isActive() { } private void activateAndAccountReadiness() { + boolean incrementReady; synchronized (this) { active = true; - channelStateMonitor.accountReadyIfNeeded(); + incrementReady = channelStateMonitor.accountReadyIfNeeded(); + } + if (incrementReady) { + incReadyChannels(false); } } private void deactivateAndAccountReadiness() { + boolean decrementReady; synchronized (this) { - channelStateMonitor.unaccountReadyIfNeeded(); + decrementReady = channelStateMonitor.unaccountReadyIfNeeded(); active = false; } + if (decrementReady) { + decReadyChannels(false); + } } - private void deactivate() { + @VisibleForTesting + void deactivateForTest() { deactivateAndAccountReadiness(); } @VisibleForTesting - void deactivateForTest() { - deactivateAndAccountReadiness(); + void setActiveStreamsForTest(int streams) { + int previous = activeStreamsCount.getAndSet(streams); + totalActiveStreams.addAndGet(streams - previous); } protected void affinityCountIncr() { @@ -2690,6 +2914,7 @@ protected void resetAffinityCount() { } protected void activeStreamsCountIncr() { + lastActivityNanos = nanoClock.get(); int actStreams = activeStreamsCount.incrementAndGet(); maxActiveStreams.accumulateAndGet(actStreams, Math::max); int totalActStreams = totalActiveStreams.incrementAndGet(); @@ -2698,6 +2923,7 @@ protected void activeStreamsCountIncr() { } protected void activeStreamsCountDecr(long startNanos, Status status, boolean fromClientSide) { + lastActivityNanos = nanoClock.get(); int actStreams = activeStreamsCount.decrementAndGet(); minActiveStreams.accumulateAndGet(actStreams, Math::min); int totalActStreams = totalActiveStreams.decrementAndGet(); @@ -2712,6 +2938,9 @@ protected void activeStreamsCountDecr(long startNanos, Status status, boolean fr if (unresponsiveDetectionEnabled) { detectUnresponsiveConnection(startNanos, status, fromClientSide); } + if (actStreams == 0 && !isActive()) { + scheduleDrain(this); + } } protected void messageReceived() { @@ -2727,10 +2956,8 @@ protected int getActiveStreamsCount() { return activeStreamsCount.get(); } - @VisibleForTesting - void setActiveStreamsForTest(int streams) { - int previous = activeStreamsCount.getAndSet(streams); - totalActiveStreams.addAndGet(streams - previous); + protected int getPickerLoad() { + return getActiveStreamsCount(); } protected long getAndResetOkCalls() { diff --git a/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannelOptions.java b/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannelOptions.java index 615ed094532f..60d18b16982c 100644 --- a/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannelOptions.java +++ b/grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpManagedChannelOptions.java @@ -219,6 +219,7 @@ public static class GcpChannelPoolOptions { private final int maxScaleUpPercent; // Maximum channels removed in one scale-down check. private final int maxScaleDownChannels; + private final Duration drainIdleGrace; // Use round-robin channel selection for affinity binding calls. private final boolean useRoundRobinOnBind; @@ -240,6 +241,7 @@ public GcpChannelPoolOptions(Builder builder) { scaleDownConsecutiveLowLoadChecks = builder.scaleDownConsecutiveLowLoadChecks; maxScaleUpPercent = builder.maxScaleUpPercent; maxScaleDownChannels = builder.maxScaleDownChannels; + drainIdleGrace = builder.drainIdleGrace; concurrentStreamsLowWatermark = builder.concurrentStreamsLowWatermark; useRoundRobinOnBind = builder.useRoundRobinOnBind; affinityKeyLifetime = builder.affinityKeyLifetime; @@ -287,6 +289,10 @@ public int getMaxScaleDownChannels() { return maxScaleDownChannels; } + public Duration getDrainIdleGrace() { + return drainIdleGrace; + } + public int getConcurrentStreamsLowWatermark() { return concurrentStreamsLowWatermark; } @@ -323,7 +329,7 @@ public String toString() { "{maxSize: %d, minSize: %d, initSize: %d, minRpcPerChannel: %d, " + "maxRpcPerChannel: %d, scaleDownInterval: %s, scaleUpCooldown: %s, " + "scaleDownConsecutiveLowLoadChecks: %d, maxScaleUpPercent: %d, " - + "maxScaleDownChannels: %d, " + + "maxScaleDownChannels: %d, drainIdleGrace: %s, " + "concurrentStreamsLowWatermark: %d, useRoundRobinOnBind: %s, " + "affinityKeyLifetime: %s, cleanupInterval: %s, channelPickStrategy: %s}", getMaxSize(), @@ -336,6 +342,7 @@ public String toString() { getScaleDownConsecutiveLowLoadChecks(), getMaxScaleUpPercent(), getMaxScaleDownChannels(), + getDrainIdleGrace(), getConcurrentStreamsLowWatermark(), isUseRoundRobinOnBind(), getAffinityKeyLifetime(), @@ -354,6 +361,7 @@ public static class Builder { private int scaleDownConsecutiveLowLoadChecks = 3; private int maxScaleUpPercent = 30; private int maxScaleDownChannels = 2; + private Duration drainIdleGrace = Duration.ofMinutes(1); private int concurrentStreamsLowWatermark = GcpManagedChannel.DEFAULT_MAX_STREAM; private boolean useRoundRobinOnBind = false; private Duration affinityKeyLifetime = Duration.ZERO; @@ -377,6 +385,7 @@ public Builder(GcpChannelPoolOptions options) { this.scaleDownConsecutiveLowLoadChecks = options.getScaleDownConsecutiveLowLoadChecks(); this.maxScaleUpPercent = options.getMaxScaleUpPercent(); this.maxScaleDownChannels = options.getMaxScaleDownChannels(); + this.drainIdleGrace = options.getDrainIdleGrace(); this.concurrentStreamsLowWatermark = options.getConcurrentStreamsLowWatermark(); this.useRoundRobinOnBind = options.isUseRoundRobinOnBind(); this.affinityKeyLifetime = options.getAffinityKeyLifetime(); @@ -433,10 +442,10 @@ public Builder setInitSize(int initSize) { * channel or across the pool average signals a background scale-up worker. * *

Every scaleDownInterval, after consecutive low-load checks, the - * longest-connected channels (by connectedSinceNanos) are removed from selection, bounded by - * maxScaleDownChannels. A removed channel is shut down on a later check once its in-flight - * calls reach zero. A READY removed channel can be reused by a later scale-up before it - * closes. + * least-loaded channels are marked draining and removed from selection, bounded by + * maxScaleDownChannels; fewer affinity bindings and then older allocations break load ties. + * Draining channels receive no new picks or affinity binds and close after their in-flight + * calls reach zero and drainIdleGrace elapses. Scale-up always creates new channels. * * @param minRpcPerChannel minimum desired average concurrent calls per channel. * @param maxRpcPerChannel maximum desired average concurrent calls per channel. @@ -507,6 +516,18 @@ public Builder setMaxScaleDownChannels(int channels) { return this; } + /** + * Sets how long an idle draining channel remains open for sticky affinity references and + * in-flight calls before its delegate closes. Defaults to one minute. + */ + public Builder setDrainIdleGrace(Duration drainIdleGrace) { + Preconditions.checkNotNull(drainIdleGrace, "Drain idle grace must not be null."); + Preconditions.checkArgument( + !drainIdleGrace.isNegative(), "Drain idle grace must not be negative."); + this.drainIdleGrace = drainIdleGrace; + return this; + } + /** * Sets the concurrent streams low watermark. If every channel in the pool has at least this * amount of concurrent streams then a new channel will be created in the pool unless the pool diff --git a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelDrainingTest.java b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelDrainingTest.java new file mode 100644 index 000000000000..571da13dd377 --- /dev/null +++ b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelDrainingTest.java @@ -0,0 +1,347 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.grpc; + +import static com.google.common.truth.Truth.assertThat; +import static org.awaitility.Awaitility.await; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.grpc.GcpManagedChannel.ChannelRef; +import com.google.cloud.grpc.GcpManagedChannelOptions.GcpChannelPoolOptions; +import com.google.common.util.concurrent.MoreExecutors; +import io.grpc.ManagedChannel; +import io.grpc.Status; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +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.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; +import org.junit.After; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class GcpManagedChannelDrainingTest { + private final ExecutorService executor = MoreExecutors.newDirectExecutorService(); + private GcpManagedChannel pool; + + @After + public void tearDown() { + if (pool != null) { + pool.shutdownNow(); + } + executor.shutdownNow(); + } + + @Test + public void idleDrainedChannelWaitsForGraceBeforeClose() { + AtomicLong clock = new AtomicLong(System.nanoTime()); + pool = newPool(Duration.ofMinutes(1)); + clock.set(System.nanoTime()); + pool.setNanoClock(clock::get); + + scaleDown(); + ChannelRef draining = pool.removedChannelRefs.iterator().next(); + pool.finishDrain(draining); + assertThat(draining.getChannel().isShutdown()).isFalse(); + + clock.addAndGet(Duration.ofMinutes(1).plusNanos(1).toNanos()); + pool.finishDrain(draining); + assertThat(draining.getChannel().isShutdown()).isTrue(); + } + + @Test + public void activeDrainedChannelClosesAfterFinalStream() { + pool = newPool(Duration.ZERO); + ChannelRef victim = pool.channelRefs.get(0); + victim.activeStreamsCountIncr(); + for (int i = 1; i < pool.channelRefs.size(); i++) { + pool.channelRefs.get(i).setActiveStreamsForTest(2); + } + + scaleDown(); + assertThat(pool.removedChannelRefs).contains(victim); + assertThat(victim.getChannel().isShutdown()).isFalse(); + + victim.activeStreamsCountDecr(System.nanoTime(), Status.OK, false); + await().atMost(Duration.ofSeconds(5)).until(() -> victim.getChannel().isShutdown()); + } + + @Test + public void scaledUpPoolBuildsFreshChannelWhileRemovedChannelDrains() { + pool = newPool(Duration.ofMinutes(1)); + ChannelRef victim = pool.channelRefs.get(0); + ((GcpManagedChannelTest.FakeManagedChannel) victim.getChannel()) + .setState(io.grpc.ConnectivityState.READY); + + scaleDown(); + assertThat(pool.removedChannelRefs).contains(victim); + + ChannelRef fresh = pool.createNewChannel(); + pool.finishDrain(victim); + + assertThat(fresh).isNotSameInstanceAs(victim); + assertThat(pool.channelRefs).contains(fresh); + assertThat(pool.channelRefs).doesNotContain(victim); + assertThat(pool.removedChannelRefs).contains(victim); + assertThat(victim.isActive()).isFalse(); + assertThat(victim.getChannel().isShutdown()).isFalse(); + } + + @Test + public void scaleDownPrefersChannelWithFewerAffinitiesWhenLoadsTie() { + pool = newPool(Duration.ofMinutes(1)); + ChannelRef affiliated = pool.channelRefs.get(0); + pool.bind(affiliated, Collections.singletonList("session")); + + scaleDown(); + + assertThat(pool.channelRefs).contains(affiliated); + assertThat(pool.removedChannelRefs).doesNotContain(affiliated); + assertThat(affiliated.getAffinityCount()).isEqualTo(1); + } + + @Test + public void finishDrainShutsDelegateDownOutsidePoolMonitor() { + AtomicLong clock = new AtomicLong(System.nanoTime()); + AtomicReference poolReference = new AtomicReference<>(); + pool = + newPool( + Duration.ofMinutes(1), () -> new LockCheckingManagedChannel(executor, poolReference)); + poolReference.set(pool); + clock.set(System.nanoTime()); + pool.setNanoClock(clock::get); + + scaleDown(); + ChannelRef draining = pool.removedChannelRefs.iterator().next(); + LockCheckingManagedChannel delegate = (LockCheckingManagedChannel) draining.getChannel(); + clock.addAndGet(Duration.ofMinutes(1).plusNanos(1).toNanos()); + + pool.finishDrain(draining); + + assertThat(delegate.isShutdown()).isTrue(); + assertThat(delegate.shutdownWithPoolMonitorHeld.get()).isFalse(); + } + + @Test + public void rejectedDrainScheduleShutsDelegatesDownOutsidePoolMonitor() { + AtomicReference poolReference = new AtomicReference<>(); + List delegates = new CopyOnWriteArrayList<>(); + pool = + newPool( + Duration.ofMinutes(1), + () -> { + LockCheckingManagedChannel delegate = + new LockCheckingManagedChannel(executor, poolReference); + delegates.add(delegate); + return delegate; + }); + poolReference.set(pool); + ScheduledThreadPoolExecutor rejectingScheduler = new ScheduledThreadPoolExecutor(1); + rejectingScheduler.shutdown(); + pool.setDrainSchedulerForTest(rejectingScheduler); + + scaleDown(); + + int shutdownCount = 0; + for (LockCheckingManagedChannel delegate : delegates) { + if (delegate.isShutdown()) { + shutdownCount++; + assertThat(delegate.shutdownWithPoolMonitorHeld.get()).isFalse(); + } + } + assertThat(shutdownCount).isEqualTo(2); + assertThat(pool.removedChannelRefs).isEmpty(); + } + + @Test + public void staleDrainTaskCannotActOnBehalfOfNewerTask() { + ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class); + List tasks = new ArrayList<>(); + when(scheduler.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))) + .thenAnswer( + invocation -> { + tasks.add(invocation.getArgument(0)); + return mock(ScheduledFuture.class); + }); + pool = newPool(Duration.ZERO); + pool.setDrainSchedulerForTest(scheduler); + scaleDown(); + assertThat(tasks).hasSize(2); + ChannelRef draining = pool.removedChannelRefs.iterator().next(); + + pool.scheduleDrain(draining); + assertThat(tasks).hasSize(3); + assertThat(pool.drainTaskCountForTest()).isEqualTo(2); + + tasks.get(0).run(); + tasks.get(1).run(); + + assertThat(draining.getChannel().isShutdown()).isFalse(); + assertThat(pool.removedChannelRefs).contains(draining); + assertThat(pool.drainTaskCountForTest()).isEqualTo(1); + + tasks.get(2).run(); + + assertThat(draining.getChannel().isShutdown()).isTrue(); + assertThat(pool.removedChannelRefs).doesNotContain(draining); + assertThat(pool.drainTaskCountForTest()).isEqualTo(0); + } + + @Test + public void inlineDrainTaskFinishesAfterItsFutureIsPublished() { + ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class); + when(scheduler.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class))) + .thenAnswer( + invocation -> { + invocation.getArgument(0).run(); + return mock(ScheduledFuture.class); + }); + pool = newPool(Duration.ZERO); + pool.setDrainSchedulerForTest(scheduler); + + scaleDown(); + + assertThat(pool.removedChannelRefs).isEmpty(); + assertThat(pool.drainTaskCountForTest()).isEqualTo(0); + } + + @Test + public void affinityKeyRebindsAwayFromDrainingChannel() { + pool = newPool(Duration.ofMinutes(1)); + ChannelRef victim = pool.channelRefs.get(0); + pool.bind(victim, Collections.singletonList("session")); + for (int i = 1; i < pool.channelRefs.size(); i++) { + pool.channelRefs.get(i).setActiveStreamsForTest(1); + } + + scaleDown(); + + assertThat(pool.getChannelRef("session")).isNotSameInstanceAs(victim); + assertThat(pool.affinityKeyToChannelRef.get("session").isActive()).isTrue(); + } + + @Test + public void powerOfTwoUsesCandidateRetryBoundBeforeFullScan() { + pool = newPool(Duration.ofMinutes(1)); + List candidates = new CopyOnWriteArrayList<>(pool.channelRefs); + int initialSize = candidates.size(); + ChannelRef firstInactive = candidates.get(0); + ChannelRef secondInactive = candidates.get(1); + ChannelRef fallback = candidates.get(2); + firstInactive.deactivateForTest(); + secondInactive.deactivateForTest(); + AtomicBoolean shrunk = new AtomicBoolean(); + AtomicInteger nextSample = new AtomicInteger(); + pool.setCandidateIndexPickerForTest( + bound -> { + if (shrunk.compareAndSet(false, true)) { + candidates.remove(firstInactive); + candidates.remove(secondInactive); + } + nextSample.incrementAndGet(); + return 0; + }); + + assertThat(pool.pickFromCandidates(candidates)).isSameInstanceAs(fallback); + assertThat(nextSample.get()).isEqualTo(4 * initialSize); + } + + @Test + public void pickerRetriesWhenChannelDeactivatesBeforeValidation() { + pool = newPool(Duration.ofMinutes(1)); + AtomicReference deactivated = new AtomicReference<>(); + pool.setPickerValidationHookForTest( + candidate -> { + deactivated.set(candidate); + candidate.deactivateForTest(); + pool.channelRefs.remove(candidate); + pool.removedChannelRefs.add(candidate); + }); + + ChannelRef picked = pool.getChannelRef(null); + picked.activeStreamsCountIncr(); + + assertThat(picked).isNotSameInstanceAs(deactivated.get()); + assertThat(pool.channelRefs).doesNotContain(deactivated.get()); + assertThat(deactivated.get().getActiveStreamsCount()).isEqualTo(0); + assertThat(picked.getActiveStreamsCount()).isEqualTo(1); + } + + private void scaleDown() { + pool.checkScaleDown(); + pool.checkScaleDown(); + pool.checkScaleDown(); + } + + private GcpManagedChannel newPool(Duration drainIdleGrace) { + return newPool(drainIdleGrace, () -> new GcpManagedChannelTest.FakeManagedChannel(executor)); + } + + private GcpManagedChannel newPool( + Duration drainIdleGrace, Supplier channelFactory) { + GcpChannelPoolOptions options = + GcpChannelPoolOptions.newBuilder() + .setInitSize(4) + .setMinSize(2) + .setMaxSize(4) + .setDynamicScaling(10, 20, Duration.ofMinutes(1)) + .setScaleDownConsecutiveLowLoadChecks(3) + .setMaxScaleDownChannels(2) + .setDrainIdleGrace(drainIdleGrace) + .build(); + return (GcpManagedChannel) + GcpManagedChannelBuilder.forDelegateBuilder( + new GcpManagedChannelTest.FakeManagedChannelBuilder(channelFactory)) + .withOptions( + GcpManagedChannelOptions.newBuilder().withChannelPoolOptions(options).build()) + .build(); + } + + private static final class LockCheckingManagedChannel + extends GcpManagedChannelTest.FakeManagedChannel { + private final AtomicReference poolReference; + private final AtomicBoolean shutdownWithPoolMonitorHeld = new AtomicBoolean(); + + private LockCheckingManagedChannel( + ExecutorService executor, AtomicReference poolReference) { + super(executor); + this.poolReference = poolReference; + } + + @Override + public ManagedChannel shutdown() { + GcpManagedChannel currentPool = poolReference.get(); + shutdownWithPoolMonitorHeld.set(currentPool != null && Thread.holdsLock(currentPool)); + return super.shutdown(); + } + } +} diff --git a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelHotChannelReproducerTest.java b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelHotChannelReproducerTest.java index f25794edaea8..79eae75919fb 100644 --- a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelHotChannelReproducerTest.java +++ b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelHotChannelReproducerTest.java @@ -201,9 +201,6 @@ public void affinityReferencesRedistributeAfterDrainingChannelsShutdown() throws affinityRefs.add(affinityRef); originalIds.add(channelId); } - for (ChannelRef channelRef : pool.channelRefs) { - channelRef.activeStreamsCountIncr(); - } invokeScaleDownCheck(pool, 3); assertThat(pool.channelRefs).hasSize(2); @@ -277,6 +274,10 @@ public void scaleDownMarksAtMostTwoOfFortyEightLiveReferencesPerCheck() throws E handle.setChannelIdForTest(channelRef.getId()); handles.add(handle); } + for (ChannelRef channelRef : pool.channelRefs) { + channelRef.activeStreamsCountDecr(System.nanoTime(), Status.OK, false); + } + invokeScaleDownCheck(pool, 2); assertThat(pool.channelRefs).hasSize(48); invokeScaleDownCheck(pool, 1); @@ -381,8 +382,8 @@ private RunResult runScenario(long seed, Variant variant, LoadShape loadShape) t runConstantLoad(pool, callersExecutor, variant, sessionNames, seed, callers); } - // Final ramp restores a full active pool. Reset delegate counters so removed/reused channel - // history does not manufacture skew in the measured hold period. + // Final ramp restores a full active pool. Reset delegate counters so prior channel history + // does not manufacture skew in the measured hold period. runWave(pool, callersExecutor, variant, sessionNames, seed + 10_000, callers, 5, 10_000); delegateBuilder.resetMeasurements(); diff --git a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelOptionsTest.java b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelOptionsTest.java index 64e6479c5824..2a9aa1b2b6e8 100644 --- a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelOptionsTest.java +++ b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelOptionsTest.java @@ -254,6 +254,7 @@ public void channelPoolOptionsToStringIncludesEveryKnob() { assertThat(options).contains("scaleDownConsecutiveLowLoadChecks:"); assertThat(options).contains("maxScaleUpPercent:"); assertThat(options).contains("maxScaleDownChannels:"); + assertThat(options).contains("drainIdleGrace:"); assertThat(options).contains("concurrentStreamsLowWatermark:"); assertThat(options).contains("useRoundRobinOnBind:"); assertThat(options).contains("affinityKeyLifetime:"); diff --git a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelScaleUpWorkerTest.java b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelScaleUpWorkerTest.java index 394d184e3652..f801e59b070a 100644 --- a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelScaleUpWorkerTest.java +++ b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelScaleUpWorkerTest.java @@ -226,6 +226,18 @@ public void scaleUpClampsToMaxSize() { await().atMost(Duration.ofSeconds(5)).until(() -> pool.getNumberOfChannels() == 5); } + @Test + public void scaleUpCountsOnlyActiveChannels() { + pool = newPool(2, 5); + ChannelRef hot = pool.channelRefs.get(0); + pool.channelRefs.get(1).deactivateForTest(); + hot.setActiveStreamsForTest(6); + + hot.activeStreamsCountIncr(); + + await().atMost(Duration.ofSeconds(1)).until(() -> pool.getNumberOfChannels() == 4); + } + @Test public void shutdownDuringScaleUpClosesUnpublishedChannel() throws Exception { AtomicInteger builds = new AtomicInteger(); diff --git a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelTest.java b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelTest.java index fa652b48b321..e6ebb17fe087 100644 --- a/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelTest.java +++ b/grpc-gcp-java/src/test/java/com/google/cloud/grpc/GcpManagedChannelTest.java @@ -55,7 +55,6 @@ import io.opencensus.metrics.LabelValue; import java.io.File; import java.io.InputStream; -import java.lang.reflect.Field; import java.lang.reflect.Method; import java.net.URL; import java.time.Duration; @@ -64,9 +63,11 @@ import java.util.Collections; import java.util.LinkedList; import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -99,6 +100,15 @@ public final class GcpManagedChannelTest { private final List logRecords = new LinkedList<>(); + @Test + public void ceilDivReturnsZeroForNonPositiveNumerator() throws Exception { + Method ceilDiv = GcpManagedChannel.class.getDeclaredMethod("ceilDiv", long.class, int.class); + ceilDiv.setAccessible(true); + + assertThat(ceilDiv.invoke(null, -1L, 3)).isEqualTo(0); + assertThat(ceilDiv.invoke(null, 0L, 3)).isEqualTo(0); + } + private String lastLogMessage() { return lastLogMessage(1); } @@ -277,7 +287,9 @@ public void testGetChannelRefInitializationWithMinSize() throws InterruptedExcep GcpManagedChannelBuilder.forDelegateBuilder(builder).withOptions(options).build(); // Should have 2 channels since the beginning. assertThat(gcpChannel.channelRefs.size()).isEqualTo(2); - TimeUnit.MILLISECONDS.sleep(50); + await() + .atMost(Duration.ofSeconds(1)) + .until(() -> gcpChannel.getState(false) != ConnectivityState.IDLE); // The connection establishment must have been started on these two channels. assertThat(gcpChannel.getState(false)) .isAnyOf( @@ -347,7 +359,7 @@ public void testChannelAffinityRefInvalidChannelIdPicksAvailableChannel() throws gcpChannel = createPoolWithFakeReadyChannels(executorService, 2); ChannelAffinityRef affinityRef = new ChannelAffinityRef(); - setChannelAffinityRefState(affinityRef, 1000); + affinityRef.setChannelIdForTest(999); ChannelRef selected = gcpChannel.getChannelRefByAffinityRef(affinityRef); ChannelRef next = gcpChannel.getChannelRefByAffinityRef(affinityRef); @@ -403,23 +415,43 @@ public void fallbackUsesChannelIdMapAfterPoolHasIndexGap() { } @Test - public void readyAccountingRemainsExactWhenReadyChannelIsReused() { + public void readyAccountingRemainsExactWhenFreshChannelReplacesReadyDrainingChannel() { resetGcpChannel(); ExecutorService executorService = Executors.newSingleThreadExecutor(); try { - gcpChannel = createPoolWithFakeReadyChannels(executorService, 2); + gcpChannel = + new GcpManagedChannel( + new FakeManagedChannelBuilder( + () -> { + FakeManagedChannel channel = new FakeManagedChannel(executorService); + channel.setState(ConnectivityState.READY); + return channel; + }), + ApiConfig.getDefaultInstance(), + GcpManagedChannelOptions.newBuilder() + .withChannelPoolOptions( + GcpChannelPoolOptions.newBuilder() + .setInitSize(2) + .setMinSize(2) + .setMaxSize(3) + .build()) + .build()); assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(2); - ChannelRef reused = gcpChannel.channelRefs.get(0); - gcpChannel.channelRefs.remove(reused); - reused.deactivateForTest(); - gcpChannel.removedChannelRefs.add(reused); + ChannelRef draining = gcpChannel.channelRefs.get(0); + gcpChannel.channelRefs.remove(draining); + draining.deactivateForTest(); + gcpChannel.removedChannelRefs.add(draining); assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(1); - assertThat(gcpChannel.createNewChannel()).isSameInstanceAs(reused); + ChannelRef fresh = gcpChannel.createNewChannel(); + assertThat(fresh).isNotSameInstanceAs(draining); + assertThat(gcpChannel.channelRefs).contains(fresh); + assertThat(gcpChannel.channelRefs).doesNotContain(draining); + assertThat(gcpChannel.removedChannelRefs).contains(draining); assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(2); - reused.deactivateForTest(); - reused.deactivateForTest(); + fresh.deactivateForTest(); + fresh.deactivateForTest(); assertThat(gcpChannel.readyChannelCountForTest()).isEqualTo(1); } finally { gcpChannel.shutdownNow(); @@ -428,7 +460,7 @@ public void readyAccountingRemainsExactWhenReadyChannelIsReused() { } @Test - public void readyRemovedChannelIsReusedWhenNewerRemovedChannelIsNotReady() { + public void readyRemovedChannelIsNeverReused() { resetGcpChannel(); ExecutorService executorService = MoreExecutors.newDirectExecutorService(); try { @@ -437,36 +469,19 @@ public void readyRemovedChannelIsReusedWhenNewerRemovedChannelIsNotReady() { new FakeManagedChannelBuilder(() -> new FakeManagedChannel(executorService)), ApiConfig.getDefaultInstance(), GcpManagedChannelOptions.newBuilder().build()); - ChannelRef olderReady = - gcpChannel.new ChannelRef(new FakeManagedChannel(executorService)) { - @Override - protected long getConnectedSinceNanos() { - return 1; - } - - @Override - protected ConnectivityState getState() { - return ConnectivityState.READY; - } - }; - ChannelRef newerNotReady = - gcpChannel.new ChannelRef(new FakeManagedChannel(executorService)) { - @Override - protected long getConnectedSinceNanos() { - return 2; - } - - @Override - protected ConnectivityState getState() { - return ConnectivityState.IDLE; - } - }; - olderReady.deactivateForTest(); - newerNotReady.deactivateForTest(); - gcpChannel.removedChannelRefs.add(olderReady); - gcpChannel.removedChannelRefs.add(newerNotReady); - - assertThat(gcpChannel.createNewChannel()).isSameInstanceAs(olderReady); + ChannelRef firstDraining = gcpChannel.new ChannelRef(new FakeManagedChannel(executorService)); + ChannelRef secondDraining = + gcpChannel.new ChannelRef(new FakeManagedChannel(executorService)); + firstDraining.deactivateForTest(); + secondDraining.deactivateForTest(); + gcpChannel.removedChannelRefs.add(firstDraining); + gcpChannel.removedChannelRefs.add(secondDraining); + + ChannelRef fresh = gcpChannel.createNewChannel(); + assertThat(fresh).isNotSameInstanceAs(firstDraining); + assertThat(fresh).isNotSameInstanceAs(secondDraining); + assertThat(gcpChannel.channelRefs).contains(fresh); + assertThat(gcpChannel.removedChannelRefs).containsExactly(firstDraining, secondDraining); } finally { gcpChannel.shutdownNow(); executorService.shutdownNow(); @@ -474,7 +489,7 @@ protected ConnectivityState getState() { } @Test - public void testChannelAffinityRefRemovedChannelPicksAvailableChannel() throws Exception { + public void testChannelAffinityRefRemovedOpenChannelStaysStickyUntilShutdown() throws Exception { resetGcpChannel(); ExecutorService executorService = Executors.newSingleThreadExecutor(); try { @@ -483,21 +498,189 @@ public void testChannelAffinityRefRemovedChannelPicksAvailableChannel() throws E ChannelRef removed = gcpChannel.getChannelRefByAffinityRef(affinityRef); gcpChannel.channelRefs.remove(removed); - deactivateChannelRef(removed); + removed.deactivateForTest(); ChannelRef selected = gcpChannel.getChannelRefByAffinityRef(affinityRef); assertThat(selected).isSameInstanceAs(removed); + assertThat(selected.isActive()).isFalse(); + removed.getChannel().shutdownNow(); - selected = gcpChannel.getChannelRefByAffinityRef(affinityRef); + ChannelRef afterShutdown = gcpChannel.getChannelRefByAffinityRef(affinityRef); ChannelRef next = gcpChannel.getChannelRefByAffinityRef(affinityRef); + assertThat(afterShutdown).isNotSameInstanceAs(removed); + assertThat(afterShutdown).isIn(gcpChannel.channelRefs); + assertThat(afterShutdown.isActive()).isTrue(); + assertThat(next).isSameInstanceAs(afterShutdown); + } finally { + gcpChannel.shutdownNow(); + executorService.shutdownNow(); + } + } - assertThat(selected).isNotSameInstanceAs(removed); - assertThat(selected).isIn(gcpChannel.channelRefs); + @Test + public void inactiveAffinityRemovalUsesConcurrentRebind() { + resetGcpChannel(); + ExecutorService channelExecutor = MoreExecutors.newDirectExecutorService(); + try { + gcpChannel = createPoolWithFakeReadyChannels(channelExecutor, 3); + ChannelRef inactive = gcpChannel.channelRefs.get(0); + ChannelRef rebound = gcpChannel.channelRefs.get(1); + ChannelRef alternative = gcpChannel.channelRefs.get(2); + String key = "session"; + gcpChannel.bind(inactive, Collections.singletonList(key)); + inactive.deactivateForTest(); + rebound.setActiveStreamsForTest(10); + gcpChannel.setCandidateIndexPickerForTest(bound -> 2); + gcpChannel.setInactiveMappingRemovalHookForTest( + () -> gcpChannel.bind(rebound, Collections.singletonList(key))); + + ChannelRef selected = gcpChannel.getChannelRef(key); + + assertThat(selected).isSameInstanceAs(rebound); + assertThat(selected).isNotSameInstanceAs(alternative); + assertThat(gcpChannel.affinityKeyToChannelRef).containsEntry(key, rebound); + assertThat(gcpChannel.affinityKeyLastUsed).containsKey(key); + } finally { + gcpChannel.shutdownNow(); + channelExecutor.shutdownNow(); + } + } + + @Test + public void inactiveAffinityRemovalRetriesWhenConcurrentRebindBecomesInactive() { + resetGcpChannel(); + ExecutorService channelExecutor = MoreExecutors.newDirectExecutorService(); + try { + gcpChannel = createPoolWithFakeReadyChannels(channelExecutor, 3); + ChannelRef inactive = gcpChannel.channelRefs.get(0); + ChannelRef rebound = gcpChannel.channelRefs.get(1); + ChannelRef alternative = gcpChannel.channelRefs.get(2); + String key = "session"; + gcpChannel.bind(inactive, Collections.singletonList(key)); + inactive.deactivateForTest(); + gcpChannel.setCandidateIndexPickerForTest(bound -> 2); + gcpChannel.setInactiveMappingRemovalHookForTest( + () -> { + gcpChannel.bind(rebound, Collections.singletonList(key)); + rebound.deactivateForTest(); + }); + + ChannelRef selected = gcpChannel.getChannelRef(key); + + assertThat(selected).isSameInstanceAs(alternative); assertThat(selected.isActive()).isTrue(); - assertThat(next).isSameInstanceAs(selected); + assertThat(gcpChannel.affinityKeyToChannelRef).containsEntry(key, alternative); + assertThat(gcpChannel.affinityKeyLastUsed).containsKey(key); } finally { gcpChannel.shutdownNow(); - executorService.shutdownNow(); + channelExecutor.shutdownNow(); + } + } + + @Test + public void bindReturnsRedirectedActiveChannel() throws Exception { + resetGcpChannel(); + ExecutorService channelExecutor = MoreExecutors.newDirectExecutorService(); + try { + gcpChannel = createPoolWithFakeReadyChannels(channelExecutor, 2); + ChannelRef inactive = gcpChannel.channelRefs.get(0); + ChannelRef active = gcpChannel.channelRefs.get(1); + inactive.deactivateForTest(); + gcpChannel.setCandidateIndexPickerForTest(bound -> 1); + Method bind = GcpManagedChannel.class.getDeclaredMethod("bind", ChannelRef.class, List.class); + + Object bound = bind.invoke(gcpChannel, inactive, Collections.singletonList("session")); + + assertThat(bound).isSameInstanceAs(active); + assertThat(gcpChannel.affinityKeyToChannelRef).containsEntry("session", active); + } finally { + gcpChannel.shutdownNow(); + channelExecutor.shutdownNow(); + } + } + + @Test + public void affinityTouchRetriesWhenMappingIsUnboundConcurrently() throws Exception { + resetGcpChannel(); + ExecutorService channelExecutor = MoreExecutors.newDirectExecutorService(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + CountDownLatch mappingRead = new CountDownLatch(1); + CountDownLatch releaseTouch = new CountDownLatch(1); + AtomicBoolean firstClockRead = new AtomicBoolean(true); + try { + gcpChannel = createPoolWithFakeReadyChannels(channelExecutor, 2); + String key = "session"; + gcpChannel.bind(gcpChannel.channelRefs.get(0), Collections.singletonList(key)); + gcpChannel.setNanoClock( + () -> { + if (firstClockRead.compareAndSet(true, false)) { + mappingRead.countDown(); + try { + releaseTouch.await(5, TimeUnit.SECONDS); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + } + return 42L; + }); + + Future resolution = caller.submit(() -> gcpChannel.getChannelRef(key)); + assertThat(mappingRead.await(5, TimeUnit.SECONDS)).isTrue(); + gcpChannel.unbind(Collections.singletonList(key)); + releaseTouch.countDown(); + ChannelRef selected = resolution.get(5, TimeUnit.SECONDS); + + assertThat(gcpChannel.affinityKeyToChannelRef).containsEntry(key, selected); + assertThat(gcpChannel.affinityKeyLastUsed).containsKey(key); + } finally { + releaseTouch.countDown(); + caller.shutdownNow(); + gcpChannel.shutdownNow(); + channelExecutor.shutdownNow(); + } + } + + @Test + public void affinityTouchDoesNotOverwriteNewerConcurrentRebindTimestamp() throws Exception { + resetGcpChannel(); + ExecutorService channelExecutor = MoreExecutors.newDirectExecutorService(); + ExecutorService caller = Executors.newSingleThreadExecutor(); + CountDownLatch lookupClockRead = new CountDownLatch(1); + CountDownLatch releaseLookup = new CountDownLatch(1); + AtomicBoolean firstClockRead = new AtomicBoolean(true); + try { + gcpChannel = createPoolWithFakeReadyChannels(channelExecutor, 1); + String key = "session"; + ChannelRef channelRef = gcpChannel.channelRefs.get(0); + gcpChannel.bind(channelRef, Collections.singletonList(key)); + gcpChannel.setNanoClock( + () -> { + if (firstClockRead.compareAndSet(true, false)) { + lookupClockRead.countDown(); + try { + releaseLookup.await(5, TimeUnit.SECONDS); + } catch (InterruptedException failure) { + Thread.currentThread().interrupt(); + throw new AssertionError(failure); + } + return 10L; + } + return 20L; + }); + + Future resolution = caller.submit(() -> gcpChannel.getChannelRef(key)); + assertThat(lookupClockRead.await(5, TimeUnit.SECONDS)).isTrue(); + gcpChannel.bind(channelRef, Collections.singletonList(key)); + releaseLookup.countDown(); + + assertThat(resolution.get(5, TimeUnit.SECONDS)).isSameInstanceAs(channelRef); + assertThat(gcpChannel.affinityKeyLastUsed).containsEntry(key, 20L); + } finally { + releaseLookup.countDown(); + caller.shutdownNow(); + gcpChannel.shutdownNow(); + channelExecutor.shutdownNow(); } } @@ -522,19 +705,6 @@ private GcpManagedChannel createPoolWithFakeReadyChannels( .build(); } - private void setChannelAffinityRefState(ChannelAffinityRef affinityRef, int state) - throws Exception { - Field stateField = ChannelAffinityRef.class.getDeclaredField("state"); - stateField.setAccessible(true); - ((AtomicInteger) stateField.get(affinityRef)).set(state); - } - - private void deactivateChannelRef(ChannelRef channelRef) throws Exception { - Method deactivate = ChannelRef.class.getDeclaredMethod("deactivate"); - deactivate.setAccessible(true); - deactivate.invoke(channelRef); - } - @Test public void testGetChannelRefPickUpSmallest() { // This test verifies deterministic smallest-stream selection (LINEAR_SCAN behavior). @@ -639,7 +809,8 @@ public void testPickLeastBusyStillPrefersLessBusyChannels() { gcpChannel.channelRefs.add(gcpChannel.new ChannelRef(channel, i, 0)); } - // Pick 100 times. Channel 0 wins only when both samples select it. + // Pick 100 times. Channel 0 (50 streams) should almost never be picked because + // any random pair that includes an idle channel will prefer the idle one. int busyPicks = 0; for (int i = 0; i < 100; i++) { ChannelRef picked = gcpChannel.getChannelRef(null); @@ -647,6 +818,7 @@ public void testPickLeastBusyStillPrefersLessBusyChannels() { busyPicks++; } } + // Sampling is with replacement, so the busy channel wins only when sampled twice. assertThat(busyPicks).isLessThan(10); } @@ -658,8 +830,8 @@ public void testPickLeastBusyStillPrefersLessBusyChannels() { public void testPickLeastBusyWithDynamicScaleUp() throws InterruptedException { final int minSize = 2; final int maxSize = 6; - final int minRpcPerChannel = 2; - final int maxRpcPerChannel = 5; + final int minRpcPerChannel = 5; + final int maxRpcPerChannel = 7; final Duration scaleDownInterval = Duration.ofMillis(50); final ExecutorService executorService = Executors.newSingleThreadExecutor(); @@ -696,12 +868,12 @@ public void testPickLeastBusyWithDynamicScaleUp() throws InterruptedException { // One more call triggers scale-up. pool.getChannelRef(null).activeStreamsCountIncr(); - await().atMost(Duration.ofSeconds(5)).until(() -> pool.getNumberOfChannels() == minSize + 2); + await().atMost(Duration.ofSeconds(5)).until(() -> pool.getNumberOfChannels() > minSize); + assertThat(pool.getNumberOfChannels()).isEqualTo(minSize + 1); - // Mark the new channels as READY. - for (int i = minSize; i < pool.getNumberOfChannels(); i++) { - ((FakeManagedChannel) pool.channelRefs.get(i).getChannel()).setState(ConnectivityState.READY); - } + // Mark the new channel as READY. + ((FakeManagedChannel) pool.channelRefs.get(minSize).getChannel()) + .setState(ConnectivityState.READY); // Now pick many times without incrementing. The new (less busy) channel should be favored, // but picks should still be distributed across channels. @@ -749,10 +921,7 @@ public void testPickLeastBusySingleChannel() { } } - /** - * With only 2 channels, power-of-two degenerates to comparing both — should always pick the less - * busy one. - */ + /** With only 2 channels, sampling with replacement strongly prefers the less busy one. */ @Test public void testPickLeastBusyTwoChannels() { resetGcpChannel(); @@ -760,14 +929,15 @@ public void testPickLeastBusyTwoChannels() { ManagedChannel ch1 = builder.build(); gcpChannel.channelRefs.add(gcpChannel.new ChannelRef(ch0, 0, 10)); gcpChannel.channelRefs.add(gcpChannel.new ChannelRef(ch1, 1, 3)); - AtomicInteger sample = new AtomicInteger(); - gcpChannel.setCandidateIndexPickerForTest(ignored -> sample.getAndIncrement() % 2); - // With 2 channels, both are always selected, so the one with fewer streams always wins. + int lessBusyPicks = 0; for (int i = 0; i < 100; i++) { ChannelRef picked = gcpChannel.getChannelRef(null); - assertThat(picked).isEqualTo(gcpChannel.channelRefs.get(1)); + if (picked == gcpChannel.channelRefs.get(1)) { + lessBusyPicks++; + } } + assertThat(lessBusyPicks).isGreaterThan(60); } /** @@ -804,7 +974,7 @@ public void testLinearScanStrategyAlwaysPicksFirstOnTie() { } } - /** Verifies that POWER_OF_TWO does not prefer a warm channel when stream counts are tied. */ + /** Verifies that POWER_OF_TWO does not add a warmth bias when stream counts are tied. */ @Test public void testPowerOfTwoDoesNotPreferWarmChannelOnTie() throws Exception { resetGcpChannel(); @@ -824,7 +994,7 @@ public void testPowerOfTwoDoesNotPreferWarmChannelOnTie() throws Exception { ChannelRef warmChannel = gcpChannel.channelRefs.get(5); warmChannel.messageReceived(); - // Pick many times. Warmth does not affect tied samples. + // Repeated picks keep the first sample on ties, so warmth does not bias selection. int warmPicks = 0; final int numPicks = 1000; for (int i = 0; i < numPicks; i++) { @@ -834,8 +1004,8 @@ public void testPowerOfTwoDoesNotPreferWarmChannelOnTie() throws Exception { } } - assertThat(warmPicks).isGreaterThan(numPicks * 6 / 100); - assertThat(warmPicks).isLessThan(numPicks * 14 / 100); + assertThat(warmPicks).isAtLeast(numPicks * 5 / 100); + assertThat(warmPicks).isLessThan(numPicks * 15 / 100); } private void assertFallbacksMetric( @@ -882,9 +1052,6 @@ public void testGetChannelRefWithFallback() { GcpMetricsOptions.newBuilder().withMetricRegistry(fakeRegistry).build()) .build()) .build(); - AtomicInteger fallbackSample = new AtomicInteger(); - pool.setCandidateIndexPickerForTest( - bound -> Math.floorMod(fallbackSample.getAndIncrement(), bound)); final int currentIndex = GcpManagedChannel.channelPoolIndex.get(); final String poolIndex = String.format("pool-%d", currentIndex); @@ -940,25 +1107,23 @@ public void testGetChannelRefWithFallback() { assertEquals(2, chRef.getId()); assertEquals(3, pool.getNumberOfChannels()); - // Now we reached max pool size. Let's bring channel 2 to the low watermark and channel 1 to the - // low watermark + 1 streams. - for (int i = 0; i < lowWatermark; i++) { + // Now we reached max pool size. Bring channel 2 just below the configured watermark and + // channel 1 above it. + for (int i = 0; i < lowWatermark - 1; i++) { pool.channelRefs.get(2).activeStreamsCountIncr(); } pool.channelRefs.get(1).activeStreamsCountIncr(); - // As we reached max size and cannot create new channels and having ready channels with low - // watermark and low watermark + 1 streams, power-of-two selection can return either ready - // channel because it samples with replacement. + // Channel 2 remains eligible because its load is below the configured watermark. assertEquals(lowWatermark + 1, pool.channelRefs.get(1).getActiveStreamsCount()); - assertEquals(lowWatermark, pool.channelRefs.get(2).getActiveStreamsCount()); + assertEquals(lowWatermark - 1, pool.channelRefs.get(2).getActiveStreamsCount()); chRef = pool.getChannelRef(null); - assertThat(chRef.getId()).isAnyOf(1, 2); + assertEquals(2, chRef.getId()); assertEquals(3, pool.getNumberOfChannels()); - // This was the third fallback from non-ready channel 0 to a ready channel. - assertFallbacksMetric(fakeRegistry, 3, 0); + // Both creating and subsequently selecting channel 2 are successful fallbacks. + assertFallbacksMetric(fakeRegistry, 4, 0); // Let's bring channel 1 to max streams and mark channel 2 as not ready. - for (int i = 0; i < MAX_STREAM - lowWatermark; i++) { + for (int i = 0; i < MAX_STREAM - (lowWatermark - 1); i++) { pool.channelRefs.get(2).activeStreamsCountIncr(); } pool.processChannelStateChange(1, ConnectivityState.CONNECTING); @@ -977,7 +1142,7 @@ public void testGetChannelRefWithFallback() { assertThat(logRecords.size()).isEqualTo(++logCount); assertThat(lastLogMessage()).isEqualTo(poolIndex + ": Failed to find fallback for channel 0"); assertThat(lastLogLevel()).isEqualTo(Level.FINEST); - assertFallbacksMetric(fakeRegistry, 3, 1); + assertFallbacksMetric(fakeRegistry, 4, 1); // Let's have an affinity key and bind it to channel 0. final String key = "ABC"; @@ -992,9 +1157,12 @@ public void testGetChannelRefWithFallback() { assertThat(logRecords.size()).isEqualTo(++logCount); assertThat(lastLogMessage()).isEqualTo(poolIndex + ": Failed to find fallback for channel 0"); assertThat(lastLogLevel()).isEqualTo(Level.FINEST); - assertFallbacksMetric(fakeRegistry, 3, 2); + assertFallbacksMetric(fakeRegistry, 4, 2); - // Let's return channel 1 to a ready state. + // Return channel 1 below the configured watermark and to a ready state. + while (pool.channelRefs.get(1).getActiveStreamsCount() >= lowWatermark) { + pool.channelRefs.get(1).activeStreamsCountDecr(System.nanoTime(), Status.OK, false); + } pool.processChannelStateChange(1, ConnectivityState.READY); logCount = logRecords.size(); // Now we have a fallback candidate. @@ -1004,7 +1172,7 @@ public void testGetChannelRefWithFallback() { assertThat(logRecords.size()).isEqualTo(++logCount); assertThat(lastLogMessage()).isEqualTo(poolIndex + ": Setting fallback channel: 0 -> 1"); assertThat(lastLogLevel()).isEqualTo(Level.FINEST); - assertFallbacksMetric(fakeRegistry, 4, 2); + assertFallbacksMetric(fakeRegistry, 5, 2); // Let's briefly bring channel 2 to ready state. pool.processChannelStateChange(2, ConnectivityState.READY); @@ -1018,7 +1186,7 @@ public void testGetChannelRefWithFallback() { assertThat(logRecords.size()).isEqualTo(++logCount); assertThat(lastLogMessage()).isEqualTo(poolIndex + ": Using fallback channel: 0 -> 1"); assertThat(lastLogLevel()).isEqualTo(Level.FINEST); - assertFallbacksMetric(fakeRegistry, 5, 2); + assertFallbacksMetric(fakeRegistry, 6, 2); pool.processChannelStateChange(2, ConnectivityState.CONNECTING); // Let's bring channel 1 back to connecting state. @@ -1032,7 +1200,7 @@ public void testGetChannelRefWithFallback() { assertThat(logRecords.size()).isEqualTo(++logCount); assertThat(lastLogMessage()).isEqualTo(poolIndex + ": Failed to find fallback for channel 0"); assertThat(lastLogLevel()).isEqualTo(Level.FINEST); - assertFallbacksMetric(fakeRegistry, 5, 3); + assertFallbacksMetric(fakeRegistry, 6, 3); // Finally, we bring both channel 1 and channel 0 to the ready state and we should get the // original channel 0 for the key without any fallbacks happening. @@ -1042,7 +1210,7 @@ public void testGetChannelRefWithFallback() { chRef = pool.getChannelRef(key); assertEquals(0, chRef.getId()); assertThat(logRecords.size()).isEqualTo(logCount); - assertFallbacksMetric(fakeRegistry, 5, 3); + assertFallbacksMetric(fakeRegistry, 6, 3); } @Test @@ -1071,8 +1239,6 @@ public void testBindUnbindKey() { ChannelRef cf2 = gcpChannel.new ChannelRef(builder.build(), 0, 4); gcpChannel.channelRefs.add(cf1); gcpChannel.channelRefs.add(cf2); - AtomicInteger sample = new AtomicInteger(); - gcpChannel.setCandidateIndexPickerForTest(ignored -> sample.getAndIncrement() % 2); gcpChannel.bind(cf1, Collections.singletonList("key1")); @@ -1132,20 +1298,18 @@ public void testUsingKeyWithoutBinding() { ChannelRef cf2 = gcpChannel.new ChannelRef(builder.build(), 0, 4); gcpChannel.channelRefs.add(cf1); gcpChannel.channelRefs.add(cf2); - AtomicInteger sample = new AtomicInteger(); - gcpChannel.setCandidateIndexPickerForTest(ignored -> sample.getAndIncrement() % 2); final String key = "non-binded-key"; ChannelRef channelRef = gcpChannel.getChannelRef(key); - // Should bind on the fly to the least busy channel, which is 2. - assertThat(channelRef.getId()).isEqualTo(2); + // Power-of-two binds on the fly to its sampled winner. + int boundChannelId = channelRef.getId(); + assertThat(gcpChannel.affinityKeyToChannelRef.get(key)).isSameInstanceAs(channelRef); cf1.activeStreamsCountDecr(System.nanoTime(), Status.OK, true); cf1.activeStreamsCountDecr(System.nanoTime(), Status.OK, true); channelRef = gcpChannel.getChannelRef(key); - // Even after channel 1 now has less active streams (3) the channel 2 is still mapped for the - // same key. - assertThat(channelRef.getId()).isEqualTo(2); + // A load change does not move the existing binding. + assertThat(channelRef.getId()).isEqualTo(boundChannelId); } @Test @@ -1451,6 +1615,8 @@ public void testLogMetrics() throws Exception { .build()) .build()) .build(); + AtomicLong nanoClock = new AtomicLong(System.nanoTime()); + pool.setNanoClock(nanoClock::get); try { final int currentIndex = GcpManagedChannel.channelPoolIndex.get(); @@ -1462,7 +1628,7 @@ public void testLogMetrics() throws Exception { // Simulate channel connecting. channels.get(i).setState(ConnectivityState.CONNECTING); waitForStateCallbacks(executorService); - TimeUnit.MILLISECONDS.sleep(10); + nanoClock.addAndGet(Duration.ofMillis(10).toNanos()); // For the last one... if (i == streams.length - 1) { @@ -1474,12 +1640,12 @@ public void testLogMetrics() throws Exception { channels.get(j).setState(ConnectivityState.CONNECTING); } waitForStateCallbacks(executorService); - TimeUnit.MILLISECONDS.sleep(100); + nanoClock.addAndGet(Duration.ofMillis(110).toNanos()); // And this will be a failed fallback (no ready channels). pool.getChannelRef(null); // Simulate unresponsive connection. - long startNanos = System.nanoTime(); + long startNanos = nanoClock.get(); final Status deStatus = Status.fromCode(Code.DEADLINE_EXCEEDED); ref.activeStreamsCountIncr(); ref.activeStreamsCountDecr(startNanos, deStatus, false); @@ -1487,12 +1653,12 @@ public void testLogMetrics() throws Exception { ref.activeStreamsCountDecr(startNanos, deStatus, false); // Simulate unresponsive connection with more dropped calls. - startNanos = System.nanoTime(); + startNanos = nanoClock.get(); ref.activeStreamsCountIncr(); ref.activeStreamsCountDecr(startNanos, deStatus, false); ref.activeStreamsCountIncr(); ref.activeStreamsCountDecr(startNanos, deStatus, false); - TimeUnit.MILLISECONDS.sleep(110); + nanoClock.addAndGet(Duration.ofMillis(110).toNanos()); ref.activeStreamsCountIncr(); ref.activeStreamsCountDecr(startNanos, deStatus, false); } @@ -1691,6 +1857,8 @@ public void testUnresponsiveDetection() throws InterruptedException { GcpMetricsOptions.newBuilder().withMetricRegistry(fakeRegistry).build()) .build()) .build(); + AtomicLong nanoClock = new AtomicLong(System.nanoTime()); + pool.setNanoClock(nanoClock::get); int currentIndex = GcpManagedChannel.channelPoolIndex.get(); String poolIndex = String.format("pool-%d", currentIndex); final AtomicInteger idleCounter = new AtomicInteger(); @@ -1698,10 +1866,10 @@ public void testUnresponsiveDetection() throws InterruptedException { ChannelRef chRef = pool.new ChannelRef(channel); assertEquals(0, idleCounter.get()); - TimeUnit.MILLISECONDS.sleep(105); + nanoClock.addAndGet(Duration.ofMillis(105).toNanos()); // Report 3 deadline exceeded errors after 100 ms. - long startNanos = System.nanoTime(); + long startNanos = nanoClock.get(); final Status deStatus = Status.fromCode(Code.DEADLINE_EXCEEDED); chRef.activeStreamsCountDecr(startNanos, deStatus, false); assertEquals(0, idleCounter.get()); @@ -1777,34 +1945,36 @@ public void testUnresponsiveDetection() throws InterruptedException { + " = 1\\d\\d"); // Any message from the server must reset the dropped requests count and timestamp. - TimeUnit.MILLISECONDS.sleep(105); - startNanos = System.nanoTime(); + nanoClock.addAndGet(Duration.ofMillis(105).toNanos()); + startNanos = nanoClock.get(); chRef.activeStreamsCountDecr(startNanos, deStatus, false); assertEquals(1, idleCounter.get()); chRef.activeStreamsCountDecr(startNanos, deStatus, false); assertEquals(1, idleCounter.get()); // A message received from the server. + nanoClock.incrementAndGet(); chRef.messageReceived(); chRef.activeStreamsCountDecr(startNanos, deStatus, false); // No idle increment expected because dropped requests count and timestamp were reset. assertEquals(1, idleCounter.get()); // Any non-deadline exceeded response must reset the dropped requests count and timestamp. - TimeUnit.MILLISECONDS.sleep(105); - startNanos = System.nanoTime(); + nanoClock.addAndGet(Duration.ofMillis(105).toNanos()); + startNanos = nanoClock.get(); chRef.activeStreamsCountDecr(startNanos, deStatus, false); assertEquals(1, idleCounter.get()); chRef.activeStreamsCountDecr(startNanos, deStatus, false); assertEquals(1, idleCounter.get()); // Response with UNAVAILABLE status received from the server. final Status unavailableStatus = Status.fromCode(Code.UNAVAILABLE); + nanoClock.incrementAndGet(); chRef.activeStreamsCountDecr(startNanos, unavailableStatus, false); chRef.activeStreamsCountDecr(startNanos, deStatus, false); // No idle increment expected because dropped requests count and timestamp were reset. assertEquals(1, idleCounter.get()); // Even if dropped requests count is reached, it must also respect 100 ms configured. - startNanos = System.nanoTime(); + startNanos = nanoClock.get(); chRef.activeStreamsCountDecr(startNanos, deStatus, false); assertEquals(1, idleCounter.get()); chRef.activeStreamsCountDecr(startNanos, deStatus, false); @@ -1813,7 +1983,7 @@ public void testUnresponsiveDetection() throws InterruptedException { // Even it's third deadline exceeded no idle increment is expected because 100ms has not pass. assertEquals(1, idleCounter.get()); - TimeUnit.MILLISECONDS.sleep(105); + nanoClock.addAndGet(Duration.ofMillis(105).toNanos()); // Any subsequent deadline exceeded after 100ms must trigger the reconnection. chRef.activeStreamsCountDecr(startNanos, deStatus, false); assertEquals(2, idleCounter.get()); @@ -1861,9 +2031,7 @@ public void testStateNotifications() throws InterruptedException { gcpChannel.notifyWhenStateChanged( ConnectivityState.SHUTDOWN, () -> immediateCallbackCalled.set(true)); - TimeUnit.MILLISECONDS.sleep(2); - - assertThat(immediateCallbackCalled.get()).isTrue(); + await().atMost(Duration.ofSeconds(1)).untilTrue(immediateCallbackCalled); // Subscribe for notification when leaving IDLE state. final AtomicReference newState = new AtomicReference<>(); @@ -1887,10 +2055,12 @@ public void run() { // Make sure it was IDLE; assertThat(currentState).isEqualTo(ConnectivityState.IDLE); - TimeUnit.MILLISECONDS.sleep(25); - - assertThat(newState.get()) - .isAnyOf(ConnectivityState.CONNECTING, ConnectivityState.TRANSIENT_FAILURE); + await() + .atMost(Duration.ofSeconds(1)) + .untilAsserted( + () -> + assertThat(newState.get()) + .isAnyOf(ConnectivityState.CONNECTING, ConnectivityState.TRANSIENT_FAILURE)); } @Test @@ -2045,6 +2215,8 @@ public void testAffinityKeysCleanup() throws InterruptedException { .build()) .build()) .build(); + AtomicLong nanoClock = new AtomicLong(System.nanoTime()); + pool.setNanoClock(nanoClock::get); final String liveKey = "live-key"; ChannelRef ch0 = pool.getChannelRef(liveKey); @@ -2072,15 +2244,18 @@ public void testAffinityKeysCleanup() throws InterruptedException { assertThat(pool.getChannelRef(expKey).getId()).isEqualTo(2); // Halfway through affinity lifetime we use the live key again. - TimeUnit.MILLISECONDS.sleep(100); + nanoClock.addAndGet(Duration.ofMillis(100).toNanos()); ch0 = pool.getChannelRef(liveKey); // Make sure affinity still works. assertThat(ch0.getId()).isEqualTo(0); // Wait the remaining time and check that there is still affinity for the live key // but no affinity for the expired key. + nanoClock.addAndGet(Duration.ofMillis(150).toNanos()); - TimeUnit.MILLISECONDS.sleep(150); + await() + .atMost(Duration.ofSeconds(1)) + .until(() -> !pool.affinityKeyToChannelRef.containsKey(expKey)); assertThat(pool.affinityKeyToChannelRef.keySet().size()).isEqualTo(1); assertThat(pool.affinityKeyToChannelRef.get(liveKey)).isEqualTo(ch0); @@ -2096,7 +2271,7 @@ public void testAffinityKeysCleanup() throws InterruptedException { } @Test - public void testDynamicChannelPool() { + public void testDynamicChannelPool() throws InterruptedException { ExecutorService executorService = Executors.newSingleThreadExecutor(); GcpManagedChannel pool = null; try { @@ -2108,25 +2283,29 @@ public void testDynamicChannelPool() { GcpManagedChannelOptions.newBuilder() .withChannelPoolOptions( GcpChannelPoolOptions.newBuilder() - .setInitSize(4) + .setInitSize(2) .setMinSize(2) .setMaxSize(4) - .setDynamicScaling(2, 5, Duration.ofMinutes(1)) - .setScaleDownConsecutiveLowLoadChecks(3) - .setMaxScaleDownChannels(2) - .setChannelPickStrategy( - GcpManagedChannelOptions.ChannelPickStrategy.LINEAR_SCAN) + .setDynamicScaling(2, 5, Duration.ofMillis(20)) + .setScaleUpCooldown(Duration.ofNanos(1)) + .setScaleDownConsecutiveLowLoadChecks(1) + .setDrainIdleGrace(Duration.ZERO) .build()) .build()) .build(); - for (ChannelRef channelRef : pool.channelRefs) { - channelRef.activeStreamsCountIncr(); + GcpManagedChannel monitoredPool = pool; + + ChannelRef hot = pool.channelRefs.get(0); + for (int i = 0; i < 7; i++) { + hot.activeStreamsCountIncr(); } + await().atMost(Duration.ofSeconds(5)).until(() -> monitoredPool.getNumberOfChannels() == 3); + assertThat(pool.getNumberOfChannels()).isEqualTo(3); - pool.checkScaleDown(); - pool.checkScaleDown(); - assertThat(pool.getNumberOfChannels()).isEqualTo(4); - pool.checkScaleDown(); + while (hot.getActiveStreamsCount() > 0) { + hot.activeStreamsCountDecr(System.nanoTime(), Status.OK, false); + } + await().atMost(Duration.ofSeconds(5)).until(() -> monitoredPool.getNumberOfChannels() == 2); assertThat(pool.getNumberOfChannels()).isEqualTo(2); } finally { if (pool != null) { @@ -2137,7 +2316,7 @@ public void testDynamicChannelPool() { } @Test - public void testDynamicChannelPoolWithAffinity() { + public void testDynamicChannelPoolWithAffinity() throws InterruptedException { ExecutorService executorService = Executors.newSingleThreadExecutor(); GcpManagedChannel pool = null; try { @@ -2149,27 +2328,26 @@ public void testDynamicChannelPoolWithAffinity() { GcpManagedChannelOptions.newBuilder() .withChannelPoolOptions( GcpChannelPoolOptions.newBuilder() - .setInitSize(4) - .setMinSize(2) - .setMaxSize(4) - .setDynamicScaling(2, 5, Duration.ofMinutes(1)) - .setScaleDownConsecutiveLowLoadChecks(3) - .setMaxScaleDownChannels(2) + .setInitSize(2) + .setMinSize(1) + .setMaxSize(2) + .setDynamicScaling(1, 3, Duration.ofMillis(20)) + .setScaleDownConsecutiveLowLoadChecks(1) + .setDrainIdleGrace(Duration.ofSeconds(5)) .build()) .build()) .build(); + GcpManagedChannel monitoredPool = pool; + ChannelRef victim = pool.channelRefs.get(0); pool.bind(victim, Collections.singletonList("session")); - for (ChannelRef channelRef : pool.channelRefs) { - channelRef.activeStreamsCountIncr(); - } - - pool.checkScaleDown(); - pool.checkScaleDown(); - pool.checkScaleDown(); + pool.channelRefs.get(1).activeStreamsCountIncr(); + await().atMost(Duration.ofSeconds(5)).until(() -> monitoredPool.getNumberOfChannels() == 1); - assertThat(pool.getNumberOfChannels()).isEqualTo(2); + assertThat(pool.getNumberOfChannels()).isEqualTo(1); assertThat(pool.affinityKeyToChannelRef).doesNotContainKey("session"); + assertThat(pool.affinityKeyLastUsed).doesNotContainKey("session"); + assertThat(victim.getAffinityCount()).isEqualTo(0); } finally { if (pool != null) { pool.shutdownNow(); @@ -2315,13 +2493,6 @@ public ManagedChannel shutdownNow() { @Override public boolean awaitTermination(long timeout, TimeUnit unit) { - if (this.state == ConnectivityState.SHUTDOWN) { - return true; - } - try { - unit.sleep(timeout); - } catch (InterruptedException e) { - } return this.state == ConnectivityState.SHUTDOWN; }