Support scheduled commands in current context

Adds support for scheduling commands to run at a later time on another
thread pool in the current thread's context:

```java
Runnable someCommand = () -> {System.err.println("Demo");};
someCommand = threadPool.getThreadContext().preserveContext(someCommand);
threadPool.schedule(timeValueMinutes(1), Names.GENERAL, someCommand);
```

This happens automatically for calls to `threadPool.execute` but `schedule`
and `scheduleWithFixedDelay` don't do that, presumably because scheduled
tasks are usually context-less. Rather than preserve the current context
on all scheduled tasks this just makes it possible to preserve it using
the syntax above.

To make this all go it moves the Runnables that wrap the commands from
EsThreadPoolExecutor into ThreadContext.

This, or something like it, is required to support reindex throttling.

Author: nik9000

buildSrc/src/main/resources/checkstyle_suppressions.xml

@@ -390,7 +390,6 @@
   <suppress files="core[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]common[/\\]util[/\\]CollectionUtils.java" checks="LineLength" />
   <suppress files="core[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]common[/\\]util[/\\]ExtensionPoint.java" checks="LineLength" />
   <suppress files="core[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]common[/\\]util[/\\]concurrent[/\\]EsExecutors.java" checks="LineLength" />
-  <suppress files="core[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]common[/\\]util[/\\]concurrent[/\\]EsThreadPoolExecutor.java" checks="LineLength" />
   <suppress files="core[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]common[/\\]util[/\\]concurrent[/\\]PrioritizedEsThreadPoolExecutor.java" checks="LineLength" />
   <suppress files="core[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]common[/\\]util[/\\]concurrent[/\\]ThreadBarrier.java" checks="LineLength" />
   <suppress files="core[/\\]src[/\\]main[/\\]java[/\\]org[/\\]elasticsearch[/\\]common[/\\]util[/\\]concurrent[/\\]ThreadContext.java" checks="LineLength" />

core/src/main/java/org/elasticsearch/common/util/concurrent/EsThreadPoolExecutor.java

