Skip to content

Commit 11fb683

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Add async concurrency gate
PiperOrigin-RevId: 974934413
1 parent 6124cb0 commit 11fb683

5 files changed

Lines changed: 383 additions & 0 deletions

File tree

‎runtime/planner/BUILD.bazel‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,10 @@ java_library(
2121
visibility = ["//:internal"],
2222
exports = ["//runtime/src/main/java/dev/cel/runtime/planner:planned_program"],
2323
)
24+
25+
java_library(
26+
name = "async_gate",
27+
testonly = 1,
28+
visibility = ["//:internal"],
29+
exports = ["//runtime/src/main/java/dev/cel/runtime/planner:async_gate"],
30+
)

‎runtime/src/main/java/dev/cel/runtime/planner/AsyncGate.java‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package dev.cel.runtime.planner;
16+
17+
import com.google.errorprone.annotations.CheckReturnValue;
18+
import javax.annotation.concurrent.ThreadSafe;
19+
import java.util.concurrent.Semaphore;
20+
import java.util.concurrent.atomic.AtomicBoolean;
21+
import java.util.concurrent.atomic.AtomicInteger;
22+
import org.jspecify.annotations.Nullable;
23+
24+
/**
25+
* Regulates the number of concurrent asynchronous function executions based on maxConcurrency.
26+
*
27+
* <p>A {@code maxConcurrency} value of {@code 0} or less represents unbounded concurrency (no limit
28+
* on concurrent executions).
29+
*/
30+
@ThreadSafe
31+
final class AsyncGate {
32+
33+
/** Null when {@code maxConcurrency <= 0}, indicating unbounded concurrency (no throttling). */
34+
private final @Nullable Semaphore semaphore;
35+
36+
private final AtomicInteger activeCalls;
37+
private final AtomicBoolean cancelled;
38+
39+
/**
40+
* Creates an {@link AsyncGate} regulating concurrent asynchronous calls.
41+
*
42+
* @param maxConcurrency the maximum number of concurrent executions allowed. A value of {@code 0}
43+
* or less indicates unbounded concurrency (no concurrency limit).
44+
*/
45+
static AsyncGate create(int maxConcurrency) {
46+
return new AsyncGate(maxConcurrency);
47+
}
48+
49+
/**
50+
* Attempts to acquire a concurrency slot for an asynchronous call.
51+
*
52+
* <p>Cancellation check is best-effort admission control. A thread may observe {@code
53+
* cancelled.get() == false} and acquire a permit immediately before a concurrent {@link
54+
* #cancel()} runs. Any call launched in this race window will complete safely into {@code
55+
* AsyncCompletionCoordinator.callCompleted()}, where permits are released and results discarded.
56+
*
57+
* @return true if a slot was acquired; false if the gate is cancelled or at maximum concurrency.
58+
*/
59+
@CheckReturnValue
60+
boolean tryAcquire() {
61+
if (semaphore != null && !semaphore.tryAcquire()) {
62+
return false;
63+
}
64+
// Best-effort check: if cancelled concurrently after this point, the launched task
65+
// will complete as a no-op in the completion coordinator.
66+
if (cancelled.get()) {
67+
if (semaphore != null) {
68+
semaphore.release();
69+
}
70+
return false;
71+
}
72+
activeCalls.incrementAndGet();
73+
return true;
74+
}
75+
76+
/** Releases a previously acquired concurrency slot and decrements the active call count. */
77+
void release() {
78+
while (true) {
79+
int current = activeCalls.get();
80+
if (current <= 0) {
81+
return;
82+
}
83+
if (activeCalls.compareAndSet(current, current - 1)) {
84+
if (semaphore != null) {
85+
semaphore.release();
86+
}
87+
return;
88+
}
89+
}
90+
}
91+
92+
/**
93+
* Cancels the gate, preventing future calls from acquiring permits.
94+
*
95+
* <p>Cancellation is best-effort admission control; tasks that acquired permits immediately prior
96+
* to cancellation will execute and complete as no-ops in the completion coordinator.
97+
*/
98+
void cancel() {
99+
cancelled.set(true);
100+
}
101+
102+
/** Returns true if the gate has been cancelled. */
103+
boolean isCancelled() {
104+
return cancelled.get();
105+
}
106+
107+
/** Returns the current number of active in-flight calls. */
108+
int activeCount() {
109+
return activeCalls.get();
110+
}
111+
112+
private AsyncGate(int maxConcurrency) {
113+
this.semaphore = maxConcurrency > 0 ? new Semaphore(maxConcurrency) : null;
114+
this.activeCalls = new AtomicInteger();
115+
this.cancelled = new AtomicBoolean(false);
116+
}
117+
}

