Revisit the semaphore handling

Author: kabir

server-common/src/main/java/io/a2a/server/events/EventQueue.java

@@ -309,19 +309,28 @@ void distributeToChildren(EventQueueItem item) {
             int childCount = children.size();
             if (LOGGER.isDebugEnabled()) {
                 LOGGER.debug("MainQueue[{}]: Distributing event {} to {} children",
-                            taskId, item.getEvent().getClass().getSimpleName(), childCount);
+                        taskId, item.getEvent().getClass().getSimpleName(), childCount);
             }
             children.forEach(child -> {
                 LOGGER.debug("MainQueue[{}]: Enqueueing event {} to child queue",
-                            taskId, item.getEvent().getClass().getSimpleName());
+                        taskId, item.getEvent().getClass().getSimpleName());
                 child.internalEnqueueItem(item);
             });
             if (LOGGER.isDebugEnabled()) {
                 LOGGER.debug("MainQueue[{}]: Completed distribution of {} to {} children",
-                            taskId, item.getEvent().getClass().getSimpleName(), childCount);
+                        taskId, item.getEvent().getClass().getSimpleName(), childCount);
             }
         }
 
+        /**
+         * Release the semaphore after event processing is complete.
+         * Called by MainEventBusProcessor in finally block to ensure release even on exceptions.
+         * Balances the acquire() in enqueueEvent() - protects MainEventBus throughput.
+         */
+        void releaseSemaphore() {
+            semaphore.release();
+        }
+
         /**
          * Get the count of active child queues.
          * Used for testing to verify reference counting mechanism.
@@ -391,17 +400,12 @@ public void enqueueItem(EventQueueItem item) {
 
         private void internalEnqueueItem(EventQueueItem item) {
             // Internal method called by MainEventBusProcessor to add to local queue
+            // Note: Semaphore is managed by parent MainQueue (acquire/release), not ChildQueue
             Event event = item.getEvent();
             if (isClosed()) {
                 LOGGER.warn("ChildQueue is closed. Event will not be enqueued. {} {}", this, event);
                 return;
             }
-            try {
-                semaphore.acquire();
-            } catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
-                throw new RuntimeException("Unable to acquire the semaphore to enqueue the event", e);
-            }
             queue.add(item);
             LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this);
         }
@@ -418,7 +422,7 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl
                     if (item != null) {
                         Event event = item.getEvent();
                         LOGGER.debug("Dequeued event item (no wait) {} {}", this, event instanceof Throwable ? event.toString() : event);
-                        semaphore.release();
+                        // Note: Semaphore is managed by parent MainQueue, not released here
                     }
                     return item;
                 }
@@ -428,7 +432,7 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl
                     if (item != null) {
                         Event event = item.getEvent();
                         LOGGER.debug("Dequeued event item (waiting) {} {}", this, event instanceof Throwable ? event.toString() : event);
-                        semaphore.release();
+                        // Note: Semaphore is managed by parent MainQueue, not released here
                     } else {
                         LOGGER.trace("Dequeue timeout (null) from ChildQueue {}", System.identityHashCode(this));
                     }

server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java

@@ -140,32 +140,40 @@ private void processEvent(MainEventBusContext context) {
         LOGGER.debug("MainEventBusProcessor: Processing event for task {}: {} (queue type: {})",
                     taskId, event.getClass().getSimpleName(), eventQueue.getClass().getSimpleName());
 
-        // Step 1: Update TaskStore FIRST (persistence before clients see it)
-        updateTaskStore(taskId, event);
-
-        // Step 2: Send push notification AFTER persistence (ensures notification sees latest state)
-        sendPushNotification(taskId);
-
-        // Step 3: Then distribute to ChildQueues (clients see it AFTER persistence + notification)
-        if (eventQueue instanceof EventQueue.MainQueue mainQueue) {
-            int childCount = mainQueue.getChildCount();
-            LOGGER.debug("MainEventBusProcessor: Distributing event to {} children for task {}", childCount, taskId);
-            mainQueue.distributeToChildren(context.eventQueueItem());
-            LOGGER.debug("MainEventBusProcessor: Distributed event {} to {} children for task {}",
-                        event.getClass().getSimpleName(), childCount, taskId);
-        } else {
-            LOGGER.warn("MainEventBusProcessor: Expected MainQueue but got {} for task {}",
-                    eventQueue.getClass().getSimpleName(), taskId);
-        }
+        try {
+            // Step 1: Update TaskStore FIRST (persistence before clients see it)
+            updateTaskStore(taskId, event);
+
+            // Step 2: Send push notification AFTER persistence (ensures notification sees latest state)
+            sendPushNotification(taskId);
+
+            // Step 3: Then distribute to ChildQueues (clients see it AFTER persistence + notification)
+            if (eventQueue instanceof EventQueue.MainQueue mainQueue) {
+                int childCount = mainQueue.getChildCount();
+                LOGGER.debug("MainEventBusProcessor: Distributing event to {} children for task {}", childCount, taskId);
+                mainQueue.distributeToChildren(context.eventQueueItem());
+                LOGGER.debug("MainEventBusProcessor: Distributed event {} to {} children for task {}",
+                            event.getClass().getSimpleName(), childCount, taskId);
+            } else {
+                LOGGER.warn("MainEventBusProcessor: Expected MainQueue but got {} for task {}",
+                        eventQueue.getClass().getSimpleName(), taskId);
+            }
 
-        LOGGER.debug("MainEventBusProcessor: Completed processing event for task {}", taskId);
+            LOGGER.debug("MainEventBusProcessor: Completed processing event for task {}", taskId);
 
-        // Step 4: Notify callback after all processing is complete
-        callback.onEventProcessed(taskId, event);
+            // Step 4: Notify callback after all processing is complete
+            callback.onEventProcessed(taskId, event);
 
-        // Step 5: If this is a final event, notify task finalization
-        if (isFinalEvent(event)) {
-            callback.onTaskFinalized(taskId);
+            // Step 5: If this is a final event, notify task finalization
+            if (isFinalEvent(event)) {
+                callback.onTaskFinalized(taskId);
+            }
+        } finally {
+            // ALWAYS release semaphore, even if processing fails
+            // Balances the acquire() in MainQueue.enqueueEvent()
+            if (eventQueue instanceof EventQueue.MainQueue mainQueue) {
+                mainQueue.releaseSemaphore();
+            }
         }
     }