Skip to content

Commit 40639ff

Browse files
l46kokcopybara-github
authored andcommitted
Add async concurrency gate
PiperOrigin-RevId: 974934413
1 parent abff7d1 commit 40639ff

5 files changed

Lines changed: 1239 additions & 0 deletions

File tree

‎runtime/planner/BUILD.bazel‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,15 @@ 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+
visibility = ["//:internal"],
28+
exports = ["//runtime/src/main/java/dev/cel/runtime/planner:async_gate"],
29+
)
30+
31+
cel_android_library(
32+
name = "async_gate_android",
33+
visibility = ["//:internal"],
34+
exports = ["//runtime/src/main/java/dev/cel/runtime/planner:async_gate_android"],
35+
)
Lines changed: 305 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,305 @@
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.base.Preconditions.checkNotNull;
18+
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
19+
20+
import com.google.common.annotations.VisibleForTesting;
21+
import com.google.common.util.concurrent.FutureCallback;
22+
import com.google.common.util.concurrent.Futures;
23+
import com.google.common.util.concurrent.ListenableFuture;
24+
import com.google.errorprone.annotations.CanIgnoreReturnValue;
25+
import javax.annotation.concurrent.ThreadSafe;
26+
import java.util.Queue;
27+
import java.util.concurrent.CancellationException;
28+
import java.util.concurrent.ConcurrentLinkedQueue;
29+
import java.util.concurrent.Executor;
30+
import java.util.concurrent.Semaphore;
31+
import java.util.concurrent.atomic.AtomicBoolean;
32+
import org.jspecify.annotations.Nullable;
33+
34+
/** Regulates the number of concurrent asynchronous function executions based on maxConcurrency. */
35+
@ThreadSafe
36+
final class AsyncGate {
37+
38+
private final @Nullable Semaphore semaphore;
39+
40+
// CEL-Internal-4
41+
private final Queue<QueuedTask> pendingTasks;
42+
43+
private final AtomicBoolean cancelled;
44+
private final ThreadLocal<Boolean> isDraining;
45+
46+
/** Represents an asynchronous task subject to concurrency gating. */
47+
interface GatedTask {
48+
/**
49+
* Starts the asynchronous task.
50+
*
51+
* <p>The concurrency gate holds a permit until the returned {@link ListenableFuture} completes.
52+
*/
53+
ListenableFuture<?> start();
54+
55+
/**
56+
* Invoked instead of {@link #start} when the task cannot be admitted (e.g. the gate is
57+
* cancelled or the executor rejects execution).
58+
*/
59+
default void reject(Throwable cause) {}
60+
}
61+
62+
static final class QueuedTask {
63+
private final GatedTask task;
64+
private final Executor executor;
65+
66+
QueuedTask(GatedTask task, Executor executor) {
67+
this.task = checkNotNull(task, "task must not be null");
68+
this.executor = checkNotNull(executor, "executor must not be null");
69+
}
70+
}
71+
72+
/** Cancels the gate, rejecting all pending queued tasks and preventing new dispatches. */
73+
void cancel() {
74+
cancelled.set(true);
75+
clearPendingTasks(new CancellationException("AsyncGate is cancelled"));
76+
}
77+
78+
/**
79+
* Dispatches a gated task for execution.
80+
*
81+
* <p>If concurrency permits allow, the task is submitted to the executor immediately. Otherwise,
82+
* it is enqueued into pending tasks and will be drained when an active task completes.
83+
*
84+
* @return true if the task was admitted or queued for execution, false if rejected due to
85+
* cancellation.
86+
*/
87+
@CanIgnoreReturnValue
88+
boolean dispatch(Executor executor, GatedTask task) {
89+
checkNotNull(executor, "executor must not be null");
90+
checkNotNull(task, "task must not be null");
91+
if (cancelled.get()) {
92+
task.reject(new CancellationException("AsyncGate is cancelled"));
93+
return false;
94+
}
95+
96+
if (semaphore == null) {
97+
AtomicBoolean started = new AtomicBoolean(false);
98+
try {
99+
executor.execute(
100+
() -> {
101+
started.set(true);
102+
if (cancelled.get()) {
103+
task.reject(new CancellationException("AsyncGate is cancelled"));
104+
return;
105+
}
106+
executeTaskWithoutPermit(task);
107+
});
108+
return true;
109+
} catch (Throwable t) {
110+
if (!started.get()) {
111+
task.reject(t);
112+
return false;
113+
}
114+
throw t;
115+
}
116+
}
117+
118+
if (pendingTasks.isEmpty() && semaphore.tryAcquire()) {
119+
AtomicBoolean started = new AtomicBoolean(false);
120+
try {
121+
executor.execute(
122+
() -> {
123+
started.set(true);
124+
if (cancelled.get()) {
125+
semaphore.release();
126+
task.reject(new CancellationException("AsyncGate is cancelled"));
127+
drainPending();
128+
return;
129+
}
130+
runTaskWithPermit(task);
131+
});
132+
return true;
133+
} catch (Throwable t) {
134+
if (!started.get()) {
135+
semaphore.release();
136+
drainPending();
137+
task.reject(t);
138+
return false;
139+
}
140+
throw t;
141+
}
142+
}
143+
144+
pendingTasks.add(new QueuedTask(task, executor));
145+
drainPending();
146+
return true;
147+
}
148+
149+
private void runTaskWithPermit(GatedTask task) {
150+
ListenableFuture<?> future;
151+
try {
152+
future = checkNotNull(task.start(), "GatedTask.start() returned null");
153+
} catch (Throwable t) {
154+
semaphore.release();
155+
try {
156+
task.reject(t);
157+
} catch (Throwable e) {
158+
// Suppress secondary exception from reject callback to preserve primary cause t.
159+
t.addSuppressed(e);
160+
}
161+
drainPending();
162+
throw t;
163+
}
164+
165+
Futures.addCallback(
166+
future,
167+
new FutureCallback<Object>() {
168+
@Override
169+
public void onSuccess(@Nullable Object result) {
170+
release();
171+
}
172+
173+
@Override
174+
public void onFailure(Throwable t) {
175+
release();
176+
}
177+
178+
private void release() {
179+
semaphore.release();
180+
drainPending();
181+
}
182+
},
183+
directExecutor());
184+
}
185+
186+
@SuppressWarnings("FutureReturnValueIgnored")
187+
private void executeTaskWithoutPermit(GatedTask task) {
188+
try {
189+
checkNotNull(task.start(), "GatedTask.start() returned null");
190+
} catch (Throwable t) {
191+
try {
192+
task.reject(t);
193+
} catch (Throwable e) {
194+
// Suppress secondary exception from reject callback to preserve primary cause t.
195+
t.addSuppressed(e);
196+
}
197+
throw t;
198+
}
199+
}
200+
201+
@VisibleForTesting
202+
void drainPending() {
203+
if (isDraining.get()) {
204+
return;
205+
}
206+
isDraining.set(true);
207+
try {
208+
if (cancelled.get()) {
209+
clearPendingTasks(new CancellationException("AsyncGate is cancelled"));
210+
return;
211+
}
212+
if (semaphore == null) {
213+
return;
214+
}
215+
216+
while (!pendingTasks.isEmpty()) {
217+
if (!semaphore.tryAcquire()) {
218+
break;
219+
}
220+
QueuedTask queuedTask = pendingTasks.poll();
221+
if (queuedTask == null) {
222+
semaphore.release();
223+
break;
224+
}
225+
226+
try {
227+
queuedTask.executor.execute(
228+
() -> {
229+
if (cancelled.get()) {
230+
semaphore.release();
231+
queuedTask.task.reject(new CancellationException("AsyncGate is cancelled"));
232+
drainPending();
233+
return;
234+
}
235+
try {
236+
runTaskWithPermit(queuedTask.task);
237+
} catch (Throwable ignored) {
238+
// runTaskWithPermit already releases permit, calls reject, and drains.
239+
}
240+
});
241+
} catch (Throwable e) {
242+
semaphore.release();
243+
try {
244+
queuedTask.task.reject(e);
245+
} catch (Throwable t) {
246+
// Suppress secondary exception from reject callback to preserve primary cause e.
247+
e.addSuppressed(t);
248+
}
249+
}
250+
}
251+
} finally {
252+
isDraining.set(false);
253+
}
254+
}
255+
256+
private void clearPendingTasks(Throwable cause) {
257+
QueuedTask queuedTask;
258+
while ((queuedTask = pendingTasks.poll()) != null) {
259+
try {
260+
queuedTask.task.reject(cause);
261+
} catch (Throwable t) {
262+
// Suppress exception from reject callback to preserve primary cancellation cause.
263+
cause.addSuppressed(t);
264+
}
265+
}
266+
}
267+
268+
@VisibleForTesting
269+
boolean isCancelled() {
270+
return cancelled.get();
271+
}
272+
273+
@VisibleForTesting
274+
int availablePermits() {
275+
return semaphore != null ? semaphore.availablePermits() : -1;
276+
}
277+
278+
@VisibleForTesting
279+
int queuedTaskCount() {
280+
return pendingTasks.size();
281+
}
282+
283+
static AsyncGate create(int maxConcurrency) {
284+
return new AsyncGate(maxConcurrency);
285+
}
286+
287+
@VisibleForTesting
288+
AsyncGate(int maxConcurrency) {
289+
this(maxConcurrency > 0 ? new Semaphore(maxConcurrency) : null, new ConcurrentLinkedQueue<>());
290+
}
291+
292+
@VisibleForTesting
293+
AsyncGate(@Nullable Semaphore semaphore, Queue<QueuedTask> pendingTasks) {
294+
this.semaphore = semaphore;
295+
this.pendingTasks = checkNotNull(pendingTasks, "pendingTasks must not be null");
296+
this.cancelled = new AtomicBoolean(false);
297+
this.isDraining =
298+
new ThreadLocal<Boolean>() {
299+
@Override
300+
protected Boolean initialValue() {
301+
return Boolean.FALSE;
302+
}
303+
};
304+
}
305+
}

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

Lines changed: 27 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"],
@@ -717,6 +730,20 @@ cel_android_library(
717730
],
718731
)
719732

733+
cel_android_library(
734+
name = "async_gate_android",
735+
srcs = ["AsyncGate.java"],
736+
tags = [
737+
],
738+
deps = [
739+
"@maven//:com_google_code_findbugs_annotations",
740+
"@maven//:com_google_errorprone_error_prone_annotations",
741+
"@maven//:com_google_guava_guava",
742+
"@maven//:org_jspecify_jspecify",
743+
"@maven_android//:com_google_guava_guava",
744+
],
745+
)
746+
720747
cel_android_library(
721748
name = "activation_wrapper_android",
722749
srcs = ["ActivationWrapper.java"],

0 commit comments

Comments
 (0)