@@ -40,11 +40,14 @@ public class EsThreadPoolExecutor extends ThreadPoolExecutor {
      */
     private final String name;
 
-    EsThreadPoolExecutor(String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, ThreadContext contextHolder) {
+    EsThreadPoolExecutor(String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
+            BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, ThreadContext contextHolder) {
         this(name, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, new EsAbortPolicy(), contextHolder);
     }
 
-    EsThreadPoolExecutor(String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, XRejectedExecutionHandler handler, ThreadContext contextHolder) {
+    EsThreadPoolExecutor(String name, int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit,
+            BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, XRejectedExecutionHandler handler,
+            ThreadContext contextHolder) {
         super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
         this.name = name;
         this.contextHolder = contextHolder;
@@ -133,112 +136,10 @@ public String toString() {
     }
 
     protected Runnable wrapRunnable(Runnable command) {
-        final Runnable wrappedCommand;
-        if (command instanceof AbstractRunnable) {
-            wrappedCommand = new FilterAbstractRunnable(contextHolder, (AbstractRunnable) command);
-        } else {
-            wrappedCommand = new FilterRunnable(contextHolder, command);
-        }
-        return wrappedCommand;
+        return contextHolder.preserveContext(command);
     }
 
     protected Runnable unwrap(Runnable runnable) {
-        if (runnable instanceof FilterAbstractRunnable) {
-            return ((FilterAbstractRunnable) runnable).in;
-        } else if (runnable instanceof FilterRunnable) {
-            return ((FilterRunnable) runnable).in;
-        }
-        return runnable;
+        return contextHolder.unwrap(runnable);
     }
-
-    private  class FilterAbstractRunnable extends AbstractRunnable {
-        private final ThreadContext contextHolder;
-        private final AbstractRunnable in;
-        private final ThreadContext.StoredContext ctx;
-
-        FilterAbstractRunnable(ThreadContext contextHolder, AbstractRunnable in) {
-            this.contextHolder = contextHolder;
-            ctx = contextHolder.newStoredContext();
-            this.in = in;
-        }
-
-        @Override
-        public boolean isForceExecution() {
-            return in.isForceExecution();
-        }
-
-        @Override
-        public void onAfter() {
-            in.onAfter();
-        }
-
-        @Override
-        public void onFailure(Throwable t) {
-            in.onFailure(t);
-        }
-
-        @Override
-        public void onRejection(Throwable t) {
-            in.onRejection(t);
-        }
-
-        @Override
-        protected void doRun() throws Exception {
-            boolean whileRunning = false;
-            try (ThreadContext.StoredContext ingore = contextHolder.stashContext()){
-                ctx.restore();
-                whileRunning = true;
-                in.doRun();
-                whileRunning = false;
-            } catch (IllegalStateException ex) {
-                if (whileRunning || isShutdown() == false) {
-                    throw ex;
-                }
-                // if we hit an ISE here we have been shutting down
-                // this comes from the threadcontext and barfs if
-                // our threadpool has been shutting down
-            }
-        }
-
-        @Override
-        public String toString() {
-            return in.toString();
-        }
-
-    }
-
-    private  class FilterRunnable implements Runnable {
-        private final ThreadContext contextHolder;
-        private final Runnable in;
-        private final ThreadContext.StoredContext ctx;
-
-        FilterRunnable(ThreadContext contextHolder, Runnable in) {
-            this.contextHolder = contextHolder;
-            ctx = contextHolder.newStoredContext();
-            this.in = in;
-        }
-
-        @Override
-        public void run() {
-            boolean whileRunning = false;
-            try (ThreadContext.StoredContext ingore = contextHolder.stashContext()){
-                ctx.restore();
-                whileRunning = true;
-                in.run();
-                whileRunning = false;
-            } catch (IllegalStateException ex) {
-                if (whileRunning || isShutdown() == false) {
-                    throw ex;
-                }
-                // if we hit an ISE here we have been shutting down
-                // this comes from the threadcontext and barfs if
-                // our threadpool has been shutting down
-            }
-        }
-        @Override
-        public String toString() {
-            return in.toString();
-        }
-    }
-
 }

core/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java

@@ -200,6 +200,36 @@ public <T> T getTransient(String key) {
         return (T) threadLocal.get().transientHeaders.get(key);
     }
 
+    /**
+     * Saves the current thread context and wraps command in a Runnable that restores that context before running command. If
+     * <code>command</code> has already been passed through this method then it is returned unaltered rather than wrapped twice.
+     */
+    public Runnable preserveContext(Runnable command) {
+        if (command instanceof ContextPreservingAbstractRunnable) {
+            return command;
+        }
+        if (command instanceof ContextPreservingRunnable) {
+            return command;
+        }
+        if (command instanceof AbstractRunnable) {
+            return new ContextPreservingAbstractRunnable((AbstractRunnable) command);
+        }
+        return new ContextPreservingRunnable(command);
+    }
+
+    /**
+     * Unwraps a command that was previously wrapped by {@link #preserveContext(Runnable)}.
+     */
+    public Runnable unwrap(Runnable command) {
+        if (command instanceof ContextPreservingAbstractRunnable) {
+            return ((ContextPreservingAbstractRunnable) command).unwrap();
+        }
+        if (command instanceof ContextPreservingRunnable) {
+            return ((ContextPreservingRunnable) command).unwrap();
+        }
+        return command;
+    }
+
     public interface StoredContext extends AutoCloseable {
         @Override
         void close();
@@ -356,4 +386,104 @@ public void close() {
             }
         }
     }
+
+    /**
+     * Wraps a Runnable to preserve the thread context.
+     */
+    class ContextPreservingRunnable implements Runnable {
+        private final Runnable in;
+        private final ThreadContext.StoredContext ctx;
+
+        ContextPreservingRunnable(Runnable in) {
+            ctx = newStoredContext();
+            this.in = in;
+        }
+
+        @Override
+        public void run() {
+            boolean whileRunning = false;
+            try (ThreadContext.StoredContext ingore = stashContext()){
+                ctx.restore();
+                whileRunning = true;
+                in.run();
+                whileRunning = false;
+            } catch (IllegalStateException ex) {
+                if (whileRunning || threadLocal.closed.get() == false) {
+                    throw ex;
+                }
+                // if we hit an ISE here we have been shutting down
+                // this comes from the threadcontext and barfs if
+                // our threadpool has been shutting down
+            }
+        }
+
+        @Override
+        public String toString() {
+            return in.toString();
+        }
+
+        public Runnable unwrap() {
+            return in;
+        }
+    }
+
+    /**
+     * Wraps an AbstractRunnable to preserve the thread context.
+     */
+    public class ContextPreservingAbstractRunnable extends AbstractRunnable {
+        private final AbstractRunnable in;
+        private final ThreadContext.StoredContext ctx;
+
+        private ContextPreservingAbstractRunnable(AbstractRunnable in) {
+            ctx = newStoredContext();
+            this.in = in;
+        }
+
+        @Override
+        public boolean isForceExecution() {
+            return in.isForceExecution();
+        }
+
+        @Override
+        public void onAfter() {
+            in.onAfter();
+        }
+
+        @Override
+        public void onFailure(Throwable t) {
+            in.onFailure(t);
+        }
+
+        @Override
+        public void onRejection(Throwable t) {
+            in.onRejection(t);
+        }
+
+        @Override
+        protected void doRun() throws Exception {
+            boolean whileRunning = false;
+            try (ThreadContext.StoredContext ingore = stashContext()){
+                ctx.restore();
+                whileRunning = true;
+                in.doRun();
+                whileRunning = false;
+            } catch (IllegalStateException ex) {
+                if (whileRunning || threadLocal.closed.get() == false) {
+                    throw ex;
+                }
+                // if we hit an ISE here we have been shutting down
+                // this comes from the threadcontext and barfs if
+                // our threadpool has been shutting down
+            }
+        }
+
+        @Override
+        public String toString() {
+            return in.toString();
+        }
+
+        public AbstractRunnable unwrap() {
+            return in;
+        }
+    }
 }

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

@@ -341,10 +341,18 @@ public ThreadPoolStats stats() {
         return new ThreadPoolStats(stats);
     }
 
+    /**
+     * Get the generic executor. This executor's {@link Executor#execute(Runnable)} method will run the Runnable it is given in
+     * the {@link ThreadContext} of the thread that queues it.
+     */
     public Executor generic() {
         return executor(Names.GENERIC);
     }
 
+    /**
+     * Get the executor with the given name. This executor's {@link Executor#execute(Runnable)} method will run the Runnable it is given in
+     * the {@link ThreadContext} of the thread that queues it.
+     */
     public Executor executor(String name) {
         Executor executor = executors.get(name).executor();
         if (executor == null) {
@@ -357,10 +365,31 @@ public ScheduledExecutorService scheduler() {
         return this.scheduler;
     }
 
+    /**
+     * Schedules a periodic action that always runs on the scheduler thread.
+     *
+     * @param command the action to take
+     * @param interval the delay interval
+     * @return a ScheduledFuture who's get will return when the task is complete and throw an exception if it is canceled
+     */
     public ScheduledFuture<?> scheduleWithFixedDelay(Runnable command, TimeValue interval) {
         return scheduler.scheduleWithFixedDelay(new LoggingRunnable(command), interval.millis(), interval.millis(), TimeUnit.MILLISECONDS);
     }
 
+    /**
+     * Schedules a one-shot command to run after a given delay. The command is not run in the context of the calling thread. To preserve the
+     * context of the calling thread you may call <code>threadPool.getThreadContext().preserveContext</code> on the runnable before passing
+     * it to this method.
+     *
+     * @param delay delay before the task executes
+     * @param name the name of the thread pool on which to execute this task. SAME means "execute on the scheduler thread" which changes the
+     *        meaning of the ScheduledFuture returned by this method. In that case the ScheduledFuture will complete only when the command
+     *        completes.
+     * @param command the command to run
+     * @return a ScheduledFuture who's get will return when the task is has been added to its target thread pool and throw an exception if
+     *         the task is canceled before it was added to its target thread pool. Once the task has been added to its target thread pool
+     *         the ScheduledFuture will cannot interact with it.
+     */
     public ScheduledFuture<?> schedule(TimeValue delay, String name, Runnable command) {
         if (!Names.SAME.equals(name)) {
             command = new ThreadedRunnable(command, executor(name));