diff --git a/runtime/planner/BUILD.bazel b/runtime/planner/BUILD.bazel index 860d413a0..2781a2e22 100644 --- a/runtime/planner/BUILD.bazel +++ b/runtime/planner/BUILD.bazel @@ -21,3 +21,10 @@ java_library( visibility = ["//:internal"], exports = ["//runtime/src/main/java/dev/cel/runtime/planner:planned_program"], ) + +java_library( + name = "async_gate", + testonly = 1, + visibility = ["//:internal"], + exports = ["//runtime/src/main/java/dev/cel/runtime/planner:async_gate"], +) diff --git a/runtime/src/main/java/dev/cel/runtime/planner/AsyncGate.java b/runtime/src/main/java/dev/cel/runtime/planner/AsyncGate.java new file mode 100644 index 000000000..d284e5286 --- /dev/null +++ b/runtime/src/main/java/dev/cel/runtime/planner/AsyncGate.java @@ -0,0 +1,112 @@ +// 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 com.google.common.annotations.VisibleForTesting; +import com.google.errorprone.annotations.CheckReturnValue; +import javax.annotation.concurrent.ThreadSafe; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.jspecify.annotations.Nullable; + +/** + * Regulates the number of concurrent asynchronous function executions based on maxConcurrency. + * + *

