From 38fa705e36e4fb3dbc57d1292263fdf33b6008c0 Mon Sep 17 00:00:00 2001 From: Sean Huh Date: Wed, 16 Sep 2026 13:09:59 -0700 Subject: [PATCH] Remove cycleId tracking from completion coordinator PiperOrigin-RevId: 982682079 --- .../planner/AsyncCompletionCoordinator.java | 111 +++++++++--------- .../AsyncCompletionCoordinatorTest.java | 105 ++++++++--------- 2 files changed, 103 insertions(+), 113 deletions(-) diff --git a/runtime/src/main/java/dev/cel/runtime/planner/AsyncCompletionCoordinator.java b/runtime/src/main/java/dev/cel/runtime/planner/AsyncCompletionCoordinator.java index 149212934..e95de7a30 100644 --- a/runtime/src/main/java/dev/cel/runtime/planner/AsyncCompletionCoordinator.java +++ b/runtime/src/main/java/dev/cel/runtime/planner/AsyncCompletionCoordinator.java @@ -87,9 +87,7 @@ enum WaitResult { private final ThreadLocal> continuationTrampoline; - @GuardedBy("lock") - private long cycleId; - + /** Monotonic staleness token bumped on every drain and snapshot. */ @GuardedBy("lock") private long debounceGeneration; @@ -118,14 +116,17 @@ void callCompleted(CelAsyncCall call) { CompletionSnapshot snapshot = null; synchronized (lock) { + // Check cancellation before activeCount() so late-finishing calls do not trip the unbalanced + // fail-safe on an evaluation that was already cancelled. + if (isCancelled) { + gate.release(); + return; + } // Check activeCount() <= 0 BEFORE release() to detect unbalanced completion misuse. if (gate.activeCount() <= 0) { unbalanced = true; } else { gate.release(); - if (isCancelled) { - return; - } completedBatch.add(call); if (!isWaiting) { return; @@ -133,10 +134,7 @@ void callCompleted(CelAsyncCall call) { // Take inFlight AFTER release() to capture remaining active calls for the drain strategy. snapshot = new CompletionSnapshot( - ImmutableList.copyOf(completedBatch), - gate.activeCount(), - cycleId, - ++debounceGeneration); + ImmutableList.copyOf(completedBatch), gate.activeCount(), ++debounceGeneration); } } @@ -155,7 +153,7 @@ void callCompleted(CelAsyncCall call) { failAndCancel(t); return; } - applyDrainAction(action, snapshot.cycleId, snapshot.debounceGeneration); + applyDrainAction(action, snapshot.debounceGeneration); } /** @@ -189,10 +187,7 @@ WaitResult waitForCompletions(Runnable continuationCallback) { snapshot = new CompletionSnapshot( - ImmutableList.copyOf(completedBatch), - gate.activeCount(), - this.cycleId, - ++this.debounceGeneration); + ImmutableList.copyOf(completedBatch), gate.activeCount(), ++this.debounceGeneration); } CelAsyncDrainAction action; @@ -207,6 +202,7 @@ WaitResult waitForCompletions(Runnable continuationCallback) { } boolean reevaluateNow = false; + ScheduledFuture timerToCancel = null; synchronized (lock) { if (isCancelled) { return WaitResult.CANCELLED; @@ -218,13 +214,16 @@ WaitResult waitForCompletions(Runnable continuationCallback) { if (action.shouldReevaluate() || gate.activeCount() == 0) { // The continuation in DrainResult is intentionally not dispatched here because // WaitResult.REEVALUATE_NOW instructs the calling thread to re-evaluate synchronously. - DrainResult unused = drainAndResetUnderLock(); + timerToCancel = drainAndResetUnderLock().timer; reevaluateNow = true; } } else { return WaitResult.REGISTERED; } } + if (timerToCancel != null) { + timerToCancel.cancel(false); + } if (reevaluateNow) { return WaitResult.REEVALUATE_NOW; } @@ -240,20 +239,21 @@ WaitResult waitForCompletions(Runnable continuationCallback) { failAndCancel(e); return WaitResult.CANCELLED; } - return scheduleDebounce(delayNanos, snapshot.cycleId, snapshot.debounceGeneration) + return scheduleDebounce(delayNanos, snapshot.debounceGeneration) ? WaitResult.REGISTERED : WaitResult.CANCELLED; } - private void applyDrainAction( - CelAsyncDrainAction action, long expectedCycleId, long expectedGen) { - boolean shouldReevaluate; + private void applyDrainAction(CelAsyncDrainAction action, long expectedGen) { DrainResult drainResult = null; synchronized (lock) { + // Bail out if a newer generation has superseded this action. + if (!isCurrentUnderLock(expectedGen)) { + return; + } // Live read: check gate.activeCount() == 0 under lock so we do not schedule an unnecessary // timer if all remaining calls completed while evaluating nextAction(). - shouldReevaluate = action.shouldReevaluate() || gate.activeCount() == 0; - if (shouldReevaluate && isCurrentUnderLock(expectedCycleId, expectedGen)) { + if (action.shouldReevaluate() || gate.activeCount() == 0) { drainResult = drainAndResetUnderLock(); } } @@ -266,9 +266,6 @@ private void applyDrainAction( } return; } - if (shouldReevaluate) { - return; - } Duration waitDuration = action.waitDuration(); if (!waitDuration.isZero()) { @@ -279,13 +276,13 @@ private void applyDrainAction( failAndCancel(e); return; } - boolean unusedScheduled = scheduleDebounce(delayNanos, expectedCycleId, expectedGen); + scheduleDebounce(delayNanos, expectedGen); return; } ScheduledFuture timerToCancel = null; synchronized (lock) { - if (isCurrentUnderLock(expectedCycleId, expectedGen)) { + if (isCurrentUnderLock(expectedGen)) { timerToCancel = cancelDebounceTimerUnderLock(); } } @@ -300,21 +297,39 @@ private void applyDrainAction( * @return false if the timer could not be scheduled, in which case the coordinator has already * been failed and cancelled. */ - private boolean scheduleDebounce(long nanos, long scheduledCycleId, long scheduledGen) { + private boolean scheduleDebounce(long nanos, long scheduledGen) { + synchronized (lock) { + if (isCancelled) { + return false; + } + if (!isCurrentUnderLock(scheduledGen)) { + return true; + } + } ScheduledFuture future; try { future = options .resolveScheduledExecutorService() - .schedule(() -> onDebounceFired(scheduledCycleId, scheduledGen), nanos, NANOSECONDS); + .schedule(() -> onDebounceFired(scheduledGen), nanos, NANOSECONDS); } catch (Throwable t) { - failAndCancel(t); - return false; + boolean shouldFail; + synchronized (lock) { + if (isCancelled) { + return false; + } + shouldFail = isCurrentUnderLock(scheduledGen); + } + if (shouldFail) { + failAndCancel(t); + return false; + } + return true; } ScheduledFuture redundantFuture = null; synchronized (lock) { - if (isCurrentUnderLock(scheduledCycleId, scheduledGen)) { + if (isCurrentUnderLock(scheduledGen)) { if (debounceTimer != null) { redundantFuture = debounceTimer; } @@ -330,22 +345,19 @@ private boolean scheduleDebounce(long nanos, long scheduledCycleId, long schedul } /** - * Returns true if the coordinator is still actively waiting on the cycle and debounce generation - * that produced the in-flight action, meaning the action is not stale. + * Returns true if the coordinator is still actively waiting on the debounce generation that + * produced the in-flight action, meaning the action is not stale. */ @GuardedBy("lock") - private boolean isCurrentUnderLock(long expectedCycleId, long expectedGen) { - return !isCancelled - && isWaiting - && this.cycleId == expectedCycleId - && this.debounceGeneration == expectedGen; + private boolean isCurrentUnderLock(long expectedGen) { + return !isCancelled && isWaiting && this.debounceGeneration == expectedGen; } @VisibleForTesting - void onDebounceFired(long firedCycleId, long firedGen) { + void onDebounceFired(long firedGen) { DrainResult drainResult = null; synchronized (lock) { - if (isCurrentUnderLock(firedCycleId, firedGen)) { + if (isCurrentUnderLock(firedGen)) { drainResult = drainAndResetUnderLock(); } } @@ -429,7 +441,6 @@ private void failAndCancel(Throwable t) { @GuardedBy("lock") @CheckReturnValue private DrainResult drainAndResetUnderLock() { - cycleId++; debounceGeneration++; isWaiting = false; completedBatch.clear(); @@ -474,13 +485,6 @@ boolean hasScheduledDebounceTimer() { } } - @VisibleForTesting - long cycleId() { - synchronized (lock) { - return cycleId; - } - } - @VisibleForTesting long debounceGeneration() { synchronized (lock) { @@ -498,14 +502,12 @@ boolean isCancelled() { private static final class CompletionSnapshot { final ImmutableList batch; final int inFlight; - final long cycleId; final long debounceGeneration; private CompletionSnapshot( - ImmutableList batch, int inFlight, long cycleId, long debounceGeneration) { + ImmutableList batch, int inFlight, long debounceGeneration) { this.batch = checkNotNull(batch, "batch must not be null"); this.inFlight = inFlight; - this.cycleId = cycleId; this.debounceGeneration = debounceGeneration; } } @@ -539,10 +541,5 @@ protected Deque initialValue() { return new ArrayDeque<>(); } }; - this.isWaiting = false; - this.isCancelled = false; - this.failureReported = false; - this.cycleId = 0; - this.debounceGeneration = 0; } } diff --git a/runtime/src/test/java/dev/cel/runtime/planner/AsyncCompletionCoordinatorTest.java b/runtime/src/test/java/dev/cel/runtime/planner/AsyncCompletionCoordinatorTest.java index 8c2073f43..80df9c84b 100644 --- a/runtime/src/test/java/dev/cel/runtime/planner/AsyncCompletionCoordinatorTest.java +++ b/runtime/src/test/java/dev/cel/runtime/planner/AsyncCompletionCoordinatorTest.java @@ -89,8 +89,7 @@ public void waitForCompletions_whenNoCallsInFlightAndEmptyBatch_returnsNoOutstan } @Test - public void - waitForCompletions_afterDrainConsumedBatchWithNoNewDispatch_returnsNoOutstandingWork() { + public void waitForCompletions_afterBatchConsumed_returnsNoOutstandingWork() { AsyncGate gate = AsyncGate.create(1); CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder() @@ -127,8 +126,7 @@ public void waitForCompletions_whenCallsInFlightAndEmptyBatch_returnsRegistered( } @Test - public void - waitForCompletions_whenDrainStrategySatisfiedImmediately_returnsReevaluateNowWithoutDispatch() { + public void waitForCompletions_drainSatisfiedImmediately_returnsReevaluateNow() { AsyncGate gate = AsyncGate.create(1); CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder() @@ -150,8 +148,7 @@ public void waitForCompletions_whenCallsInFlightAndEmptyBatch_returnsRegistered( } @Test - public void - waitForCompletions_whenDrainStrategyReturnsReevaluateWithInFlightCalls_returnsReevaluateNow() { + public void waitForCompletions_drainReevaluatesWithCallsInFlight_returnsReevaluateNow() { CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder() .setDrainStrategy(CelAsyncDrainStrategy.drainNone()) @@ -309,8 +306,7 @@ public void callCompleted_whenWaitingWithDrainAllStrategy_waitsWhileCallsRemainI } @Test - public void - callCompleted_whenWaitingWithDrainAllStrategy_triggersContinuationWhenFinalCallCompletes() { + public void callCompleted_drainAllAndFinalCallCompletes_triggersContinuation() { CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder() .setDrainStrategy(CelAsyncDrainStrategy.drainAll()) @@ -331,8 +327,7 @@ public void callCompleted_whenWaitingWithDrainAllStrategy_waitsWhileCallsRemainI } @Test - public void - callCompleted_whenWaitingWithDrainNoneStrategy_triggersContinuationWhileCallsRemainInFlight() { + public void callCompleted_drainNoneWithCallsInFlight_triggersContinuation() { CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder() .setDrainStrategy(CelAsyncDrainStrategy.drainNone()) @@ -383,8 +378,7 @@ public void callCompleted_whenDebounceTimerPending_resetsDebounceTimerForSliding } @Test - public void - callCompleted_whenDebounceTimerPendingAndNextActionIsWaitForMore_cancelsPendingDebounceTimer() { + public void callCompleted_pendingTimerAndNextActionWaitForMore_cancelsDebounceTimer() { TrackingScheduler scheduler = new TrackingScheduler(); try { AtomicInteger callCount = new AtomicInteger(); @@ -449,29 +443,6 @@ public void onDebounceFired_whenWaiting_triggersContinuation() { } } - @Test - public void onDebounceFired_whenCycleMismatch_doesNotExecuteContinuation() { - AtomicInteger executedCount = new AtomicInteger(); - Executor rejectingNullExecutor = - task -> { - requireNonNull(task, "task must not be null"); - executedCount.incrementAndGet(); - task.run(); - }; - CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder().build(); - AsyncGate gate = AsyncGate.create(1); - AsyncCompletionCoordinator coordinator = - AsyncCompletionCoordinator.create(options, gate, rejectingNullExecutor, t -> {}); - acquirePermit(gate); - registerWait(coordinator, () -> {}); - - coordinator.onDebounceFired(coordinator.cycleId() - 1, coordinator.debounceGeneration()); - - assertThat(executedCount.get()).isEqualTo(0); - assertThat(coordinator.isWaiting()).isTrue(); - assertThat(coordinator.hasContinuation()).isTrue(); - } - @Test public void onDebounceFired_whenDebounceGenerationMismatch_doesNotExecuteContinuation() { ScheduledThreadPoolExecutor scheduler = new ScheduledThreadPoolExecutor(1); @@ -493,7 +464,7 @@ public void onDebounceFired_whenDebounceGenerationMismatch_doesNotExecuteContinu long staleGen = coordinator.debounceGeneration(); coordinator.callCompleted(DUMMY_CALL); - coordinator.onDebounceFired(coordinator.cycleId(), staleGen); + coordinator.onDebounceFired(staleGen); assertThat(continuationRan.get()).isEqualTo(0); assertThat(coordinator.isWaiting()).isTrue(); @@ -646,8 +617,7 @@ public ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) } @Test - public void - scheduleDebounce_whenCoordinatorCancelledConcurrently_cancelsScheduledFutureWithoutInterrupt() { + public void scheduleDebounce_cancelledConcurrently_cancelsFutureWithoutInterrupt() { AtomicBoolean cancelledInsideScheduler = new AtomicBoolean(false); AsyncCompletionCoordinator[] coordinatorHolder = new AsyncCompletionCoordinator[1]; TrackingScheduler scheduler = @@ -856,24 +826,23 @@ public void staleTimerFromPreviousPass_doesNotTriggerContinuationOnSubsequentPas } @Test - public void drainAndReset_incrementsCycleIdAndClearsContinuation() { + public void drainAndReset_incrementsGenerationAndClearsContinuation() { CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder().build(); AsyncGate gate = AsyncGate.create(1); AsyncCompletionCoordinator coordinator = AsyncCompletionCoordinator.create(options, gate, Runnable::run, t -> {}); acquirePermit(gate); registerWait(coordinator, () -> {}); - long initialCycleId = coordinator.cycleId(); + long initialGeneration = coordinator.debounceGeneration(); coordinator.callCompleted(DUMMY_CALL); - assertThat(coordinator.cycleId()).isGreaterThan(initialCycleId); + assertThat(coordinator.debounceGeneration()).isGreaterThan(initialGeneration); assertThat(coordinator.hasContinuation()).isFalse(); } @Test - public void - waitForCompletions_lastCallCompletesDuringDrainStrategyEval_runsContinuationAndReturnsRegistered() { + public void waitForCompletions_lastCallCompletesDuringDrainEval_runsContinuationAndRegisters() { AtomicReference coordinatorRef = new AtomicReference<>(); CelAsyncDrainStrategy racingStrategy = new RacingDrainStrategy(coordinatorRef); CelAsyncEvaluationOptions options = @@ -1051,8 +1020,7 @@ public void waitForCompletions_whenCancelledDuringDrainStrategyEvaluation_return } @Test - public void - waitForCompletions_whenDrainStrategyReturnsWaitForMore_registersWithoutSchedulingTimer() { + public void waitForCompletions_drainWaitsForMore_registersWithoutTimer() { CelAsyncDrainStrategy waitForMoreStrategy = (batch, active) -> CelAsyncDrainAction.waitForMore(); CelAsyncEvaluationOptions options = @@ -1271,8 +1239,7 @@ public void applyDrainAction_whenGenerationStaleAndActionIsWaitForMore_keepsNewe } @Test - public void - failAndCancel_whenFailureCallbackThrows_suppressesCallbackExceptionAndCancelsCoordinator() { + public void failAndCancel_failureCallbackThrows_suppressesExceptionAndCancels() { RuntimeException strategyError = new RuntimeException("strategy failed"); CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder() @@ -1328,7 +1295,7 @@ public void dispatchContinuation_whenFailureCallbackThrows_doesNotEscapeAndSuppr } @Test - public void callCompleted_whenCancelledAndNoCallsInFlight_invokesFailureCallback() { + public void callCompleted_whenCancelled_doesNotInvokeFailureCallback() { AsyncGate gate = AsyncGate.create(1); CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder().build(); AtomicReference capturedFailure = new AtomicReference<>(); @@ -1338,6 +1305,36 @@ public void callCompleted_whenCancelledAndNoCallsInFlight_invokesFailureCallback coordinator.callCompleted(DUMMY_CALL); + assertThat(capturedFailure.get()).isNull(); + } + + @Test + public void callCompleted_whenCancelledWithCallInFlight_releasesPermitWithoutFailing() { + AsyncGate gate = AsyncGate.create(1); + CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder().build(); + AtomicReference capturedFailure = new AtomicReference<>(); + AsyncCompletionCoordinator coordinator = + AsyncCompletionCoordinator.create(options, gate, Runnable::run, capturedFailure::set); + acquirePermit(gate); + coordinator.cancel(); + + coordinator.callCompleted(DUMMY_CALL); + + assertThat(capturedFailure.get()).isNull(); + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(coordinator.hasPendingBatch()).isFalse(); + } + + @Test + public void callCompleted_whenNoCallsInFlight_invokesFailureCallback() { + AsyncGate gate = AsyncGate.create(1); + CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder().build(); + AtomicReference capturedFailure = new AtomicReference<>(); + AsyncCompletionCoordinator coordinator = + AsyncCompletionCoordinator.create(options, gate, Runnable::run, capturedFailure::set); + + coordinator.callCompleted(DUMMY_CALL); + Throwable failure = capturedFailure.get(); assertThat(failure).isInstanceOf(IllegalStateException.class); assertThat(failure).hasMessageThat().contains("callCompleted called with no active calls"); @@ -1425,13 +1422,12 @@ public void failAndCancel_whenCalledMultipleTimes_invokesFailureCallbackOnlyOnce } @Test - public void create_initialState_hasZeroGenerationsAndCleanDefaults() { + public void create_initialState_hasZeroGenerationAndCleanDefaults() { CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder().build(); AsyncGate gate = AsyncGate.create(1); AsyncCompletionCoordinator coordinator = AsyncCompletionCoordinator.create(options, gate, Runnable::run, t -> {}); - assertThat(coordinator.cycleId()).isEqualTo(0); assertThat(coordinator.debounceGeneration()).isEqualTo(0); assertThat(coordinator.isWaiting()).isFalse(); assertThat(coordinator.isCancelled()).isFalse(); @@ -1441,8 +1437,7 @@ public void create_initialState_hasZeroGenerationsAndCleanDefaults() { } @Test - public void - callCompleted_whenDebounceTimerPendingAndNextActionIsReevaluate_cancelsPendingDebounceTimerWithoutInterrupt() { + public void callCompleted_pendingTimerAndNextActionReevaluate_cancelsTimerWithoutInterrupt() { TrackingScheduler scheduler = new TrackingScheduler(); try { AtomicInteger callCount = new AtomicInteger(); @@ -1475,8 +1470,7 @@ public void create_initialState_hasZeroGenerationsAndCleanDefaults() { } @Test - public void - waitForCompletions_whenPendingBatchAndStrategyWaits_schedulesDebounceTimerAndReturnsRegistered() { + public void waitForCompletions_pendingBatchAndStrategyWaits_schedulesTimerAndRegisters() { TrackingScheduler scheduler = new TrackingScheduler(); try { CelAsyncEvaluationOptions options = @@ -1502,8 +1496,7 @@ public void create_initialState_hasZeroGenerationsAndCleanDefaults() { } @Test - public void - waitForCompletions_whenWaitingWithDrainAllStrategyAndCallsInFlight_returnsRegistered() { + public void waitForCompletions_drainAllWithCallsInFlight_returnsRegistered() { AsyncGate gate = AsyncGate.create(2); CelAsyncEvaluationOptions options = CelAsyncEvaluationOptions.builder()