Add SingleThreadExecutor.executeOrRun and issue reads inline from the ledger's own thread

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

@@ -306,58 +306,81 @@ protected ExecutorService getBoundedExecutor(ExecutorService executor) {
     }
 
     protected ExecutorService addExecutorDecorators(ExecutorService executor) {
-        return new ForwardingExecutorService() {
-            @Override
-            protected ExecutorService delegate() {
-                return executor;
-            }
+        checkArgument(executor instanceof ThreadBoundExecutor, "Expected a ThreadBoundExecutor, got %s", executor);
+        return new DecoratedThread((ThreadBoundExecutor) executor);
+    }
 
-            @Override
-            public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks)
-                    throws InterruptedException {
-                return super.invokeAll(timedCallables(tasks));
-            }
+    /**
+     * One of the pool's threads with the task tracing and MDC decorators applied, which keeps exposing the thread
+     * identity of the underlying executor.
+     */
+    private class DecoratedThread extends ForwardingExecutorService implements ThreadBoundExecutor {
+        private final ThreadBoundExecutor executor;
 
-            @Override
-            public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks,
-                                                 long timeout, TimeUnit unit)
-                    throws InterruptedException {
-                return super.invokeAll(timedCallables(tasks), timeout, unit);
-            }
+        DecoratedThread(ThreadBoundExecutor executor) {
+            this.executor = executor;
+        }
 
-            @Override
-            public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
-                    throws InterruptedException, ExecutionException {
-                return super.invokeAny(timedCallables(tasks));
-            }
+        @Override
+        protected ExecutorService delegate() {
+            return executor;
+        }
 
-            @Override
-            public <T> T invokeAny(Collection<? extends Callable<T>> tasks,
-                                   long timeout, TimeUnit unit)
-                    throws InterruptedException, ExecutionException, TimeoutException {
-                return super.invokeAny(timedCallables(tasks), timeout, unit);
-            }
+        @Override
+        public boolean isCurrentThread() {
+            return executor.isCurrentThread();
+        }
 
-            @Override
-            public void execute(Runnable command) {
-                super.execute(timedRunnable(command));
-            }
+        @Override
+        public void executeOrRun(Runnable r) {
+            executor.executeOrRun(timedRunnable(r));
+        }
 
-            @Override
-            public <T> Future<T> submit(Callable<T> task) {
-                return super.submit(timedCallable(task));
-            }
+        @Override
+        public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks)
+                throws InterruptedException {
+            return super.invokeAll(timedCallables(tasks));
+        }
 
-            @Override
-            public Future<?> submit(Runnable task) {
-                return super.submit(timedRunnable(task));
-            }
+        @Override
+        public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks,
+                                             long timeout, TimeUnit unit)
+                throws InterruptedException {
+            return super.invokeAll(timedCallables(tasks), timeout, unit);
+        }
 
-            @Override
-            public <T> Future<T> submit(Runnable task, T result) {
-                return super.submit(timedRunnable(task), result);
-            }
-        };
+        @Override
+        public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
+                throws InterruptedException, ExecutionException {
+            return super.invokeAny(timedCallables(tasks));
+        }
+
+        @Override
+        public <T> T invokeAny(Collection<? extends Callable<T>> tasks,
+                               long timeout, TimeUnit unit)
+                throws InterruptedException, ExecutionException, TimeoutException {
+            return super.invokeAny(timedCallables(tasks), timeout, unit);
+        }
+
+        @Override
+        public void execute(Runnable command) {
+            super.execute(timedRunnable(command));
+        }
+
+        @Override
+        public <T> Future<T> submit(Callable<T> task) {
+            return super.submit(timedCallable(task));
+        }
+
+        @Override
+        public Future<?> submit(Runnable task) {
+            return super.submit(timedRunnable(task));
+        }
+
+        @Override
+        public <T> Future<T> submit(Runnable task, T result) {
+            return super.submit(timedRunnable(task), result);
+        }
     }
 
     /**

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

@@ -24,7 +24,6 @@
 import java.util.List;
 import java.util.concurrent.AbstractExecutorService;
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
@@ -45,7 +44,7 @@
  * proceed with the next tasks.
  */
 @CustomLog
