Skip to content

Commit 1ffc02c

Browse files
serhiy-bzhezytskyyserhiy-bzhezytskyy
authored
fix: preserve push notification order per task (#1156)
Chain pushes onto a per-task tail future instead of submitting them independently to the ForkJoinPool, so each one starts only after the previous one for the same task finishes. Different tasks stay independent. Register whenComplete cleanup after compute() returns: with a synchronous executor the future completes inside compute(), so a callback registered there fires before the future is stored -- a no-op remove that leaks one entry per task. Fixes #775
1 parent fc5eb07 commit 1ffc02c

2 files changed

Lines changed: 229 additions & 8 deletions

File tree

‎server-common/src/main/java/org/a2aproject/sdk/server/events/MainEventBusProcessor.java‎

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package org.a2aproject.sdk.server.events;
22

33
import java.util.concurrent.CompletableFuture;
4+
import java.util.concurrent.ConcurrentHashMap;
5+
import java.util.concurrent.Executor;
6+
import java.util.concurrent.ForkJoinPool;
47
import java.util.concurrent.atomic.AtomicReference;
58

69
import jakarta.annotation.Nullable;
@@ -75,6 +78,17 @@ private record UpdateResult(boolean isFinal, @Nullable Task taskSnapshot) {}
7578
*/
7679
private volatile @Nullable java.util.concurrent.Executor pushNotificationExecutor = null;
7780

81+
/**
82+
* Per-task tail of the push-notification chain. {@link #sendPushNotification} is only ever
83+
* called from this class's own single-threaded processing loop, so pushes for a given task
84+
* are chained onto this exact sequence rather than submitted independently -- each one only
85+
* starts once the previous one for the same task has finished, which is what guarantees they
86+
* reach the client in the order they were produced even though delivery itself happens off
87+
* this thread. A completed chain removes its own entry (compare-and-remove against the future
88+
* it added) so the map does not grow forever.
89+
*/
90+
private final ConcurrentHashMap<String, CompletableFuture<Void>> pushNotificationChains = new ConcurrentHashMap<>();
91+
7892
private MainEventBus eventBus;
7993

8094
private TaskStore taskStore;
@@ -135,6 +149,13 @@ public void setPushNotificationExecutor(java.util.concurrent.Executor executor)
135149
this.pushNotificationExecutor = executor;
136150
}
137151

