Skip to content

Add async concurrency gate #1202

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Sep 14, 2026
Merged

Add async concurrency gate #1202

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions runtime/planner/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
)
117 changes: 117 additions & 0 deletions runtime/src/main/java/dev/cel/runtime/planner/AsyncGate.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
// 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.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.
*
* <p>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.
*
* <p>Cancellation check is best-effort admission control. A thread may observe {@code
* cancelled.get() == false} and acquire a permit immediately before a concurrent {@link
* #cancel()} runs. Any call launched in this race window will complete safely into {@code
* AsyncCompletionCoordinator.callCompleted()}, where permits are released and results discarded.
*
* @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;
}
// Best-effort check: if cancelled concurrently after this point, the launched task
// will complete as a no-op in the completion coordinator.
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 future calls from acquiring permits.
*
* <p>Cancellation is best-effort admission control; tasks that acquired permits immediately prior
* to cancellation will execute and complete as no-ops in the completion coordinator.
*/
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();
}

private AsyncGate(int maxConcurrency) {
this.semaphore = maxConcurrency > 0 ? new Semaphore(maxConcurrency) : null;
this.activeCalls = new AtomicInteger();
this.cancelled = new AtomicBoolean(false);
}
}
13 changes: 13 additions & 0 deletions runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
245 changes: 245 additions & 0 deletions runtime/src/test/java/dev/cel/runtime/planner/AsyncGateTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,245 @@
// 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);
}

@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);
}

@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);
}

@Test
public void tryAcquire_negativeMaxConcurrency_treatedAsUnbounded() {
AsyncGate gate = AsyncGate.create(-1);

boolean acquired = gate.tryAcquire();

assertThat(acquired).isTrue();
assertThat(gate.activeCount()).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);
}

@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.tryAcquire()).isTrue();
}

@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.activeCount()).isEqualTo(0);
assertThat(gate.tryAcquire()).isTrue();
assertThat(gate.tryAcquire()).isTrue();
assertThat(gate.tryAcquire()).isFalse();
}

@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.activeCount()).isEqualTo(0);
assertThat(gate.tryAcquire()).isTrue();
assertThat(gate.tryAcquire()).isFalse();
}

@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.activeCount()).isEqualTo(0);
assertThat(gate.isCancelled()).isFalse();
}

@Test
public void create_withMaxInteger_initializesCorrectly() {
AsyncGate gate = AsyncGate.create(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();
}
}
}
Loading
Loading