-public class SingleThreadExecutor extends AbstractExecutorService implements ExecutorService, Runnable {
+public class SingleThreadExecutor extends AbstractExecutorService implements ThreadBoundExecutor, Runnable {
 
     private static final int MAX_DRAIN_BATCH_SIZE = 1024;
 
@@ -120,7 +119,7 @@ public void run() {
                 for (int i = 0; i < n; i++) {
                     Runnable task = localTasks[i];
                     localTasks[i] = null;
-                    if (!safeRunTask(task)) {
+                    if (!runQueuedTask(task)) {
                         return;
                     }
                 }
@@ -129,7 +128,7 @@ public void run() {
             // Clear the queue in orderly shutdown
             Runnable task;
             while ((task = queue.poll()) != null) {
-                safeRunTask(task);
+                runQueuedTask(task);
             }
         } catch (InterruptedException ie) {
             // Exit loop when interrupted
@@ -142,6 +141,19 @@ public void run() {
         }
     }
 
+    private boolean runQueuedTask(Runnable r) {
+        try {
+            return safeRunTask(r);
+        } finally {
+            decrementPendingTaskCount(1);
+        }
+    }
+
+    /**
+     * Runs a task, logging and counting a failure instead of propagating it.
+     *
+     * @return false when the task was interrupted
+     */
     private boolean safeRunTask(Runnable r) {
         try {
             r.run();
@@ -154,8 +166,6 @@ private boolean safeRunTask(Runnable r) {
                 tasksFailed.increment();
                 log.error().exception(t).log("Error while running task");
             }
-        } finally {
-            decrementPendingTaskCount(1);
         }
 
         return true;
@@ -220,6 +230,30 @@ public void execute(Runnable r) {
         executeRunnableOrList(r, null);
     }
 
+    @Override
+    public boolean isCurrentThread() {
+        return Thread.currentThread() == runner;
+    }
+
+    /**
+     * {@inheritDoc}
+     *
+     * <p>Failures of an inline run are logged and counted like those of queued tasks.
+     */
+    @Override
+    public void executeOrRun(Runnable r) {
+        if (state != State.Running) {
+            throw new RejectedExecutionException("Executor is shutting down");
+        }
+
+        if (isCurrentThread()) {
+            tasksCount.increment();
+            safeRunTask(r);
+        } else {
+            execute(r);
+        }
+    }
+
     @VisibleForTesting
     void executeRunnableOrList(Runnable runnable, List<Runnable> runnableList) {
         if (state != State.Running) {

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

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.bookkeeper.common.util;
+
+import java.util.concurrent.ExecutorService;
+
+/**
+ * An {@link ExecutorService} backed by a single thread, which can tell whether the caller is already on that
+ * thread and then run a task inline instead of queueing it.
+ *
+ * <p>Implemented by {@link SingleThreadExecutor} and by the threads an {@link OrderedExecutor} hands out from
+ * {@link OrderedExecutor#chooseThread(long)}, which are decorated when task tracing or MDC preservation is enabled.
+ */
+public interface ThreadBoundExecutor extends ExecutorService {
+
+    /**
+     * Whether the calling thread is the thread of this executor.
+     */
+    boolean isCurrentThread();
+
+    /**
+     * Runs the task inline when called from this executor's own thread, otherwise submits it like
+     * {@link #execute(Runnable)}. Like {@code execute}, it rejects the task once the executor is shut down.
+     *
+     * <p>The inline run bypasses the queue: a task submitted this way from the executor thread runs before the
+     * tasks already queued, nested inside the task that submitted it. Use it only where that reordering is
+     * acceptable.
+     */
+    void executeOrRun(Runnable r);
+}

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

@@ -22,14 +22,17 @@
 package org.apache.bookkeeper.common.util;
 
 import static org.hamcrest.Matchers.hasItem;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.AdditionalAnswers.answerVoid;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.spy;
 
 import java.util.Queue;
 import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.TimeUnit;
 import lombok.CustomLog;
@@ -82,6 +85,29 @@ public void tearDown() throws Exception {
         ThreadContext.clearMap();
     }
 
+    @Test
+    public void testDecoratedThreadsAreThreadBound() throws Exception {
+        for (OrderedExecutor executor : new OrderedExecutor[] {
+                OrderedExecutor.newBuilder().name("traced").numThreads(2).traceTaskExecution(true).build(),
+                OrderedExecutor.newBuilder().name("mdc").numThreads(2).preserveMdcForTaskExecution(true).build()}) {
+            try {
+                ThreadBoundExecutor thread = (ThreadBoundExecutor) executor.chooseThread(10);
+                assertFalse(thread.isCurrentThread());
+
+                // From the thread itself, executeOrRun runs the task before returning
+                CompletableFuture<Boolean> ranInline = new CompletableFuture<>();
+                thread.execute(() -> {
+                    boolean[] ran = new boolean[1];
+                    thread.executeOrRun(() -> ran[0] = thread.isCurrentThread());
+                    ranInline.complete(ran[0]);
+                });
+                assertTrue(ranInline.get(10, TimeUnit.SECONDS));
+            } finally {
+                executor.shutdown();
+            }
+        }
+    }
+
     @Test
     public void testMDCInvokeOrdered() throws Exception {
         OrderedExecutor executor = OrderedExecutor.newBuilder()