‎runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,19 @@ java_library(
187187
],
188188
)
189189

190+
java_library(
191+
name = "async_gate",
192+
srcs = ["AsyncGate.java"],
193+
tags = [
194+
],
195+
deps = [
196+
"@maven//:com_google_code_findbugs_annotations",
197+
"@maven//:com_google_errorprone_error_prone_annotations",
198+
"@maven//:com_google_guava_guava",
199+
"@maven//:org_jspecify_jspecify",
200+
],
201+
)
202+
190203
java_library(
191204
name = "activation_wrapper",
192205
srcs = ["ActivationWrapper.java"],

‎runtime/src/test/java/dev/cel/runtime/planner/AsyncGateTest.java‎

Lines changed: 245 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package dev.cel.runtime.planner;
16+
17+
import static com.google.common.truth.Truth.assertThat;
18+
import static java.util.concurrent.TimeUnit.SECONDS;
19+
20+
import java.util.concurrent.CountDownLatch;
21+
import java.util.concurrent.ExecutorService;
22+
import java.util.concurrent.Executors;
23+
import java.util.concurrent.atomic.AtomicInteger;
24+
import org.junit.Test;
25+
import org.junit.runner.RunWith;
26+
import org.junit.runners.JUnit4;
27+
28+
@RunWith(JUnit4.class)
29+
public final class AsyncGateTest {
30+
31+
@Test
32+
public void tryAcquire_withAvailablePermits_returnsTrueAndIncrementsActiveCount() {
33+
AsyncGate gate = AsyncGate.create(2);
34+
35+
boolean firstAcquired = gate.tryAcquire();
36+
boolean secondAcquired = gate.tryAcquire();
37+
38+
assertThat(firstAcquired).isTrue();
39+
assertThat(secondAcquired).isTrue();
40+
assertThat(gate.activeCount()).isEqualTo(2);
41+
}
42+
43+
@Test
44+
public void tryAcquire_atMaxConcurrency_returnsFalseAndDoesNotIncrementActiveCount() {
45+
AsyncGate gate = AsyncGate.create(1);
46+
assertThat(gate.tryAcquire()).isTrue();
47+
48+
boolean acquired = gate.tryAcquire();
49+
50+
assertThat(acquired).isFalse();
51+
assertThat(gate.activeCount()).isEqualTo(1);
52+
}
53+
54+
@Test
55+
public void tryAcquire_unbounded_alwaysSucceeds() {
56+
AsyncGate gate = AsyncGate.create(0);
57+
58+
boolean first = gate.tryAcquire();
59+
boolean second = gate.tryAcquire();
60+
boolean third = gate.tryAcquire();
61+
62+
assertThat(first).isTrue();
63+
assertThat(second).isTrue();
64+
assertThat(third).isTrue();
65+
assertThat(gate.activeCount()).isEqualTo(3);
66+
}
67+
68+
@Test
69+
public void tryAcquire_negativeMaxConcurrency_treatedAsUnbounded() {
70+
AsyncGate gate = AsyncGate.create(-1);
71+
72+
boolean acquired = gate.tryAcquire();
73+
74+
assertThat(acquired).isTrue();
75+
assertThat(gate.activeCount()).isEqualTo(1);
76+
}
77+
78+
@Test
79+
public void tryAcquire_whenCancelled_returnsFalse() {
80+
AsyncGate gate = AsyncGate.create(2);
81+
gate.cancel();
82+
83+
boolean acquired = gate.tryAcquire();
84+
85+
assertThat(acquired).isFalse();
86+
assertThat(gate.activeCount()).isEqualTo(0);
87+
}
88+
89+
@Test
90+
public void tryAcquire_unboundedWhenCancelled_returnsFalse() {
91+
AsyncGate gate = AsyncGate.create(0);
92+
gate.cancel();
93+
94+
boolean acquired = gate.tryAcquire();
95+
96+
assertThat(acquired).isFalse();
97+
assertThat(gate.activeCount()).isEqualTo(0);
98+
}
99+
100+
@Test
101+
public void release_decrementsActiveCountAndFreesPermit() {
102+
AsyncGate gate = AsyncGate.create(2);
103+
assertThat(gate.tryAcquire()).isTrue();
104+
assertThat(gate.tryAcquire()).isTrue();
105+
106+
gate.release();
107+
108+
assertThat(gate.activeCount()).isEqualTo(1);
109+
assertThat(gate.tryAcquire()).isTrue();
110+
}
111+
112+
@Test
113+
public void release_unbounded_decrementsActiveCount() {
114+
AsyncGate gate = AsyncGate.create(0);
115+
assertThat(gate.tryAcquire()).isTrue();
116+
117+
gate.release();
118+
119+
assertThat(gate.activeCount()).isEqualTo(0);
120+
}
121+
122+
@Test
123+
public void release_allowsSubsequentTryAcquire() {
124+
AsyncGate gate = AsyncGate.create(1);
125+
assertThat(gate.tryAcquire()).isTrue();
126+
127+
gate.release();
128+
129+
assertThat(gate.tryAcquire()).isTrue();
130+
assertThat(gate.activeCount()).isEqualTo(1);
131+
}
132+
133+
@Test
134+
public void release_withoutPriorAcquire_doesNotExceedMaxConcurrency() {
135+
AsyncGate gate = AsyncGate.create(2);
136+
137+
gate.release();
138+
139+
assertThat(gate.activeCount()).isEqualTo(0);
140+
assertThat(gate.tryAcquire()).isTrue();
141+
assertThat(gate.tryAcquire()).isTrue();
142+
assertThat(gate.tryAcquire()).isFalse();
143+
}
144+
145+
@Test
146+
public void release_withoutPriorAcquire_doesNotUnderflowActiveCount() {
147+
AsyncGate gate = AsyncGate.create(0);
148+
149+
gate.release();
150+
gate.release();
151+
152+
assertThat(gate.activeCount()).isEqualTo(0);
153+
}
154+
155+
@Test
156+
public void release_calledMoreThanAcquires_onlyReleasesAcquiredPermits() {
157+
AsyncGate gate = AsyncGate.create(1);
158+
assertThat(gate.tryAcquire()).isTrue();
159+
160+
gate.release();
161+
gate.release();
162+
163+
assertThat(gate.activeCount()).isEqualTo(0);
164+
assertThat(gate.tryAcquire()).isTrue();
165+
assertThat(gate.tryAcquire()).isFalse();
166+
}
167+
168+
@Test
169+
public void cancel_setsIsCancelledToTrue() {
170+
AsyncGate gate = AsyncGate.create(1);
171+
172+
gate.cancel();
173+
174+
assertThat(gate.isCancelled()).isTrue();
175+
}
176+
177+
@Test
178+
public void cancel_idempotent() {
179+
AsyncGate gate = AsyncGate.create(1);
180+
181+
gate.cancel();
182+
gate.cancel();
183+
184+
assertThat(gate.isCancelled()).isTrue();
185+
}
186+
187+
@Test
188+
public void create_factoryMethod_returnsConfiguredGate() {
189+
AsyncGate gate = AsyncGate.create(5);
190+
191+
assertThat(gate.activeCount()).isEqualTo(0);
192+
assertThat(gate.isCancelled()).isFalse();
193+
}
194+
195+
@Test
196+
public void create_withMaxInteger_initializesCorrectly() {
197+
AsyncGate gate = AsyncGate.create(Integer.MAX_VALUE);
198+
199+
assertThat(gate.tryAcquire()).isTrue();
200+
assertThat(gate.activeCount()).isEqualTo(1);
201+
}
202+
203+
@Test
204+
public void concurrentTryAcquireAndRelease_neverExceedsMaxConcurrency()
205+
throws InterruptedException {
206+
int maxConcurrency = 4;
207+
int taskCount = 32;
208+
AsyncGate gate = AsyncGate.create(maxConcurrency);
209+
AtomicInteger peakConcurrency = new AtomicInteger();
210+
ExecutorService executor = Executors.newFixedThreadPool(8);
211+
CountDownLatch startLatch = new CountDownLatch(1);
212+
CountDownLatch doneLatch = new CountDownLatch(taskCount);
213+
214+
try {
215+
for (int i = 0; i < taskCount; i++) {
216+
executor.execute(
217+
() -> {
218+
try {
219+
startLatch.await();
220+
while (!gate.tryAcquire()) {
221+
Thread.sleep(1);
222+
}
223+
int current = gate.activeCount();
224+
peakConcurrency.accumulateAndGet(current, Math::max);
225+
Thread.sleep(2);
226+
gate.release();
227+
} catch (InterruptedException e) {
228+
Thread.currentThread().interrupt();
229+
} finally {
230+
doneLatch.countDown();
231+
}
232+
});
233+
}
234+
235+
startLatch.countDown();
236+
boolean completed = doneLatch.await(5, SECONDS);
237+
238+
assertThat(completed).isTrue();
239+
assertThat(peakConcurrency.get()).isAtMost(maxConcurrency);
240+
assertThat(gate.activeCount()).isEqualTo(0);
241+
} finally {
242+
executor.shutdown();
243+
}
244+
}
245+
}

0 commit comments

Comments
 (0)