From 491a17930b1d0dd130bb8c6410c69f167e748305 Mon Sep 17 00:00:00 2001 From: Sean Huh Date: Fri, 18 Sep 2026 17:32:38 -0700 Subject: [PATCH] Add AsyncCallStateTracker for call deduplication and state management PiperOrigin-RevId: 984153900 --- .../java/dev/cel/runtime/RuntimeEquality.java | 8 +- .../planner/AsyncCallStateTracker.java | 298 +++++++ .../java/dev/cel/runtime/planner/BUILD.bazel | 13 +- .../dev/cel/runtime/RuntimeEqualityTest.java | 22 +- .../planner/AsyncCallStateTrackerTest.java | 828 ++++++++++++++++++ .../java/dev/cel/runtime/planner/BUILD.bazel | 4 + 6 files changed, 1165 insertions(+), 8 deletions(-) create mode 100644 runtime/src/main/java/dev/cel/runtime/planner/AsyncCallStateTracker.java create mode 100644 runtime/src/test/java/dev/cel/runtime/planner/AsyncCallStateTrackerTest.java diff --git a/runtime/src/main/java/dev/cel/runtime/RuntimeEquality.java b/runtime/src/main/java/dev/cel/runtime/RuntimeEquality.java index 4814dc269..a9b607e2b 100644 --- a/runtime/src/main/java/dev/cel/runtime/RuntimeEquality.java +++ b/runtime/src/main/java/dev/cel/runtime/RuntimeEquality.java @@ -135,12 +135,12 @@ public Optional findInMap(Map map, Object index) { * comparable even if they are not of the same type, where type differences are usually trivially * false. */ - @SuppressWarnings({"rawtypes", "unchecked"}) + @SuppressWarnings({"rawtypes", "unchecked", "ReferenceEquality"}) public boolean objectEquals(Object x, Object y) { if (celOptions.disableCelStandardEquality()) { return Objects.equals(x, y); } - if (x == y) { + if (x == y && !isNan(x)) { return true; } x = runtimeHelpers.adaptValue(x); @@ -278,6 +278,10 @@ private static Optional unsignedToLongLossless(UnsignedLong v) { return Optional.empty(); } + private static boolean isNan(Object value) { + return value instanceof Number && Double.isNaN(((Number) value).doubleValue()); + } + RuntimeEquality(RuntimeHelpers runtimeHelpers, CelOptions celOptions) { this.runtimeHelpers = runtimeHelpers; this.celOptions = celOptions; diff --git a/runtime/src/main/java/dev/cel/runtime/planner/AsyncCallStateTracker.java b/runtime/src/main/java/dev/cel/runtime/planner/AsyncCallStateTracker.java new file mode 100644 index 000000000..ee4b8a5a4 --- /dev/null +++ b/runtime/src/main/java/dev/cel/runtime/planner/AsyncCallStateTracker.java @@ -0,0 +1,298 @@ +// 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 dev.cel.runtime.planner; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import javax.annotation.concurrent.ThreadSafe; +import dev.cel.common.exceptions.CelRuntimeException; +import dev.cel.common.values.CelValueConverter; +import dev.cel.runtime.AccumulatedUnknowns; +import dev.cel.runtime.CelAsyncFunctionOverload; +import dev.cel.runtime.CelAsyncObserver; +import dev.cel.runtime.CelEvaluationException; +import dev.cel.runtime.InterpreterUtil; +import dev.cel.runtime.RuntimeEquality; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicLong; +import org.jspecify.annotations.Nullable; + +/** + * Tracks the registry and cache of all asynchronous function calls made during an expression + * evaluation. + */ +@ThreadSafe +// CEL-Internal-4 +final class AsyncCallStateTracker { + private final AtomicLong callIdGenerator = new AtomicLong(1); + private final ConcurrentMap> recordsByBucket = + new ConcurrentHashMap<>(); + private final ConcurrentMap recordsById = new ConcurrentHashMap<>(); + private final RuntimeEquality runtimeEquality; + + static AsyncCallStateTracker create(RuntimeEquality runtimeEquality) { + return new AsyncCallStateTracker(runtimeEquality); + } + + /** + * Returns the resolved result for a previously completed call matching {@code (exprId, + * overloadId, args)}, throws a runtime exception if the call failed or was cancelled, or + * registers and returns an {@link AccumulatedUnknowns} with the call's tracking ID if pending. + */ + Object recordOrGet( + long exprId, + String functionName, + String overloadId, + Object[] args, + CelAsyncFunctionOverload overload, + CelValueConverter celValueConverter) { + checkNotNull(functionName); + checkNotNull(overloadId); + checkNotNull(args); + checkNotNull(overload); + checkNotNull(celValueConverter); + int bucketKey = AsyncCallRecord.hashCall(exprId, overloadId, args); + CopyOnWriteArrayList bucket = recordsByBucket.get(bucketKey); + if (bucket != null) { + for (int i = 0; i < bucket.size(); i++) { + AsyncCallRecord existing = bucket.get(i); + if (existing.matches(exprId, functionName, overloadId, args, runtimeEquality)) { + return resolveRecord(existing, celValueConverter); + } + } + } + + bucket = recordsByBucket.computeIfAbsent(bucketKey, k -> new CopyOnWriteArrayList<>()); + AsyncCallRecord record = null; + synchronized (bucket) { + for (int i = 0; i < bucket.size(); i++) { + AsyncCallRecord existing = bucket.get(i); + if (existing.matches(exprId, functionName, overloadId, args, runtimeEquality)) { + record = existing; + break; + } + } + if (record == null) { + long callId = callIdGenerator.getAndIncrement(); + record = AsyncCallRecord.create(callId, exprId, functionName, overloadId, args, overload); + recordsById.put(callId, record); + bucket.add(record); + } + } + + return resolveRecord(record, celValueConverter); + } + + /** + * Launches every not-yet-started call in {@code requiredCallIds}, subject to {@code gate} + * admission control. + * + *

{@code executor} must run or reject each task; one that silently discards tasks strands the + * call's concurrency permit. + */ + void dispatchPendingCalls( + Set requiredCallIds, + Executor executor, + AsyncGate gate, + AsyncCompletionCoordinator coordinator, + @Nullable CelAsyncObserver observer) { + checkNotNull(requiredCallIds); + checkNotNull(executor); + checkNotNull(gate); + checkNotNull(coordinator); + for (Long callId : ImmutableList.sortedCopyOf(requiredCallIds)) { + AsyncCallRecord record = recordsById.get(callId); + if (record != null && record.state() == AsyncCallRecord.State.NOT_STARTED) { + tryLaunch(record, executor, gate, coordinator, observer); + } + } + } + + @VisibleForTesting + void tryLaunch( + AsyncCallRecord record, + Executor executor, + AsyncGate gate, + AsyncCompletionCoordinator coordinator, + @Nullable CelAsyncObserver observer) { + if (!gate.tryAcquire()) { + return; + } + if (!record.markRunning()) { + gate.release(); + return; + } + + try { + if (observer != null) { + observer.onCallStarted(record, ImmutableList.copyOf(record.args())); + } + executor.execute(() -> executeAsyncCall(record, coordinator, observer)); + } catch (RuntimeException e) { + handleFailure(record, e, coordinator, observer); + } + } + + private static void executeAsyncCall( + AsyncCallRecord record, + AsyncCompletionCoordinator coordinator, + @Nullable CelAsyncObserver observer) { + ListenableFuture future; + try { + if (record.isCancelled()) { + throw new CancellationException("Async call was cancelled before dispatch"); + } + future = + checkNotNull( + record.overload().applyAsync(record.args()), + "Async function '%s' returned a null ListenableFuture", + record.functionName()); + record.setInFlightFuture(future); + } catch (CelEvaluationException | RuntimeException e) { + handleFailure(record, e, coordinator, observer); + return; + } + + Futures.addCallback( + future, + new FutureCallback() { + @Override + public void onSuccess(Object result) { + handleSuccess(record, result, coordinator, observer); + } + + @Override + public void onFailure(Throwable t) { + handleFailure(record, t, coordinator, observer); + } + }, + directExecutor()); + } + + private Object resolveRecord(AsyncCallRecord record, CelValueConverter celValueConverter) { + switch (record.state()) { + case SUCCESS: + Object rawResult = record.result().orElseThrow(AssertionError::new); + return InterpreterUtil.maybeAdaptToAccumulatedUnknowns( + celValueConverter.maybeUnwrap(celValueConverter.toRuntimeValue(rawResult))); + case FAILURE: + Throwable error = record.error().orElseThrow(AssertionError::new); + if (error instanceof CelRuntimeException) { + throw (CelRuntimeException) error; + } + String errorMessage = + error.getMessage() != null ? error.getMessage() : error.getClass().getSimpleName(); + throw new IllegalArgumentException( + String.format("Async function '%s' failed: %s", record.functionName(), errorMessage), + error); + case RUNNING: + case NOT_STARTED: + return AccumulatedUnknowns.createForAsyncCall(record.exprId(), record.callId()); + case CANCELLED: + throw new CancellationException( + String.format("Async function '%s' was cancelled", record.functionName())); + } + throw new AssertionError("Unexpected record state: " + record.state()); + } + + boolean hasInFlightCalls() { + for (AsyncCallRecord record : recordsById.values()) { + if (record.state() == AsyncCallRecord.State.RUNNING) { + return true; + } + } + return false; + } + + void cancelInFlight() { + for (AsyncCallRecord record : recordsById.values()) { + record.cancelInFlight(); + } + } + + private static void handleSuccess( + AsyncCallRecord record, + @Nullable Object result, + AsyncCompletionCoordinator coordinator, + @Nullable CelAsyncObserver observer) { + if (result == null) { + handleFailure( + record, + new NullPointerException( + String.format("Async function '%s' returned a null result", record.functionName())), + coordinator, + observer); + return; + } + record.complete(result); + reportCompletion(record, result, /* error= */ null, coordinator, observer); + } + + private static void handleFailure( + AsyncCallRecord record, + Throwable error, + AsyncCompletionCoordinator coordinator, + @Nullable CelAsyncObserver observer) { + record.fail(error); + reportCompletion(record, /* result= */ null, error, coordinator, observer); + } + + /** + * Notifies the observer and completion coordinator of a launched call's terminal outcome at most + * once. + */ + private static void reportCompletion( + AsyncCallRecord record, + @Nullable Object result, + @Nullable Throwable error, + AsyncCompletionCoordinator coordinator, + @Nullable CelAsyncObserver observer) { + if (!record.markCompletionReported()) { + return; + } + try { + if (observer != null) { + observer.onCallFinished(record, result, error); + } + } finally { + coordinator.callCompleted(record); + } + } + + @VisibleForTesting + ConcurrentMap> recordsByBucket() { + return recordsByBucket; + } + + @VisibleForTesting + ConcurrentMap recordsById() { + return recordsById; + } + + private AsyncCallStateTracker(RuntimeEquality runtimeEquality) { + this.runtimeEquality = checkNotNull(runtimeEquality); + } +} diff --git a/runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel b/runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel index 009672e32..0ab4f1beb 100644 --- a/runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel +++ b/runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel @@ -219,12 +219,23 @@ java_library( java_library( name = "async_call_state_tracker", - srcs = ["AsyncCallRecord.java"], + srcs = [ + "AsyncCallRecord.java", + "AsyncCallStateTracker.java", + ], tags = [ ], deps = [ + ":async_completion_coordinator", + ":async_gate", + "//common/exceptions:runtime_exception", + "//common/values", + "//runtime:accumulated_unknowns", "//runtime:async_call", + "//runtime:async_observer", + "//runtime:evaluation_exception", "//runtime:function_overload", + "//runtime:interpreter_util", "//runtime:runtime_equality", "@maven//:com_google_code_findbugs_annotations", "@maven//:com_google_errorprone_error_prone_annotations", diff --git a/runtime/src/test/java/dev/cel/runtime/RuntimeEqualityTest.java b/runtime/src/test/java/dev/cel/runtime/RuntimeEqualityTest.java index a5617beee..50c39ed58 100644 --- a/runtime/src/test/java/dev/cel/runtime/RuntimeEqualityTest.java +++ b/runtime/src/test/java/dev/cel/runtime/RuntimeEqualityTest.java @@ -57,12 +57,24 @@ private void assertEqualityAndHashCode(RuntimeEquality runtimeEquality, Object o public void objectEquals_messageLite_throws() { RuntimeEquality runtimeEquality = RuntimeEquality.create(RuntimeHelpers.create(), CelOptions.DEFAULT); + TestAllTypes.Builder builder = TestAllTypes.newBuilder(); + TestAllTypes defaultInstance = TestAllTypes.getDefaultInstance(); // Unimplemented until CelLiteDescriptor is available. - assertThrows( - UnsupportedOperationException.class, - () -> - runtimeEquality.objectEquals( - TestAllTypes.newBuilder(), TestAllTypes.getDefaultInstance())); + UnsupportedOperationException e = + assertThrows( + UnsupportedOperationException.class, + () -> runtimeEquality.objectEquals(builder, defaultInstance)); + + assertThat(e).hasMessageThat().contains("Not implemented yet"); + } + + @Test + public void objectEquals_nanWithIdenticalReference_returnsFalse() { + RuntimeEquality runtimeEquality = + RuntimeEquality.create(RuntimeHelpers.create(), CelOptions.DEFAULT); + Double nan = Double.NaN; + + assertThat(runtimeEquality.objectEquals(nan, nan)).isFalse(); } } diff --git a/runtime/src/test/java/dev/cel/runtime/planner/AsyncCallStateTrackerTest.java b/runtime/src/test/java/dev/cel/runtime/planner/AsyncCallStateTrackerTest.java new file mode 100644 index 000000000..f222ba3e2 --- /dev/null +++ b/runtime/src/test/java/dev/cel/runtime/planner/AsyncCallStateTrackerTest.java @@ -0,0 +1,828 @@ +// 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 dev.cel.runtime.planner; + +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.Iterables.getOnlyElement; +import static com.google.common.truth.Truth.assertThat; +import static com.google.common.util.concurrent.Futures.immediateFailedFuture; +import static com.google.common.util.concurrent.Futures.immediateFuture; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.Assert.assertThrows; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.primitives.UnsignedLong; +import com.google.common.util.concurrent.ForwardingListenableFuture.SimpleForwardingListenableFuture; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import com.google.errorprone.annotations.Immutable; +import javax.annotation.concurrent.ThreadSafe; +import com.google.testing.junit.testparameterinjector.TestParameter; +import com.google.testing.junit.testparameterinjector.TestParameterInjector; +import dev.cel.common.CelOptions; +import dev.cel.common.exceptions.CelDivideByZeroException; +import dev.cel.common.exceptions.CelRuntimeException; +import dev.cel.common.values.CelValueConverter; +import dev.cel.common.values.NullValue; +import dev.cel.runtime.AccumulatedUnknowns; +import dev.cel.runtime.CelAsyncCall; +import dev.cel.runtime.CelAsyncEvaluationOptions; +import dev.cel.runtime.CelAsyncFunctionOverload; +import dev.cel.runtime.CelAsyncObserver; +import dev.cel.runtime.RuntimeEquality; +import dev.cel.runtime.RuntimeHelpers; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; + +@RunWith(TestParameterInjector.class) +@SuppressWarnings("Immutable") +public final class AsyncCallStateTrackerTest { + + private final RuntimeEquality runtimeEquality = + RuntimeEquality.create(RuntimeHelpers.create(), CelOptions.DEFAULT); + private final Executor directExecutor = directExecutor(); + private final AsyncCallStateTracker tracker = AsyncCallStateTracker.create(runtimeEquality); + private final AsyncGate gate = AsyncGate.create(1); + private final AsyncCompletionCoordinator coordinator = newCoordinator(gate, directExecutor); + private final RecordingObserver observer = new RecordingObserver(); + + @Test + public void dispatchPendingCalls_onlyLaunchesRequiredCallIds() throws Exception { + AtomicBoolean call1Executed = new AtomicBoolean(false); + AtomicBoolean call2Executed = new AtomicBoolean(false); + AccumulatedUnknowns unk1 = + recordCall( + 1L, + "func1", + "a", + args -> { + call1Executed.set(true); + return immediateFuture("res1"); + }); + recordCall( + 2L, + "func2", + "b", + args -> { + call2Executed.set(true); + return immediateFuture("res2"); + }); + + tracker.dispatchPendingCalls( + unk1.callIds(), directExecutor, gate, coordinator, /* observer= */ null); + + assertThat(call1Executed.get()).isTrue(); + assertThat(call2Executed.get()).isFalse(); + } + + @Test + public void dispatchPendingCalls_cancelledWhileQueued_abortsOverloadAndNotifiesObserver( + @TestParameter boolean withObserver) throws Exception { + List queuedTasks = new ArrayList<>(); + AtomicBoolean overloadExecuted = new AtomicBoolean(false); + AccumulatedUnknowns unknowns = + recordDefaultCall( + args -> { + overloadExecuted.set(true); + return immediateFuture("ok"); + }); + tracker.dispatchPendingCalls( + unknowns.callIds(), queuedTasks::add, gate, coordinator, withObserver ? observer : null); + + tracker.cancelInFlight(); + queuedTasks.get(0).run(); + + assertThat(overloadExecuted.get()).isFalse(); + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(tracker.hasInFlightCalls()).isFalse(); + if (withObserver) { + assertThat(observer.startedCalls()).hasSize(1); + assertThat(getOnlyElement(observer.finishedCalls()).error) + .isInstanceOf(CancellationException.class); + } + } + + @Test + public void dispatchPendingCalls_whenGateFull_defersUntilPermitReleased( + @TestParameter boolean releaseAndRetry) throws Exception { + checkState(gate.tryAcquire(), "Failed to acquire permit"); + List queuedTasks = new ArrayList<>(); + AtomicBoolean overloadCalled = new AtomicBoolean(false); + AccumulatedUnknowns unknowns = + recordDefaultCall( + args -> { + overloadCalled.set(true); + return immediateFuture("ok"); + }); + + tracker.dispatchPendingCalls( + unknowns.callIds(), queuedTasks::add, gate, coordinator, /* observer= */ null); + if (releaseAndRetry) { + gate.release(); + tracker.dispatchPendingCalls( + unknowns.callIds(), queuedTasks::add, gate, coordinator, /* observer= */ null); + queuedTasks.get(0).run(); + } + + assertThat(overloadCalled.get()).isEqualTo(releaseAndRetry); + assertThat(gate.activeCount()).isEqualTo(releaseAndRetry ? 0 : 1); + assertThat(tracker.hasInFlightCalls()).isFalse(); + } + + @Test + public void recordOrGet_existingKey_reusesCallIdAndAllocatesNewIdForDistinctKey() + throws Exception { + AccumulatedUnknowns first = recordCall(10L, "fn", "x", args -> SettableFuture.create()); + AccumulatedUnknowns second = recordCall(10L, "fn", "x", args -> SettableFuture.create()); + AccumulatedUnknowns third = recordCall(20L, "fn", "y", args -> SettableFuture.create()); + + assertThat(first.callIds()).containsExactly(1L); + assertThat(second.callIds()).containsExactly(1L); + assertThat(third.callIds()).containsExactly(2L); + } + + @Test + public void recordOrGet_bucketHashCollision_disambiguatesViaMatches() throws Exception { + // Both arguments are complex types (lists) so hashArg yields COMPLEX_HASH_MARKER for both, + // causing a bucket collision under the same (exprId, overloadId). + AccumulatedUnknowns first = + recordCall(10L, "fn", ImmutableList.of("a"), args -> SettableFuture.create()); + AccumulatedUnknowns second = + recordCall(10L, "fn", ImmutableList.of("b"), args -> SettableFuture.create()); + AccumulatedUnknowns firstAgain = + recordCall(10L, "fn", ImmutableList.of("a"), args -> SettableFuture.create()); + AccumulatedUnknowns secondAgain = + recordCall(10L, "fn", ImmutableList.of("b"), args -> SettableFuture.create()); + + assertThat(first.callIds()).containsExactly(1L); + assertThat(second.callIds()).containsExactly(2L); + assertThat(firstAgain.callIds()).containsExactly(1L); + assertThat(secondAgain.callIds()).containsExactly(2L); + } + + @Test + public void recordOrGet_celEqualArguments_reusesCallId() throws Exception { + AccumulatedUnknowns first = recordCall(10L, "fn", 1L, args -> SettableFuture.create()); + AccumulatedUnknowns second = recordCall(10L, "fn", 1.0d, args -> SettableFuture.create()); + + assertThat(first.callIds()).containsExactly(1L); + assertThat(second.callIds()).containsExactly(1L); + } + + @Test + public void recordOrGet_nanArguments_reusesCallId() throws Exception { + AccumulatedUnknowns first = recordCall(10L, "fn", Double.NaN, args -> SettableFuture.create()); + AccumulatedUnknowns second = recordCall(10L, "fn", Float.NaN, args -> SettableFuture.create()); + + assertThat(first.callIds()).containsExactly(1L); + assertThat(second.callIds()).containsExactly(1L); + } + + @Test + public void recordOrGet_signedZeroArguments_reusesCallId() throws Exception { + AccumulatedUnknowns first = recordCall(10L, "fn", 0.0d, args -> SettableFuture.create()); + AccumulatedUnknowns second = recordCall(10L, "fn", -0.0d, args -> SettableFuture.create()); + + assertThat(first.callIds()).containsExactly(1L); + assertThat(second.callIds()).containsExactly(1L); + } + + @Test + public void recordOrGet_unsignedLongAndLongArguments_reusesCallIdWhenEqual() throws Exception { + AccumulatedUnknowns first = + recordCall(10L, "fn", UnsignedLong.valueOf(42L), args -> SettableFuture.create()); + AccumulatedUnknowns second = recordCall(10L, "fn", 42L, args -> SettableFuture.create()); + + assertThat(first.callIds()).containsExactly(1L); + assertThat(second.callIds()).containsExactly(1L); + } + + @Test + public void recordOrGet_optimisticRead_doesNotBlockOnBucketLock() throws Exception { + CelAsyncFunctionOverload overload = args -> SettableFuture.create(); + // Seeds a second entry in the same bucket so the optimistic read iterates past one element. + recordCall(1L, "fn", ImmutableList.of("a"), overload); + AccumulatedUnknowns second = recordCall(1L, "fn", ImmutableList.of("b"), overload); + int bucketKey = + AsyncCallRecord.hashCall(1L, "fn_overload", new Object[] {ImmutableList.of("b")}); + CopyOnWriteArrayList bucket = tracker.recordsByBucket().get(bucketKey); + checkNotNull(bucket, "Bucket must not be null"); + checkState(bucket.size() >= 2, "Bucket must contain at least two colliding calls"); + + CountDownLatch lockAcquired = new CountDownLatch(1); + CountDownLatch releaseLock = new CountDownLatch(1); + Thread blockerThread = + new Thread( + () -> { + synchronized (bucket) { + lockAcquired.countDown(); + try { + releaseLock.await(5, SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + }); + SettableFuture readResult = SettableFuture.create(); + Thread readerThread = + new Thread( + () -> { + try { + // Read the second colliding record to ensure iteration covers multiple elements + // lock-free + readResult.set(recordCall(1L, "fn", ImmutableList.of("b"), overload)); + } catch (Throwable t) { + readResult.setException(t); + } + }); + try { + blockerThread.start(); + checkState(lockAcquired.await(5, SECONDS), "blockerThread failed to acquire bucket lock"); + + readerThread.start(); + AccumulatedUnknowns result = readResult.get(1, SECONDS); + + assertThat(result.callIds()).containsExactlyElementsIn(second.callIds()); + } finally { + releaseLock.countDown(); + blockerThread.join(5000); + readerThread.join(5000); + } + } + + @Test + public void recordOrGet_concurrentRegistrationRace_reusesExistingRecordInSlowPath() + throws Exception { + CelAsyncFunctionOverload overload = args -> SettableFuture.create(); + Object[] args = new Object[] {1L}; + int bucketKey = AsyncCallRecord.hashCall(1L, "fn_overload", args); + CopyOnWriteArrayList bucket = new CopyOnWriteArrayList<>(); + tracker.recordsByBucket().put(bucketKey, bucket); + + CountDownLatch blockerLocked = new CountDownLatch(1); + CountDownLatch populateAndRelease = new CountDownLatch(1); + Thread blockerThread = + new Thread( + () -> { + synchronized (bucket) { + blockerLocked.countDown(); + try { + populateAndRelease.await(5, SECONDS); + AsyncCallRecord preExisting = + AsyncCallRecord.create(99L, 1L, "fn", "fn_overload", args, overload); + tracker.recordsById().put(99L, preExisting); + bucket.add(preExisting); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + }); + SettableFuture callerResult = SettableFuture.create(); + Thread callerThread = + new Thread( + () -> { + try { + callerResult.set( + (AccumulatedUnknowns) + tracker.recordOrGet( + 1L, + "fn", + "fn_overload", + args, + overload, + CelValueConverter.getDefaultInstance())); + } catch (Throwable t) { + callerResult.setException(t); + } + }); + try { + blockerThread.start(); + checkState(blockerLocked.await(5, SECONDS), "blockerThread failed to acquire lock"); + + callerThread.start(); + long deadline = System.currentTimeMillis() + 5000; + while (callerThread.getState() != Thread.State.BLOCKED) { + if (callerThread.getState() == Thread.State.TERMINATED) { + callerResult.get(); + throw new AssertionError("callerThread terminated unexpectedly without blocking"); + } + if (System.currentTimeMillis() > deadline) { + throw new AssertionError( + "callerThread never entered BLOCKED state; state is " + callerThread.getState()); + } + Thread.sleep(10); + } + populateAndRelease.countDown(); + AccumulatedUnknowns unknowns = callerResult.get(5, SECONDS); + + assertThat(unknowns.callIds()).containsExactly(99L); + } finally { + populateAndRelease.countDown(); + blockerThread.join(5000); + callerThread.join(5000); + } + } + + @Test + public void recordOrGet_concurrentBucketHashCollision_registersAllCallsSafely() throws Exception { + List results = new CopyOnWriteArrayList<>(); + + runConcurrently( + 16, + () -> { + int id = results.size(); + results.add( + recordCall( + 10L, "fn", ImmutableList.of("arg_" + id), args -> SettableFuture.create())); + }); + + assertThat(results).hasSize(16); + assertThat(tracker.hasInFlightCalls()).isFalse(); + } + + private enum OverloadFailureMode { + THROWS_SYNCHRONOUSLY, + FAILED_FUTURE, + RETURNS_NULL_FUTURE, + FUTURE_COMPLETES_WITH_NULL + } + + @Test + public void dispatchPendingCalls_overloadFails_notifiesObserverAndReleasesPermit( + @TestParameter OverloadFailureMode failureMode, @TestParameter boolean withObserver) + throws Exception { + RuntimeException expectedError = new RuntimeException("fail"); + AccumulatedUnknowns unknowns = + recordDefaultCall( + args -> { + switch (failureMode) { + case THROWS_SYNCHRONOUSLY: + throw expectedError; + case FAILED_FUTURE: + return immediateFailedFuture(expectedError); + case RETURNS_NULL_FUTURE: + return null; + case FUTURE_COMPLETES_WITH_NULL: + return immediateFuture(null); + } + throw new AssertionError(); + }); + + tracker.dispatchPendingCalls( + unknowns.callIds(), directExecutor, gate, coordinator, withObserver ? observer : null); + + boolean isNullFailure = + failureMode == OverloadFailureMode.RETURNS_NULL_FUTURE + || failureMode == OverloadFailureMode.FUTURE_COMPLETES_WITH_NULL; + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(tracker.hasInFlightCalls()).isFalse(); + if (withObserver) { + assertThat(getOnlyElement(observer.startedArgs())).containsExactly("x"); + Throwable recordedError = getOnlyElement(observer.finishedCalls()).error; + if (isNullFailure) { + assertThat(recordedError).isInstanceOf(NullPointerException.class); + } else { + assertThat(recordedError).isSameInstanceAs(expectedError); + } + } + IllegalArgumentException evalException = + assertThrows( + IllegalArgumentException.class, + () -> getDefaultCall(args -> immediateFuture("unused"))); + if (isNullFailure) { + assertThat(evalException).hasCauseThat().isInstanceOf(NullPointerException.class); + } else { + assertThat(evalException).hasCauseThat().isSameInstanceAs(expectedError); + } + } + + @Test + public void dispatchPendingCalls_executorRejection_releasesPermitAndFailsRecord() + throws Exception { + AccumulatedUnknowns unknowns = recordDefaultCall(args -> immediateFuture("done")); + Executor rejectingExecutor = + cmd -> { + throw new RejectedExecutionException("pool full"); + }; + + tracker.dispatchPendingCalls( + unknowns.callIds(), rejectingExecutor, gate, coordinator, /* observer= */ null); + + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(tracker.hasInFlightCalls()).isFalse(); + } + + @Test + public void recordOrGet_afterSuccess_returnsResolvedValueAndNotifiesObserver( + @TestParameter boolean nullSentinelValues) throws Exception { + Object arg = nullSentinelValues ? NullValue.NULL_VALUE : "x"; + Object expectedResult = nullSentinelValues ? NullValue.NULL_VALUE : "syncSuccess"; + SettableFuture future = SettableFuture.create(); + AccumulatedUnknowns unknowns = recordCall(10L, "fn", arg, args -> future); + tracker.dispatchPendingCalls(unknowns.callIds(), directExecutor, gate, coordinator, observer); + future.set(expectedResult); + + Object result = recordOrGetCall(10L, "fn", arg, args -> future); + + FinishedCall finished = getOnlyElement(observer.finishedCalls()); + assertThat(result).isEqualTo(expectedResult); + assertThat(getOnlyElement(observer.startedCalls()).call.functionName()).isEqualTo("fn"); + assertThat(getOnlyElement(observer.startedArgs())).containsExactly(arg); + assertThat(finished.result).isEqualTo(expectedResult); + assertThat(finished.call.functionName()).isEqualTo("fn"); + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(tracker.hasInFlightCalls()).isFalse(); + } + + @Test + public void dispatchPendingCalls_concurrentRace_threadContentionHandledSafely() throws Exception { + AtomicInteger callsDispatched = new AtomicInteger(0); + AccumulatedUnknowns unknowns = + recordDefaultCall( + args -> { + callsDispatched.incrementAndGet(); + return immediateFuture("result"); + }); + + runConcurrently(8, () -> dispatchDefaultCalls(unknowns)); + + assertThat(callsDispatched.get()).isEqualTo(1); + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(tracker.hasInFlightCalls()).isFalse(); + } + + @Test + public void cancelInFlight_beforeFutureCompletes_releasesPermitAndNotifiesObserver( + @TestParameter boolean succeedsAfterCancel) throws Exception { + SettableFuture underlyingFuture = SettableFuture.create(); + dispatchCallWithObserver(args -> nonCancellableFuture(underlyingFuture)); + RuntimeException lateError = new RuntimeException("late_failure"); + + tracker.cancelInFlight(); + if (succeedsAfterCancel) { + underlyingFuture.set("late_success"); + } else { + underlyingFuture.setException(lateError); + } + + assertThat(observer.startedCalls()).hasSize(1); + FinishedCall finished = getOnlyElement(observer.finishedCalls()); + assertThat(finished.result).isEqualTo(succeedsAfterCancel ? "late_success" : null); + assertThat(finished.error).isEqualTo(succeedsAfterCancel ? null : lateError); + assertThat(finished.call.functionName()).isEqualTo("fn"); + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(tracker.hasInFlightCalls()).isFalse(); + } + + @Test + public void recordOrGet_whenInFlight_returnsAccumulatedUnknownsWithSameCallId() throws Exception { + SettableFuture pendingFuture = SettableFuture.create(); + AccumulatedUnknowns initial = recordDefaultCall(args -> pendingFuture); + dispatchDefaultCalls(initial); + + AccumulatedUnknowns whileRunning = recordDefaultCall(args -> pendingFuture); + + assertThat(whileRunning.callIds()).containsExactlyElementsIn(initial.callIds()); + assertThat(tracker.hasInFlightCalls()).isTrue(); + } + + @Test + public void recordOrGet_concurrentRegistrationSameKey_deduplicatesToSingleCallId() + throws Exception { + List results = new CopyOnWriteArrayList<>(); + + runConcurrently(16, () -> results.add(recordDefaultCall(args -> immediateFuture("done")))); + + assertThat(results).hasSize(16); + long canonicalCallId = getOnlyElement(results.get(0).callIds()); + for (AccumulatedUnknowns result : results) { + assertThat(result.callIds()).containsExactly(canonicalCallId); + } + } + + @Test + public void tryLaunch_whenRecordCannotTransitionToRunning_releasesPermitWithoutDispatch( + @TestParameter boolean alreadyRunning) { + AtomicInteger tasksExecuted = new AtomicInteger(0); + AsyncCallRecord record = + defaultRecord( + args -> { + tasksExecuted.incrementAndGet(); + return immediateFuture("done"); + }); + if (alreadyRunning) { + checkState(record.markRunning()); + } else { + record.cancelInFlight(); + } + + tracker.tryLaunch(record, directExecutor, gate, coordinator, observer); + + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(tasksExecuted.get()).isEqualTo(0); + assertThat(observer.startedCalls()).isEmpty(); + } + + @Test + public void tryLaunch_whenFutureNotifiesListenersTwice_releasesGatePermitOnce() { + AsyncGate twoPermitGate = AsyncGate.create(2); + checkState(twoPermitGate.tryAcquire()); + AsyncCallRecord record = defaultRecord(args -> doubleNotifyingFuture()); + + tracker.tryLaunch( + record, + directExecutor, + twoPermitGate, + newCoordinator(twoPermitGate, directExecutor), + observer); + + assertThat(observer.finishedCalls()).hasSize(1); + assertThat(twoPermitGate.activeCount()).isEqualTo(1); + } + + @Test + public void tryLaunch_observerThrowsOnStart_failsRecordAndReleasesPermit() { + RuntimeException expected = new RuntimeException("observer start failure"); + AtomicReference reportedError = new AtomicReference<>(); + CelAsyncObserver throwingObserver = + new CelAsyncObserver() { + @Override + public void onCallStarted(CelAsyncCall call, ImmutableList args) { + throw expected; + } + + @Override + public void onCallFinished( + CelAsyncCall call, @Nullable Object result, @Nullable Throwable error) { + reportedError.set(error); + } + }; + AsyncCallRecord record = defaultRecord(args -> immediateFuture("done")); + + tracker.tryLaunch(record, directExecutor, gate, coordinator, throwingObserver); + + assertThat(record.state()).isEqualTo(AsyncCallRecord.State.FAILURE); + assertThat(record.error()).hasValue(expected); + assertThat(reportedError.get()).isSameInstanceAs(expected); + assertThat(gate.activeCount()).isEqualTo(0); + } + + @Test + public void tryLaunch_observerThrowsOnFinish_preservesResultAndReleasesPermit() { + CelAsyncObserver throwingObserver = + new CelAsyncObserver() { + @Override + public void onCallStarted(CelAsyncCall call, ImmutableList args) {} + + @Override + public void onCallFinished( + CelAsyncCall call, @Nullable Object result, @Nullable Throwable error) { + throw new RuntimeException("observer finish failure"); + } + }; + AsyncCallRecord record = defaultRecord(args -> immediateFuture("done")); + + tracker.tryLaunch(record, directExecutor, gate, coordinator, throwingObserver); + + assertThat(record.state()).isEqualTo(AsyncCallRecord.State.SUCCESS); + assertThat(record.result()).hasValue("done"); + assertThat(gate.activeCount()).isEqualTo(0); + } + + private enum FailureCase { + CHECKED_OR_RUNTIME(new IllegalArgumentException("computation failed"), "computation failed"), + NULL_MESSAGE(new IllegalStateException((String) null), "IllegalStateException"), + CEL_RUNTIME_EXCEPTION(new CelDivideByZeroException(), "/ by zero"), + CANCELLED(new CancellationException(), "was cancelled"); + + private final Throwable cause; + private final String expectedMessage; + + FailureCase(Throwable cause, String expectedMessage) { + this.cause = cause; + this.expectedMessage = expectedMessage; + } + } + + @Test + public void recordOrGet_whenRecordFailed_throwsExpectedException( + @TestParameter FailureCase failureCase) throws Exception { + ListenableFuture future = immediateFailedFuture(failureCase.cause); + if (failureCase == FailureCase.CANCELLED) { + recordDefaultCall(args -> future); + tracker.cancelInFlight(); + CancellationException e = + assertThrows(CancellationException.class, () -> recordDefaultCall(args -> future)); + assertThat(e).hasMessageThat().contains(failureCase.expectedMessage); + } else if (failureCase == FailureCase.CEL_RUNTIME_EXCEPTION) { + dispatchDefaultCalls(recordDefaultCall(args -> future)); + CelRuntimeException e = + assertThrows(CelRuntimeException.class, () -> recordDefaultCall(args -> future)); + assertThat(e).isSameInstanceAs(failureCase.cause); + } else { + dispatchDefaultCalls(recordDefaultCall(args -> future)); + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> recordDefaultCall(args -> future)); + assertThat(e).hasMessageThat().contains(failureCase.expectedMessage); + assertThat(e).hasCauseThat().isInstanceOf(failureCase.cause.getClass()); + } + } + + @Test + public void recordOrGet_nullCelValueConverter_throwsNullPointerException() { + assertThrows( + NullPointerException.class, + () -> + tracker.recordOrGet( + 10L, + "fn", + "fn_overload", + new Object[] {"x"}, + args -> immediateFuture("done"), + null)); + } + + @Test + public void dispatchPendingCalls_withUnknownCallId_doesNotThrow() { + tracker.dispatchPendingCalls( + ImmutableSet.of(9999L), directExecutor, gate, coordinator, /* observer= */ null); + + assertThat(gate.activeCount()).isEqualTo(0); + } + + private void dispatchCallWithObserver(CelAsyncFunctionOverload overload) { + tracker.dispatchPendingCalls( + recordDefaultCall(overload).callIds(), directExecutor, gate, coordinator, observer); + } + + private void dispatchDefaultCalls(AccumulatedUnknowns unknowns) { + tracker.dispatchPendingCalls( + unknowns.callIds(), directExecutor, gate, coordinator, /* observer= */ null); + } + + private AccumulatedUnknowns recordDefaultCall(CelAsyncFunctionOverload overload) { + return (AccumulatedUnknowns) getDefaultCall(overload); + } + + private Object getDefaultCall(CelAsyncFunctionOverload overload) { + return recordOrGetCall(10L, "fn", "x", overload); + } + + private AccumulatedUnknowns recordCall( + long exprId, String functionName, Object arg, CelAsyncFunctionOverload overload) { + return (AccumulatedUnknowns) recordOrGetCall(exprId, functionName, arg, overload); + } + + private Object recordOrGetCall( + long exprId, String functionName, Object arg, CelAsyncFunctionOverload overload) { + return tracker.recordOrGet( + exprId, + functionName, + functionName + "_overload", + new Object[] {arg}, + overload, + CelValueConverter.getDefaultInstance()); + } + + private static AsyncCallRecord defaultRecord(CelAsyncFunctionOverload overload) { + return AsyncCallRecord.create(100L, 10L, "fn", "fn_overload", new Object[] {"x"}, overload); + } + + private static AsyncCompletionCoordinator newCoordinator(AsyncGate gate, Executor executor) { + return AsyncCompletionCoordinator.create( + CelAsyncEvaluationOptions.defaultOptions(), gate, executor, t -> {}); + } + + private interface ThrowingRunnable { + void run() throws Exception; + } + + private static void runConcurrently(int threads, ThrowingRunnable action) + throws InterruptedException { + ExecutorService pool = Executors.newFixedThreadPool(threads); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch done = new CountDownLatch(threads); + for (int i = 0; i < threads; i++) { + pool.execute( + () -> { + try { + start.await(); + action.run(); + } catch (Exception e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + throw new AssertionError(e); + } finally { + done.countDown(); + } + }); + } + start.countDown(); + assertThat(done.await(5, SECONDS)).isTrue(); + pool.shutdown(); + } + + private static ListenableFuture nonCancellableFuture(ListenableFuture delegate) { + return new SimpleForwardingListenableFuture(delegate) { + @Override + public boolean cancel(boolean mayInterruptIfRunning) { + return false; + } + }; + } + + private static ListenableFuture doubleNotifyingFuture() { + return new SimpleForwardingListenableFuture(immediateFuture("done")) { + @Override + public void addListener(Runnable listener, Executor executor) { + super.addListener(listener, executor); + super.addListener(listener, executor); + } + }; + } + + @ThreadSafe + private static final class RecordingObserver implements CelAsyncObserver { + private final CopyOnWriteArrayList startedCalls = new CopyOnWriteArrayList<>(); + private final CopyOnWriteArrayList finishedCalls = new CopyOnWriteArrayList<>(); + + @Override + public void onCallStarted(CelAsyncCall call, ImmutableList args) { + startedCalls.add(new StartedCall(call, args)); + } + + @Override + public void onCallFinished( + CelAsyncCall call, @Nullable Object result, @Nullable Throwable error) { + finishedCalls.add(new FinishedCall(call, result, error)); + } + + ImmutableList startedCalls() { + return ImmutableList.copyOf(startedCalls); + } + + ImmutableList> startedArgs() { + return startedCalls.stream().map(s -> s.args).collect(toImmutableList()); + } + + ImmutableList finishedCalls() { + return ImmutableList.copyOf(finishedCalls); + } + } + + @Immutable + @SuppressWarnings("Immutable") + private static final class StartedCall { + private final CelAsyncCall call; + private final ImmutableList args; + + private StartedCall(CelAsyncCall call, ImmutableList args) { + this.call = call; + this.args = args; + } + } + + @Immutable + @SuppressWarnings("Immutable") + private static final class FinishedCall { + private final CelAsyncCall call; + private final @Nullable Object result; + private final @Nullable Throwable error; + + private FinishedCall(CelAsyncCall call, @Nullable Object result, @Nullable Throwable error) { + this.call = call; + this.result = result; + this.error = error; + } + } +} diff --git a/runtime/src/test/java/dev/cel/runtime/planner/BUILD.bazel b/runtime/src/test/java/dev/cel/runtime/planner/BUILD.bazel index 304fbe21d..02ca90662 100644 --- a/runtime/src/test/java/dev/cel/runtime/planner/BUILD.bazel +++ b/runtime/src/test/java/dev/cel/runtime/planner/BUILD.bazel @@ -24,6 +24,7 @@ java_library( "//common:options", "//common/ast", "//common/exceptions:divide_by_zero", + "//common/exceptions:runtime_exception", "//common/internal:cel_descriptor_pools", "//common/internal:default_message_factory", "//common/internal:dynamic_proto", @@ -41,6 +42,7 @@ java_library( "//extensions", "//parser:macro", "//runtime", + "//runtime:accumulated_unknowns", "//runtime:descriptor_type_resolver", "//runtime:dispatcher", "//runtime:function_binding", @@ -57,9 +59,11 @@ java_library( "//runtime/standard:type", "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto", "@maven//:com_google_code_findbugs_annotations", + "@maven//:com_google_errorprone_error_prone_annotations", "@maven//:com_google_guava_guava", "@maven//:com_google_testparameterinjector_test_parameter_injector", "@maven//:junit_junit", + "@maven//:org_jspecify_jspecify", ], )