152+
/**
153+
* Test-only visibility into the per-task chain map, to assert entries don't leak.
154+
*/
155+
int pushNotificationChainCount() {
156+
return pushNotificationChains.size();
157+
}
158+
138159
@SuppressWarnings("NullAway.Init")
139160
@PostConstruct
140161
void start() {
@@ -371,8 +392,14 @@ private UpdateResult updateTaskStore(String taskId, Event event, boolean isRepli
371392
* The event will be automatically wrapped in StreamResponse format by JsonUtil.
372393
* </p>
373394
* <p>
374-
* <b>NOTE:</b> Tests can inject a synchronous executor via setPushNotificationExecutor()
375-
* to ensure deterministic ordering of push notifications in the test environment.
395+
* <b>Ordering:</b> chained onto {@link #pushNotificationChains this task's own tail future}
396+
* rather than submitted independently, so pushes for the same task are always delivered in
397+
* the order they were produced, regardless of executor scheduling. Different tasks remain
398+
* fully independent -- one task's slow delivery never delays another's.
399+
* </p>
400+
* <p>
401+
* Tests can inject a synchronous executor via setPushNotificationExecutor() to make delivery
402+
* itself synchronous too, but ordering does not depend on that.
376403
* </p>
377404
*
378405
* @param taskId the task ID
@@ -393,12 +420,17 @@ private void sendPushNotification(String taskId, StreamingEventKind event, @Null
393420
}
394421
};
395422

396-
// Use custom executor if set (for tests), otherwise use default ForkJoinPool (async)
397-
if (pushNotificationExecutor != null) {
398-
pushNotificationExecutor.execute(pushTask);
399-
} else {
400-
CompletableFuture.runAsync(pushTask);
401-
}
423+
Executor executor = pushNotificationExecutor != null ? pushNotificationExecutor : ForkJoinPool.commonPool();
424+
425+
// Registered after compute() returns: with a synchronous executor, next
426+
// completes before compute()'s lambda returns, so attaching whenComplete
427+
// in there fires the removal before next is even stored -- a no-op that
428+
// leaks one entry per task.
429+
CompletableFuture<Void> next = pushNotificationChains.compute(taskId, (id, previous) -> {
430+
CompletableFuture<Void> previousOrDone = previous != null ? previous : CompletableFuture.completedFuture(null);
431+
return previousOrDone.thenRunAsync(pushTask, executor);
432+
});
433+
next.whenComplete((v, t) -> pushNotificationChains.remove(taskId, next));
402434
}
403435

404436
/**

‎server-common/src/test/java/org/a2aproject/sdk/server/events/MainEventBusProcessorPushNotificationOrderTest.java‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
package org.a2aproject.sdk.server.events;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertTrue;
5+
6+
import java.util.List;
7+
import java.util.concurrent.CopyOnWriteArrayList;
8+
import java.util.concurrent.CountDownLatch;
9+
import java.util.concurrent.TimeUnit;
10+
import java.util.concurrent.atomic.AtomicInteger;
11+
12+
import org.a2aproject.sdk.server.tasks.InMemoryTaskStore;
13+
import org.a2aproject.sdk.server.tasks.PushNotificationSender;
14+
import org.a2aproject.sdk.spec.Task;
15+
import org.a2aproject.sdk.spec.TaskState;
16+
import org.a2aproject.sdk.spec.TaskStatus;
17+
import org.a2aproject.sdk.spec.TaskStatusUpdateEvent;
18+
import org.junit.jupiter.api.AfterEach;
19+
import org.junit.jupiter.api.BeforeEach;
20+
import org.junit.jupiter.api.Test;
21+
22+
/**
23+
* Regression coverage for https://github.com/a2aproject/a2a-java/issues/775: push notifications
24+
* for the same task used to be submitted independently to the shared {@code ForkJoinPool},
25+
* so the client could observe them out of order relative to how the agent produced them.
26+
*/
27+
public class MainEventBusProcessorPushNotificationOrderTest {
28+
29+
private static final String CONTEXT_ID = "test-context";
30+
31+
private MainEventBus mainEventBus;
32+
private MainEventBusProcessor mainEventBusProcessor;
33+
private InMemoryTaskStore taskStore;
34+
private InMemoryQueueManager queueManager;
35+
36+
@BeforeEach
37+
public void setUp() {
38+
taskStore = new InMemoryTaskStore();
39+
mainEventBus = new MainEventBus();
40+
queueManager = new InMemoryQueueManager(null, mainEventBus);
41+
}
42+
43+
@AfterEach
44+
public void tearDown() {
45+
if (mainEventBusProcessor != null) {
46+
EventQueueUtil.stop(mainEventBusProcessor);
47+
}
48+
}
49+
50+
@Test
51+
public void testEventsForSameTaskDeliveredInOrder() throws InterruptedException {
52+
String taskId = "order-task";
53+
int eventCount = 20;
54+
List<Integer> deliveryOrder = new CopyOnWriteArrayList<>();
55+
CountDownLatch latch = new CountDownLatch(eventCount);
56+
57+
// The Nth push sleeps longer the SMALLER N is, so an unordered executor would very
58+
// likely deliver a later, faster push before an earlier, slower one -- exactly the
59+
// interleaving #775 reported. A correctly-ordered chain delivers 0, 1, 2, ... regardless.
60+
PushNotificationSender sender = (event, snapshot) -> {
61+
TaskStatusUpdateEvent statusEvent = (TaskStatusUpdateEvent) event;
62+
int index = extractSequence(statusEvent);
63+
try {
64+
Thread.sleep(Math.max(0, (eventCount - index) * 2L));
65+
} catch (InterruptedException e) {
66+
Thread.currentThread().interrupt();
67+
}
68+
deliveryOrder.add(index);
69+
latch.countDown();
70+
};
71+
72+
mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, sender, queueManager);
73+
EventQueueUtil.start(mainEventBusProcessor);
74+
EventQueue eventQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus)
75+
.taskId(taskId)
76+
.mainEventBus(mainEventBus)
77+
.build().tap();
78+
79+
eventQueue.enqueueEvent(Task.builder()
80+
.id(taskId)
81+
.contextId(CONTEXT_ID)
82+
.status(new TaskStatus(TaskState.TASK_STATE_SUBMITTED))
83+
.build());
84+
85+
for (int i = 0; i < eventCount; i++) {
86+
eventQueue.enqueueEvent(statusEventWithSequence(taskId, i));
87+
}
88+
89+
assertTrue(latch.await(10, TimeUnit.SECONDS), "All push notifications should complete within timeout");
90+
assertEquals(eventCount, deliveryOrder.size());
91+
for (int i = 0; i < eventCount; i++) {
92+
assertEquals(i, deliveryOrder.get(i),
93+
"Push notification " + i + " should be delivered in produced order, got " + deliveryOrder);
94+
}
95+
}
96+
97+
@Test
98+
public void testDifferentTasksAreIndependent() throws InterruptedException {
99+
String slowTaskId = "slow-task";
100+
String fastTaskId = "fast-task";
101+
CountDownLatch slowLatch = new CountDownLatch(1);
102+
CountDownLatch fastLatch = new CountDownLatch(1);
103+
AtomicInteger fastDeliveredWhileSlowBlocked = new AtomicInteger(-1);
104+
105+
PushNotificationSender sender = (event, snapshot) -> {
106+
TaskStatusUpdateEvent statusEvent = (TaskStatusUpdateEvent) event;
107+
if (statusEvent.taskId().equals(slowTaskId)) {
108+
try {
109+
Thread.sleep(2000);
110+
} catch (InterruptedException e) {
111+
Thread.currentThread().interrupt();
112+
}
113+
slowLatch.countDown();
114+
} else {
115+
fastDeliveredWhileSlowBlocked.set(slowLatch.getCount() > 0 ? 1 : 0);
116+
fastLatch.countDown();
117+
}
118+
};
119+
120+
mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, sender, queueManager);
121+
EventQueueUtil.start(mainEventBusProcessor);
122+
123+
EventQueue slowQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus)
124+
.taskId(slowTaskId).mainEventBus(mainEventBus).build().tap();
125+
EventQueue fastQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus)
126+
.taskId(fastTaskId).mainEventBus(mainEventBus).build().tap();
127+
128+
slowQueue.enqueueEvent(Task.builder().id(slowTaskId).contextId(CONTEXT_ID)
129+
.status(new TaskStatus(TaskState.TASK_STATE_SUBMITTED)).build());
130+
slowQueue.enqueueEvent(statusEventWithSequence(slowTaskId, 0));
131+
132+
fastQueue.enqueueEvent(Task.builder().id(fastTaskId).contextId(CONTEXT_ID)
133+
.status(new TaskStatus(TaskState.TASK_STATE_SUBMITTED)).build());
134+
fastQueue.enqueueEvent(statusEventWithSequence(fastTaskId, 0));
135+
136+
assertTrue(fastLatch.await(1, TimeUnit.SECONDS),
137+
"The fast task's push must not wait for the slow task's push");
138+
assertEquals(1, fastDeliveredWhileSlowBlocked.get(),
139+
"The fast task should have been delivered while the slow task was still blocked");
140+
141+
assertTrue(slowLatch.await(5, TimeUnit.SECONDS), "The slow task's push should eventually complete");
142+
}
143+
144+
@Test
145+
public void testSynchronousExecutorDoesNotLeakChainEntries() throws InterruptedException {
146+
// One entry per DISTINCT task leaks, not repeated pushes for the same one --
147+
// a single task's own map slot just gets overwritten each time. Only several
148+
// different task IDs reveal growth.
149+
int taskCount = 10;
150+
CountDownLatch latch = new CountDownLatch(taskCount);
151+
PushNotificationSender sender = (event, snapshot) -> latch.countDown();
152+
153+
mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, sender, queueManager);
154+
mainEventBusProcessor.setPushNotificationExecutor(Runnable::run);
155+
EventQueueUtil.start(mainEventBusProcessor);
156+
157+
for (int i = 0; i < taskCount; i++) {
158+
String taskId = "sync-task-" + i;
159+
EventQueue eventQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus)
160+
.taskId(taskId)
161+
.mainEventBus(mainEventBus)
162+
.build().tap();
163+
eventQueue.enqueueEvent(Task.builder()
164+
.id(taskId)
165+
.contextId(CONTEXT_ID)
166+
.status(new TaskStatus(TaskState.TASK_STATE_SUBMITTED))
167+
.build());
168+
eventQueue.enqueueEvent(statusEventWithSequence(taskId, 0));
169+
}
170+
171+
assertTrue(latch.await(5, TimeUnit.SECONDS), "All push notifications should complete");
172+
173+
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
174+
while (mainEventBusProcessor.pushNotificationChainCount() != 0 && System.nanoTime() < deadline) {
175+
Thread.sleep(10);
176+
}
177+
assertEquals(0, mainEventBusProcessor.pushNotificationChainCount(),
178+
"A completed task's chain entry must not be left in the map");
179+
}
180+
181+
private TaskStatusUpdateEvent statusEventWithSequence(String taskId, int sequence) {
182+
return new TaskStatusUpdateEvent(taskId, new TaskStatus(TaskState.TASK_STATE_WORKING, null, null),
183+
CONTEXT_ID, java.util.Map.of("sequence", sequence));
184+
}
185+
186+
private int extractSequence(TaskStatusUpdateEvent event) {
187+
return ((Number) event.metadata().get("sequence")).intValue();
188+
}
189+
}

0 commit comments

Comments
 (0)