A {@code maxConcurrency} value of {@code 0} or less represents unbounded concurrency (no limit + * on concurrent executions). + */ +@ThreadSafe +final class AsyncGate { + + /** Null when {@code maxConcurrency <= 0}, indicating unbounded concurrency (no throttling). */ + private final @Nullable Semaphore semaphore; + + private final AtomicInteger activeCalls; + private final AtomicBoolean cancelled; + + /** + * Creates an {@link AsyncGate} regulating concurrent asynchronous calls. + * + * @param maxConcurrency the maximum number of concurrent executions allowed. A value of {@code 0} + * or less indicates unbounded concurrency (no concurrency limit). + */ + static AsyncGate create(int maxConcurrency) { + return new AsyncGate(maxConcurrency); + } + + /** + * Attempts to acquire a concurrency slot for an asynchronous call. + * + * @return true if a slot was acquired; false if the gate is cancelled or at maximum concurrency. + */ + @CheckReturnValue + boolean tryAcquire() { + if (semaphore != null && !semaphore.tryAcquire()) { + return false; + } + if (cancelled.get()) { + if (semaphore != null) { + semaphore.release(); + } + return false; + } + activeCalls.incrementAndGet(); + return true; + } + + /** Releases a previously acquired concurrency slot and decrements the active call count. */ + void release() { + while (true) { + int current = activeCalls.get(); + if (current <= 0) { + return; + } + if (activeCalls.compareAndSet(current, current - 1)) { + if (semaphore != null) { + semaphore.release(); + } + return; + } + } + } + + /** Cancels the gate, preventing any future calls from acquiring permits. */ + void cancel() { + cancelled.set(true); + } + + /** Returns true if the gate has been cancelled. */ + boolean isCancelled() { + return cancelled.get(); + } + + /** Returns the current number of active in-flight calls. */ + int activeCount() { + return activeCalls.get(); + } + + /** Returns the number of currently available permits, or -1 if unbounded. */ + @VisibleForTesting + int availablePermits() { + return semaphore != null ? semaphore.availablePermits() : -1; + } + + private AsyncGate(int maxConcurrency) { + this.semaphore = maxConcurrency > 0 ? new Semaphore(maxConcurrency) : null; + this.activeCalls = new AtomicInteger(); + this.cancelled = new AtomicBoolean(false); + } +} 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 d4dbb1659..74d7d8d41 100644 --- a/runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel +++ b/runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel @@ -187,6 +187,19 @@ java_library( ], ) +java_library( + name = "async_gate", + srcs = ["AsyncGate.java"], + tags = [ + ], + deps = [ + "@maven//:com_google_code_findbugs_annotations", + "@maven//:com_google_errorprone_error_prone_annotations", + "@maven//:com_google_guava_guava", + "@maven//:org_jspecify_jspecify", + ], +) + java_library( name = "activation_wrapper", srcs = ["ActivationWrapper.java"], diff --git a/runtime/src/test/java/dev/cel/runtime/planner/AsyncGateTest.java b/runtime/src/test/java/dev/cel/runtime/planner/AsyncGateTest.java new file mode 100644 index 000000000..857297185 --- /dev/null +++ b/runtime/src/test/java/dev/cel/runtime/planner/AsyncGateTest.java @@ -0,0 +1,249 @@ +// 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.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.SECONDS; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class AsyncGateTest { + + @Test + public void tryAcquire_withAvailablePermits_returnsTrueAndIncrementsActiveCount() { + AsyncGate gate = AsyncGate.create(2); + + boolean firstAcquired = gate.tryAcquire(); + boolean secondAcquired = gate.tryAcquire(); + + assertThat(firstAcquired).isTrue(); + assertThat(secondAcquired).isTrue(); + assertThat(gate.activeCount()).isEqualTo(2); + assertThat(gate.availablePermits()).isEqualTo(0); + } + + @Test + public void tryAcquire_atMaxConcurrency_returnsFalseAndDoesNotIncrementActiveCount() { + AsyncGate gate = AsyncGate.create(1); + assertThat(gate.tryAcquire()).isTrue(); + + boolean acquired = gate.tryAcquire(); + + assertThat(acquired).isFalse(); + assertThat(gate.activeCount()).isEqualTo(1); + assertThat(gate.availablePermits()).isEqualTo(0); + } + + @Test + public void tryAcquire_unbounded_alwaysSucceeds() { + AsyncGate gate = AsyncGate.create(0); + + boolean first = gate.tryAcquire(); + boolean second = gate.tryAcquire(); + boolean third = gate.tryAcquire(); + + assertThat(first).isTrue(); + assertThat(second).isTrue(); + assertThat(third).isTrue(); + assertThat(gate.activeCount()).isEqualTo(3); + assertThat(gate.availablePermits()).isEqualTo(-1); + } + + @Test + public void tryAcquire_negativeMaxConcurrency_treatedAsUnbounded() { + AsyncGate gate = AsyncGate.create(-1); + + boolean acquired = gate.tryAcquire(); + + assertThat(acquired).isTrue(); + assertThat(gate.activeCount()).isEqualTo(1); + assertThat(gate.availablePermits()).isEqualTo(-1); + } + + @Test + public void tryAcquire_whenCancelled_returnsFalse() { + AsyncGate gate = AsyncGate.create(2); + gate.cancel(); + + boolean acquired = gate.tryAcquire(); + + assertThat(acquired).isFalse(); + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(gate.availablePermits()).isEqualTo(2); + } + + @Test + public void tryAcquire_unboundedWhenCancelled_returnsFalse() { + AsyncGate gate = AsyncGate.create(0); + gate.cancel(); + + boolean acquired = gate.tryAcquire(); + + assertThat(acquired).isFalse(); + assertThat(gate.activeCount()).isEqualTo(0); + } + + @Test + public void release_decrementsActiveCountAndFreesPermit() { + AsyncGate gate = AsyncGate.create(2); + assertThat(gate.tryAcquire()).isTrue(); + assertThat(gate.tryAcquire()).isTrue(); + + gate.release(); + + assertThat(gate.activeCount()).isEqualTo(1); + assertThat(gate.availablePermits()).isEqualTo(1); + } + + @Test + public void release_unbounded_decrementsActiveCount() { + AsyncGate gate = AsyncGate.create(0); + assertThat(gate.tryAcquire()).isTrue(); + + gate.release(); + + assertThat(gate.activeCount()).isEqualTo(0); + } + + @Test + public void release_allowsSubsequentTryAcquire() { + AsyncGate gate = AsyncGate.create(1); + assertThat(gate.tryAcquire()).isTrue(); + + gate.release(); + + assertThat(gate.tryAcquire()).isTrue(); + assertThat(gate.activeCount()).isEqualTo(1); + } + + @Test + public void release_withoutPriorAcquire_doesNotExceedMaxConcurrency() { + AsyncGate gate = AsyncGate.create(2); + + gate.release(); + + assertThat(gate.availablePermits()).isEqualTo(2); + assertThat(gate.activeCount()).isEqualTo(0); + } + + @Test + public void release_withoutPriorAcquire_doesNotUnderflowActiveCount() { + AsyncGate gate = AsyncGate.create(0); + + gate.release(); + gate.release(); + + assertThat(gate.activeCount()).isEqualTo(0); + } + + @Test + public void release_calledMoreThanAcquires_onlyReleasesAcquiredPermits() { + AsyncGate gate = AsyncGate.create(1); + assertThat(gate.tryAcquire()).isTrue(); + + gate.release(); + gate.release(); + + assertThat(gate.availablePermits()).isEqualTo(1); + assertThat(gate.activeCount()).isEqualTo(0); + } + + @Test + public void cancel_setsIsCancelledToTrue() { + AsyncGate gate = AsyncGate.create(1); + + gate.cancel(); + + assertThat(gate.isCancelled()).isTrue(); + } + + @Test + public void cancel_idempotent() { + AsyncGate gate = AsyncGate.create(1); + + gate.cancel(); + gate.cancel(); + + assertThat(gate.isCancelled()).isTrue(); + } + + @Test + public void create_factoryMethod_returnsConfiguredGate() { + AsyncGate gate = AsyncGate.create(5); + + assertThat(gate.availablePermits()).isEqualTo(5); + assertThat(gate.activeCount()).isEqualTo(0); + assertThat(gate.isCancelled()).isFalse(); + } + + @Test + public void create_withMaxInteger_initializesCorrectly() { + AsyncGate gate = AsyncGate.create(Integer.MAX_VALUE); + + assertThat(gate.availablePermits()).isEqualTo(Integer.MAX_VALUE); + assertThat(gate.tryAcquire()).isTrue(); + assertThat(gate.activeCount()).isEqualTo(1); + } + + @Test + public void concurrentTryAcquireAndRelease_neverExceedsMaxConcurrency() + throws InterruptedException { + int maxConcurrency = 4; + int taskCount = 32; + AsyncGate gate = AsyncGate.create(maxConcurrency); + AtomicInteger peakConcurrency = new AtomicInteger(); + ExecutorService executor = Executors.newFixedThreadPool(8); + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(taskCount); + + try { + for (int i = 0; i < taskCount; i++) { + executor.execute( + () -> { + try { + startLatch.await(); + while (!gate.tryAcquire()) { + Thread.sleep(1); + } + int current = gate.activeCount(); + peakConcurrency.accumulateAndGet(current, Math::max); + Thread.sleep(2); + gate.release(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + doneLatch.countDown(); + } + }); + } + + startLatch.countDown(); + boolean completed = doneLatch.await(5, SECONDS); + + assertThat(completed).isTrue(); + assertThat(peakConcurrency.get()).isAtMost(maxConcurrency); + assertThat(gate.activeCount()).isEqualTo(0); + } finally { + executor.shutdown(); + } + } +} 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 9116818dc..5ff4b4d81 100644 --- a/runtime/src/test/java/dev/cel/runtime/planner/BUILD.bazel +++ b/runtime/src/test/java/dev/cel/runtime/planner/BUILD.bazel @@ -49,6 +49,7 @@ java_library( "//runtime:runtime_helpers", "//runtime:standard_functions", "//runtime:unknown_attributes", + "//runtime/planner:async_gate", "//runtime/planner:program_planner", "//runtime/standard:type", "@cel_spec//proto/cel/expr/conformance/proto3:test_all_types_java_proto",