2.x: Add interruptible mode to Schedulers.from (ReactiveX#6370)

Author: akarnokd

src/main/java/io/reactivex/internal/schedulers/ExecutorScheduler.java

@@ -22,7 +22,7 @@
 import io.reactivex.internal.disposables.*;
 import io.reactivex.internal.functions.Functions;
 import io.reactivex.internal.queue.MpscLinkedQueue;
-import io.reactivex.internal.schedulers.ExecutorScheduler.ExecutorWorker.BooleanRunnable;
+import io.reactivex.internal.schedulers.ExecutorScheduler.ExecutorWorker.*;
 import io.reactivex.plugins.RxJavaPlugins;
 import io.reactivex.schedulers.*;
 
@@ -31,19 +31,22 @@
  */
 public final class ExecutorScheduler extends Scheduler {
 
+    final boolean interruptibleWorker;
+
     @NonNull
     final Executor executor;
 
     static final Scheduler HELPER = Schedulers.single();
 
-    public ExecutorScheduler(@NonNull Executor executor) {
+    public ExecutorScheduler(@NonNull Executor executor, boolean interruptibleWorker) {
         this.executor = executor;
+        this.interruptibleWorker = interruptibleWorker;
     }
 
     @NonNull
     @Override
     public Worker createWorker() {
-        return new ExecutorWorker(executor);
+        return new ExecutorWorker(executor, interruptibleWorker);
     }
 
     @NonNull
@@ -58,9 +61,15 @@ public Disposable scheduleDirect(@NonNull Runnable run) {
                 return task;
             }
 
-            BooleanRunnable br = new BooleanRunnable(decoratedRun);
-            executor.execute(br);
-            return br;
+            if (interruptibleWorker) {
+                InterruptibleRunnable interruptibleTask = new InterruptibleRunnable(decoratedRun, null);
+                executor.execute(interruptibleTask);
+                return interruptibleTask;
+            } else {
+                BooleanRunnable br = new BooleanRunnable(decoratedRun);
+                executor.execute(br);
+                return br;
+            }
         } catch (RejectedExecutionException ex) {
             RxJavaPlugins.onError(ex);
             return EmptyDisposable.INSTANCE;
@@ -111,6 +120,9 @@ public Disposable schedulePeriodicallyDirect(@NonNull Runnable run, long initial
     }
     /* public: test support. */
     public static final class ExecutorWorker extends Scheduler.Worker implements Runnable {
+
+        final boolean interruptibleWorker;
+
         final Executor executor;
 
         final MpscLinkedQueue<Runnable> queue;
@@ -121,9 +133,10 @@ public static final class ExecutorWorker extends Scheduler.Worker implements Run
 
         final CompositeDisposable tasks = new CompositeDisposable();
 
-        public ExecutorWorker(Executor executor) {
+        public ExecutorWorker(Executor executor, boolean interruptibleWorker) {
             this.executor = executor;
             this.queue = new MpscLinkedQueue<Runnable>();
+            this.interruptibleWorker = interruptibleWorker;
         }
 
         @NonNull
@@ -134,9 +147,24 @@ public Disposable schedule(@NonNull Runnable run) {
             }
 
             Runnable decoratedRun = RxJavaPlugins.onSchedule(run);
-            BooleanRunnable br = new BooleanRunnable(decoratedRun);
 
-            queue.offer(br);
+            Runnable task;
+            Disposable disposable;
+
+            if (interruptibleWorker) {
+                InterruptibleRunnable interruptibleTask = new InterruptibleRunnable(decoratedRun, tasks);
+                tasks.add(interruptibleTask);
+
+                task = interruptibleTask;
+                disposable = interruptibleTask;
+            } else {
+                BooleanRunnable runnableTask = new BooleanRunnable(decoratedRun);
+
+                task = runnableTask;
+                disposable = runnableTask;
+            }
+
+            queue.offer(task);
 
             if (wip.getAndIncrement() == 0) {
                 try {
@@ -149,7 +177,7 @@ public Disposable schedule(@NonNull Runnable run) {
                 }
             }
 
-            return br;
+            return disposable;
         }
 
         @NonNull
@@ -288,6 +316,97 @@ public void run() {
                 mar.replace(schedule(decoratedRun));
             }
         }
+
+        /**
+         * Wrapper for a {@link Runnable} with additional logic for handling interruption on
+         * a shared thread, similar to how Java Executors do it.
+         */
+        static final class InterruptibleRunnable extends AtomicInteger implements Runnable, Disposable {
+
+            private static final long serialVersionUID = -3603436687413320876L;
+
+            final Runnable run;
+
+            final DisposableContainer tasks;
+
+            volatile Thread thread;
+
+            static final int READY = 0;
+
+            static final int RUNNING = 1;
+
+            static final int FINISHED = 2;
+
+            static final int INTERRUPTING = 3;
+
+            static final int INTERRUPTED = 4;
+
+            InterruptibleRunnable(Runnable run, DisposableContainer tasks) {
+                this.run = run;
+                this.tasks = tasks;
+            }
+
+            @Override
+            public void run() {
+                if (get() == READY) {
+                    thread = Thread.currentThread();
+                    if (compareAndSet(READY, RUNNING)) {
+                        try {
+                            run.run();
+                        } finally {
+                            thread = null;
+                            if (compareAndSet(RUNNING, FINISHED)) {
+                                cleanup();
+                            } else {
+                                while (get() == INTERRUPTING) {
+                                    Thread.yield();
+                                }
+                                Thread.interrupted();
+                            }
+                        }
+                    } else {
+                        thread = null;
+                    }
+                }
+            }
+
+            @Override
+            public void dispose() {
+                for (;;) {
+                    int state = get();
+                    if (state >= FINISHED) {
+                        break;
+                    } else if (state == READY) {
+                        if (compareAndSet(READY, INTERRUPTED)) {
+                            cleanup();
+                            break;
+                        }
+                    } else {
+                        if (compareAndSet(RUNNING, INTERRUPTING)) {
+                            Thread t = thread;
+                            if (t != null) {
+                                t.interrupt();
+                                thread = null;
+                            }
+                            set(INTERRUPTED);
+                            cleanup();
+                            break;
+                        }
+                    }
+                }
+            }
+
+            void cleanup() {
+                if (tasks != null) {
+                    tasks.delete(this);
+                }
+            }
+
+            @Override
+            public boolean isDisposed() {
+                return get() >= FINISHED;
+            }
+        }
     }
 
     static final class DelayedRunnable extends AtomicReference<Runnable>

src/main/java/io/reactivex/schedulers/Schedulers.java

@@ -13,13 +13,13 @@
 
 package io.reactivex.schedulers;
 
+import java.util.concurrent.*;
+
 import io.reactivex.Scheduler;
-import io.reactivex.annotations.NonNull;
+import io.reactivex.annotations.*;
 import io.reactivex.internal.schedulers.*;
 import io.reactivex.plugins.RxJavaPlugins;
 
-import java.util.concurrent.*;
-
 /**
  * Static factory methods for returning standard Scheduler instances.
  * <p>
@@ -299,6 +299,9 @@ public static Scheduler single() {
      * a time delay or periodically will use the {@link #single()} scheduler for the timed waiting
      * before posting the actual task to the given executor.
      * <p>
+     * Tasks submitted to the {@link Scheduler.Worker} of this {@code Scheduler} are also not interruptible. Use the
+     * {@link #from(Executor, boolean)} overload to enable task interruption via this wrapper.
+     * <p>
      * If the provided executor supports the standard Java {@link ExecutorService} API,
      * cancelling tasks scheduled by this scheduler can be cancelled/interrupted by calling
      * {@link io.reactivex.disposables.Disposable#dispose()}. In addition, tasks scheduled with
@@ -329,7 +332,7 @@ public static Scheduler single() {
      * }
      * </code></pre>
      * <p>
-     * This type of scheduler is less sensitive to leaking {@link io.reactivex.Scheduler.Worker} instances, although
+     * This type of scheduler is less sensitive to leaking {@link Scheduler.Worker} instances, although
      * not disposing a worker that has timed/delayed tasks not cancelled by other means may leak resources and/or
      * execute those tasks "unexpectedly".
      * <p>
@@ -340,7 +343,67 @@ public static Scheduler single() {
      */
     @NonNull
     public static Scheduler from(@NonNull Executor executor) {
-        return new ExecutorScheduler(executor);
+        return new ExecutorScheduler(executor, false);
+    }
+
+    /**
+     * Wraps an {@link Executor} into a new Scheduler instance and delegates {@code schedule()}
+     * calls to it.
+     * <p>
+     * The tasks scheduled by the returned {@link Scheduler} and its {@link Scheduler.Worker}
+     * can be optionally interrupted.
+     * <p>
+     * If the provided executor doesn't support any of the more specific standard Java executor
+     * APIs, tasks scheduled with a time delay or periodically will use the
+     * {@link #single()} scheduler for the timed waiting
+     * before posting the actual task to the given executor.
+     * <p>
+     * If the provided executor supports the standard Java {@link ExecutorService} API,
+     * canceling tasks scheduled by this scheduler can be cancelled/interrupted by calling
+     * {@link io.reactivex.disposables.Disposable#dispose()}. In addition, tasks scheduled with
+     * a time delay or periodically will use the {@link #single()} scheduler for the timed waiting
+     * before posting the actual task to the given executor.
+     * <p>
+     * If the provided executor supports the standard Java {@link ScheduledExecutorService} API,
+     * canceling tasks scheduled by this scheduler can be cancelled/interrupted by calling
+     * {@link io.reactivex.disposables.Disposable#dispose()}. In addition, tasks scheduled with
+     * a time delay or periodically will use the provided executor. Note, however, if the provided
+     * {@code ScheduledExecutorService} instance is not single threaded, tasks scheduled
+     * with a time delay close to each other may end up executing in different order than
+     * the original schedule() call was issued. This limitation may be lifted in a future patch.
+     * <p>
+     * Starting, stopping and restarting this scheduler is not supported (no-op) and the provided
+     * executor's lifecycle must be managed externally:
+     * <pre><code>
+     * ExecutorService exec = Executors.newSingleThreadedExecutor();
+     * try {
+     *     Scheduler scheduler = Schedulers.from(exec, true);
+     *     Flowable.just(1)
+     *        .subscribeOn(scheduler)
+     *        .map(v -&gt; v + 1)
+     *        .observeOn(scheduler)
+     *        .blockingSubscribe(System.out::println);
+     * } finally {
+     *     exec.shutdown();
+     * }
+     * </code></pre>
+     * <p>
+     * This type of scheduler is less sensitive to leaking {@link Scheduler.Worker} instances, although
+     * not disposing a worker that has timed/delayed tasks not cancelled by other means may leak resources and/or
+     * execute those tasks "unexpectedly".
+     * <p>
+     * Note that this method returns a new {@link Scheduler} instance, even for the same {@link Executor} instance.
+     * @param executor
+     *          the executor to wrap
+     * @param interruptibleWorker if {@code true} the tasks submitted to the {@link Scheduler.Worker} will
+     * be interrupted when the task is disposed.
+     * @return the new Scheduler wrapping the Executor
+     * @since 2.2.6 - experimental
+     */
+    @NonNull
+    @Experimental
+    public static Scheduler from(@NonNull Executor executor, boolean interruptibleWorker) {
+        return new ExecutorScheduler(executor, interruptibleWorker);
     }
 
     /**