fix: resolve potential exception and blocking issue in queue capacity handling

Author: nodece

bookkeeper-common/src/main/java/org/apache/bookkeeper/common/util/SingleThreadExecutor.java

@@ -143,9 +143,6 @@ public void run() {
 
     private boolean safeRunTask(Runnable r) {
         try {
-            if (maxQueueCapacity > 0) {
-                waiterCountUpdater.decrementAndGet(this);
-            }
             r.run();
             tasksCompleted.increment();
         } catch (Throwable t) {
@@ -156,6 +153,8 @@ private boolean safeRunTask(Runnable r) {
                 tasksFailed.increment();
                 log.error("Error while running task: {}", t.getMessage(), t);
             }
+        } finally {
+            releaseQueuePermits(1);
         }
 
         return true;
@@ -219,25 +218,44 @@ public long getFailedTasksCount() {
 
     @Override
     public void execute(Runnable r) {
-        execute(r, null);
+        executeRunnableOrList(r, null);
     }
 
     @VisibleForTesting
-    void execute(Runnable r, List<Runnable> runnableList) {
+    void executeRunnableOrList(Runnable runnable, List<Runnable> runnableList) {
         if (state != State.Running) {
             throw new RejectedExecutionException("Executor is shutting down");
         }
 
+        boolean hasSingle = runnable != null;
+        boolean hasList = runnableList != null && !runnableList.isEmpty();
+
+        if (hasSingle == hasList) {
+            // Both are provided or both are missing
+            throw new IllegalArgumentException("Provide either 'runnable' or a non-empty 'runnableList', not both.");
+        }
+
         try {
             if (!rejectExecution) {
-                queue.put(r);
-                tasksCount.increment();
+                if (hasSingle) {
+                    queue.put(runnable);
+                    tasksCount.increment();
+                } else {
+                    for (Runnable task : runnableList) {
+                        queue.put(task);
+                        tasksCount.increment();
+                    }
+                }
             } else {
-                int delta = r != null ? 1 : runnableList.size();
-                validateQueueCapacity(delta);
-                if (r != null ? queue.offer(r) : queue.addAll(runnableList)) {
-                    tasksCount.add(delta);
+                int permits = runnable != null ? 1 : runnableList.size();
+                reserveQueuePermits(permits);
+                boolean success = hasSingle
+                        ? queue.offer(runnable)
+                        : queue.addAll(runnableList);
+                if (success) {
+                    tasksCount.add(permits);
                 } else {
+                    releaseQueuePermits(permits);
                     reject();
                 }
             }
@@ -246,12 +264,36 @@ void execute(Runnable r, List<Runnable> runnableList) {
         }
     }
 
-    private void validateQueueCapacity(int delta) {
-        if (maxQueueCapacity > 0 && waiterCountUpdater.addAndGet(this, delta) > maxQueueCapacity) {
+    private void reserveQueuePermits(int permits) {
+        if (maxQueueCapacity <= 0) {
+            return; // // Unlimited capacity, nothing to release
+        }
+
+        if (permits < 0) {
+            throw new IllegalArgumentException("Number of permits must be non-negative");
+        }
+
+        if (waiterCountUpdater.addAndGet(this, permits) > maxQueueCapacity) {
+            releaseQueuePermits(permits);
             reject();
         }
     }
 
+    private void releaseQueuePermits(int permits) {
+        if (maxQueueCapacity <= 0) {
+            return; // // Unlimited capacity, nothing to release
+        }
+
+        if (permits < 0) {
+            throw new IllegalArgumentException("Number of permits must be non-negative");
+        }
+
+        int waiterCount = waiterCountUpdater.addAndGet(this, -permits);
+        if (log.isDebugEnabled()) {
+            log.debug("Released {} permits, current waiter count: {}", permits, waiterCount);
+        }
+    }
+
     private void reject() {
         tasksRejected.increment();
         throw new ExecutorRejectedException("Executor queue is full");
@@ -321,6 +363,11 @@ public Number getSample() {
                 });
     }
 
+    @VisibleForTesting
+    int getWaiterCount() {
+        return waiterCountUpdater.get(this);
+    }
+
     private static class ExecutorRejectedException extends RejectedExecutionException {
 
         private ExecutorRejectedException(String msg) {

bookkeeper-common/src/test/java/org/apache/bookkeeper/common/util/TestSingleThreadExecutor.java

@@ -20,11 +20,13 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import io.netty.util.concurrent.DefaultThreadFactory;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.BrokenBarrierException;
 import java.util.concurrent.CountDownLatch;
@@ -36,6 +38,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import lombok.Cleanup;
+import org.assertj.core.util.Lists;
 import org.awaitility.Awaitility;
 import org.junit.Test;
 
@@ -122,54 +125,43 @@ public void testRejectWhenDrainToInProgressAndQueueIsEmpty() throws Exception {
         @Cleanup("shutdownNow")
         SingleThreadExecutor ste = new SingleThreadExecutor(THREAD_FACTORY, 10, true);
 
-        CyclicBarrier barrier = new CyclicBarrier(10);
-        CountDownLatch startedLatch = new CountDownLatch(1);
+        CountDownLatch waitedLatch = new CountDownLatch(1);
         List<Runnable> tasks = new ArrayList<>();
 
         for (int i = 0; i < 10; i++) {
-            int n = i;
             tasks.add(() -> {
-                if (n == 0) {
-                    startedLatch.countDown();
-                } else {
-                    try {
-                        barrier.await();
-                    } catch (InterruptedException | BrokenBarrierException e) {
-                        // ignore
-                    }
+                try {
+                    // Block remaining tasks to simulate a stuck queue.
+                    waitedLatch.await();
+                } catch (Exception e) {
+                    // ignored
                 }
             });
         }
-        ste.execute(null, tasks);
-
-        // Wait until the first task is done.
-        try {
-            startedLatch.await();
-        } catch (InterruptedException e) {
-            throw new RuntimeException(e);
-        }
+        ste.executeRunnableOrList(null, tasks);
 
-        // Next task should go through, because the runner thread has already pulled out the first and second items
-        // from the queue.
-        List<Runnable> nextTasks = new ArrayList<>();
-        nextTasks.add(() -> {
-        });
-        nextTasks.add(() -> {
-        });
-        ste.execute(null, nextTasks);
+        Awaitility.await().pollDelay(1, TimeUnit.SECONDS).untilAsserted(() -> assertEquals(10, ste.getWaiterCount()));
 
-        // Now the queue is really full and should reject tasks
-        try {
-            ste.execute(() -> {
-            });
-            fail("should have rejected the task");
-        } catch (RejectedExecutionException e) {
-            // Expected
-        }
+        // Now the queue is really full and should reject tasks.
+        assertThrows(RejectedExecutionException.class, () -> ste.execute(() -> {
+        }));
 
-        assertEquals(12, ste.getSubmittedTasksCount());
+        assertEquals(10, ste.getWaiterCount());
         assertEquals(1, ste.getRejectedTasksCount());
         assertEquals(0, ste.getFailedTasksCount());
+
+        // Now we can unblock the waited tasks.
+        waitedLatch.countDown();
+
+        // Check the tasks are completed.
+        Awaitility.await().pollDelay(1, TimeUnit.SECONDS).untilAsserted(() -> assertEquals(0, ste.getWaiterCount()));
+
+        // Invalid cases - should throw IllegalArgumentException.
+        assertThrows(IllegalArgumentException.class, () -> ste.executeRunnableOrList(null, null));
+        assertThrows(IllegalArgumentException.class, () -> ste.executeRunnableOrList(null, Collections.emptyList()));
+        assertThrows(IllegalArgumentException.class, () -> ste.executeRunnableOrList(() -> {
+        }, Lists.newArrayList(() -> {
+        })));
     }
 
     @Test