Integrate threadpool scheduling with AbstractRunnable (#106552)

Today `ThreadPool#scheduleWithFixedDelay` does not interact as expected
with `AbstractRunnable`: if the task fails or is rejected then this
isn't passed back to the relevant callback, and the task cannot specify
that it should be force-executed. This commit fixes that.

Backport of #106542 to 7.17

Author: DaveCTurner

server/src/main/java/org/elasticsearch/threadpool/Scheduler.java

@@ -106,7 +106,7 @@ static boolean awaitTermination(
      *         not be interrupted.
      */
     default Cancellable scheduleWithFixedDelay(Runnable command, TimeValue interval, String executor) {
-        return new ReschedulingRunnable(command, interval, executor, this, (e) -> {}, (e) -> {});
+        return new ReschedulingRunnable(command, interval, executor, this, e -> {}, e -> {});
     }
 
     /**
@@ -217,13 +217,25 @@ public void doRun() {
 
         @Override
         public void onFailure(Exception e) {
-            failureConsumer.accept(e);
+            try {
+                if (runnable instanceof AbstractRunnable) {
+                    ((AbstractRunnable) runnable).onFailure(e);
+                }
+            } finally {
+                failureConsumer.accept(e);
+            }
         }
 
         @Override
         public void onRejection(Exception e) {
             run = false;
-            rejectionConsumer.accept(e);
+            try {
+                if (runnable instanceof AbstractRunnable) {
+                    ((AbstractRunnable) runnable).onRejection(e);
+                }
+            } finally {
+                rejectionConsumer.accept(e);
+            }
         }
 
         @Override
@@ -238,6 +250,11 @@ public void onAfter() {
             }
         }
 
+        @Override
+        public boolean isForceExecution() {
+            return runnable instanceof AbstractRunnable && ((AbstractRunnable) runnable).isForceExecution();
+        }
+
         @Override
         public String toString() {
             return "ReschedulingRunnable{" + "runnable=" + runnable + ", interval=" + interval + '}';

server/src/main/java/org/elasticsearch/threadpool/ThreadPool.java

@@ -488,15 +488,13 @@ public void scheduleUnlessShuttingDown(TimeValue delay, String executor, Runnabl
 
     @Override
     public Cancellable scheduleWithFixedDelay(Runnable command, TimeValue interval, String executor) {
-        return new ReschedulingRunnable(command, interval, executor, this, (e) -> {
-            if (logger.isDebugEnabled()) {
-                logger.debug(() -> new ParameterizedMessage("scheduled task [{}] was rejected on thread pool [{}]", command, executor), e);
-            }
-        },
-            (e) -> logger.warn(
-                () -> new ParameterizedMessage("failed to run scheduled task [{}] on thread pool [{}]", command, executor),
-                e
-            )
+        return new ReschedulingRunnable(
+            command,
+            interval,
+            executor,
+            this,
+            e -> logger.debug(() -> new ParameterizedMessage("scheduled task [{}] was rejected on thread pool [{}]", command, executor), e),
+            e -> logger.warn(() -> new ParameterizedMessage("failed to run scheduled task [{}] on thread pool [{}]", command, executor), e)
         );
     }
 

server/src/test/java/org/elasticsearch/threadpool/ThreadPoolTests.java

@@ -11,16 +11,21 @@
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionRunnable;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
 import org.elasticsearch.common.util.concurrent.FutureUtils;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.MockLogAppender;
 
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.threadpool.ThreadPool.ESTIMATED_TIME_INTERVAL_SETTING;
 import static org.elasticsearch.threadpool.ThreadPool.LATE_TIME_INTERVAL_WARN_THRESHOLD_SETTING;
@@ -304,4 +309,146 @@ public String toString() {
             assertTrue(terminate(threadPool));
         }
     }
+
+    public void testScheduledOneShotRejection() {
+        final String name = "fixed-bounded";
+        final ThreadPool threadPool = new TestThreadPool(
+            getTestName(),
+            new FixedExecutorBuilder(Settings.EMPTY, name, between(1, 5), between(1, 5), "prefix")
+        );
+
+        final PlainActionFuture<Void> future = new PlainActionFuture<>();
+        final CountDownLatch latch = new CountDownLatch(1);
+        try {
+            blockExecution(threadPool.executor(name), latch);
+            threadPool.schedule(
+                ActionRunnable.run(future, () -> fail("should not execute")),
+                TimeValue.timeValueMillis(between(1, 100)),
+                name
+            );
+
+            expectThrows(EsRejectedExecutionException.class, () -> FutureUtils.get(future, 10, TimeUnit.SECONDS));
+        } finally {
+            latch.countDown();
+            assertTrue(terminate(threadPool));
+        }
+    }
+
+    public void testScheduledOneShotForceExecution() {
+        final String name = "fixed-bounded";
+        final ThreadPool threadPool = new TestThreadPool(
+            getTestName(),
+            new FixedExecutorBuilder(Settings.EMPTY, name, between(1, 5), between(1, 5), "prefix")
+        );
+
+        final PlainActionFuture<Void> future = new PlainActionFuture<>();
+        final CountDownLatch latch = new CountDownLatch(1);
+        try {
+            blockExecution(threadPool.executor(name), latch);
+            threadPool.schedule(forceExecution(ActionRunnable.run(future, () -> {})), TimeValue.timeValueMillis(between(1, 100)), name);
+
+            Thread.yield();
+            assertFalse(future.isDone());
+
+            latch.countDown();
+            FutureUtils.get(future, 10, TimeUnit.SECONDS); // shouldn't throw
+        } finally {
+            latch.countDown();
+            assertTrue(terminate(threadPool));
+        }
+    }
+
+    public void testScheduledFixedDelayRejection() {
+        final String name = "fixed-bounded";
+        final ThreadPool threadPool = new TestThreadPool(
+            getTestName(),
+            new FixedExecutorBuilder(Settings.EMPTY, name, between(1, 5), between(1, 5), "prefix")
+        );
+
+        final PlainActionFuture<Void> future = new PlainActionFuture<>();
+        final CountDownLatch latch = new CountDownLatch(1);
+        try {
+            threadPool.scheduleWithFixedDelay(
+                ActionRunnable.wrap(future, ignored -> Thread.yield()),
+                TimeValue.timeValueMillis(between(1, 100)),
+                name
+            );
+
+            while (future.isDone() == false) {
+                // might not block all threads the first time round if the scheduled runnable is running, so must keep trying
+                blockExecution(threadPool.executor(name), latch);
+            }
+            expectThrows(EsRejectedExecutionException.class, () -> FutureUtils.get(future));
+        } finally {
+            latch.countDown();
+            assertTrue(terminate(threadPool));
+        }
+    }
+
+    public void testScheduledFixedDelayForceExecution() {
+        final String name = "fixed-bounded";
+        final ThreadPool threadPool = new TestThreadPool(
+            getTestName(),
+            new FixedExecutorBuilder(Settings.EMPTY, name, between(1, 5), between(1, 5), "prefix")
+        );
+
+        final PlainActionFuture<Void> future = new PlainActionFuture<>();
+        final CountDownLatch latch = new CountDownLatch(1);
+        try {
+            blockExecution(threadPool.executor(name), latch);
+
+            threadPool.scheduleWithFixedDelay(
+                forceExecution(ActionRunnable.run(future, Thread::yield)),
+                TimeValue.timeValueMillis(between(1, 100)),
+                name
+            );
+
+            assertFalse(future.isDone());
+
+            latch.countDown();
+            FutureUtils.get(future, 10, TimeUnit.SECONDS); // shouldn't throw
+        } finally {
+            latch.countDown();
+            assertTrue(terminate(threadPool));
+        }
+    }
+
+    private static AbstractRunnable forceExecution(AbstractRunnable delegate) {
+        return new AbstractRunnable() {
+            @Override
+            public void onFailure(Exception e) {
+                delegate.onFailure(e);
+            }
+
+            @Override
+            protected void doRun() {
+                delegate.run();
+            }
+
+            @Override
+            public void onRejection(Exception e) {
+                delegate.onRejection(e);
+            }
+
+            @Override
+            public void onAfter() {
+                delegate.onAfter();
+            }
+
+            @Override
+            public boolean isForceExecution() {
+                return true;
+            }
+        };
+    }
+
+    private static void blockExecution(ExecutorService executor, CountDownLatch latch) {
+        while (true) {
+            try {
+                executor.execute(() -> safeAwait(latch));
+            } catch (EsRejectedExecutionException e) {
+                break;
+            }
+        }
+    }
 }