Fix ObserveOn and add ParallelMerge Scheduler overload

rxjava-core/src/main/java/rx/Observable.java

@@ -4071,6 +4071,23 @@ public static <T> Observable<Observable<T>> parallelMerge(Observable<Observable<
         return OperationParallelMerge.parallelMerge(source, parallelObservables);
     }
 
+    /**
+     * Merges an <code>Observable<Observable<T>></code> to <code>Observable<Observable<T>></code>
+     * with number of inner Observables as defined by <code>parallelObservables</code> and runs each Observable on the defined Scheduler.
+     * <p>
+     * For example, if the original <code>Observable<Observable<T>></code> has 100 Observables to be emitted and <code>parallelObservables</code>
+     * is defined as 8, the 100 will be grouped onto 8 output Observables.
+     * <p>
+     * This is a mechanism for efficiently processing N number of Observables on a smaller N number of resources (typically CPU cores).
+     * 
+     * @param parallelObservables
+     *            the number of Observables to merge into.
+     * @return an Observable of Observables constrained to number defined by <code>parallelObservables</code>.
+     */
+    public static <T> Observable<Observable<T>> parallelMerge(Observable<Observable<T>> source, int parallelObservables, Scheduler scheduler) {
+        return OperationParallelMerge.parallelMerge(source, parallelObservables, scheduler);
+    }
+
     /**
      * Returns a {@link ConnectableObservable}, which waits until its
      * {@link ConnectableObservable#connect connect} method is called before it

rxjava-core/src/main/java/rx/operators/OperationObserveOn.java

@@ -27,8 +27,10 @@
 import rx.concurrency.CurrentThreadScheduler;
 import rx.concurrency.ImmediateScheduler;
 import rx.subscriptions.CompositeSubscription;
+import rx.subscriptions.Subscriptions;
 import rx.util.functions.Action0;
 import rx.util.functions.Action1;
+import rx.util.functions.Func2;
 
 /**
  * Asynchronously notify Observers on the specified Scheduler.
@@ -44,6 +46,7 @@ public static <T> OnSubscribeFunc<T> observeOn(Observable<? extends T> source, S
     private static class ObserveOn<T> implements OnSubscribeFunc<T> {
         private final Observable<? extends T> source;
         private final Scheduler scheduler;
+        private volatile Scheduler recursiveScheduler;
 
         final ConcurrentLinkedQueue<Notification<? extends T>> queue = new ConcurrentLinkedQueue<Notification<? extends T>>();
         final AtomicInteger counter = new AtomicInteger(0);
@@ -66,7 +69,7 @@ public Subscription onSubscribe(final Observer<? super T> observer) {
             }
         }
 
-        public Subscription observeOn(final Observer<? super T> observer, Scheduler scheduler) {
+        public Subscription observeOn(final Observer<? super T> observer, final Scheduler scheduler) {
             final CompositeSubscription s = new CompositeSubscription();
 
             s.add(source.materialize().subscribe(new Action1<Notification<? extends T>>() {
@@ -80,7 +83,22 @@ public void call(Notification<? extends T> e) {
                     // it will be 0 if it's the first notification or the scheduler has finished processing work
                     // and we need to start doing it again
                     if (counter.getAndIncrement() == 0) {
-                        processQueue(s, observer);
+                        if (recursiveScheduler == null) {
+                            s.add(scheduler.schedule(null, new Func2<Scheduler, T, Subscription>() {
+
+                                @Override
+                                public Subscription call(Scheduler innerScheduler, T state) {
+                                    // record innerScheduler so 'processQueue' can use it for all subsequent executions
+                                    recursiveScheduler = innerScheduler;
+
+                                    processQueue(s, observer);
+
+                                    return Subscriptions.empty();
+                                }
+                            }));
+                        } else {
+                            processQueue(s, observer);
+                        }
                     }
 
                 }
@@ -89,8 +107,13 @@ public void call(Notification<? extends T> e) {
             return s;
         }
 
-        private void processQueue(CompositeSubscription s, final Observer<? super T> observer) {
-            s.add(scheduler.schedule(new Action1<Action0>() {
+        /**
+         * This uses 'recursiveScheduler' NOT 'scheduler' as it should reuse the same scheduler each time it processes.
+         * This means it must first get the recursiveScheduler when it first executes.
+         */
+        private void processQueue(final CompositeSubscription s, final Observer<? super T> observer) {
+
+            s.add(recursiveScheduler.schedule(new Action1<Action0>() {
                 @Override
                 public void call(Action0 self) {
                     Notification<? extends T> not = queue.poll();

rxjava-core/src/main/java/rx/operators/OperationParallelMerge.java

@@ -18,29 +18,31 @@
 import java.util.concurrent.atomic.AtomicLong;
 
 import rx.Observable;
+import rx.Scheduler;
+import rx.concurrency.Schedulers;
 import rx.observables.GroupedObservable;
 import rx.util.functions.Func1;
 
 public class OperationParallelMerge {
 
-    public static <T> Observable<Observable<T>> parallelMerge(final Observable<Observable<T>> source, final int num) {
+    public static <T> Observable<Observable<T>> parallelMerge(final Observable<Observable<T>> source, final int parallelObservables) {
+        return parallelMerge(source, parallelObservables, Schedulers.currentThread());
+    }
+
+    public static <T> Observable<Observable<T>> parallelMerge(final Observable<Observable<T>> source, final int parallelObservables, final Scheduler scheduler) {
 
         return source.groupBy(new Func1<Observable<T>, Integer>() {
             final AtomicLong rollingCount = new AtomicLong();
 
             @Override
             public Integer call(Observable<T> o) {
-                return (int) rollingCount.incrementAndGet() % num;
+                return (int) rollingCount.incrementAndGet() % parallelObservables;
             }
         }).map(new Func1<GroupedObservable<Integer, Observable<T>>, Observable<T>>() {
 
-            /**
-             * Safe to cast from GroupedObservable to Observable so suppressing warning
-             */
-            @SuppressWarnings("unchecked")
             @Override
             public Observable<T> call(GroupedObservable<Integer, Observable<T>> o) {
-                return (Observable<T>) o;
+                return Observable.merge(o).observeOn(scheduler);
             }
 
         });

rxjava-core/src/main/java/rx/operators/OperationParallelMergeTest.java

@@ -18,11 +18,16 @@
 import static org.junit.Assert.*;
 
 import java.util.List;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
 
 import org.junit.Test;
 
 import rx.Observable;
+import rx.concurrency.Schedulers;
 import rx.subjects.PublishSubject;
+import rx.util.functions.Action1;
+import rx.util.functions.Func1;
 
 public class OperationParallelMergeTest {
 
@@ -42,8 +47,89 @@ public void testParallelMerge() {
         List<? super Observable<String>> threeList = threeStreams.toList().toBlockingObservable().last();
         List<? super Observable<String>> twoList = twoStreams.toList().toBlockingObservable().last();
 
+        System.out.println("two list: " + twoList);
+        System.out.println("three list: " + threeList);
+        System.out.println("four list: " + fourList);
+
         assertEquals(4, fourList.size());
         assertEquals(3, threeList.size());
         assertEquals(2, twoList.size());
     }
+
+    @Test
+    public void testNumberOfThreads() {
+        final ConcurrentHashMap<String, String> threads = new ConcurrentHashMap<String, String>();
+        Observable.merge(getStreams())
+                .toBlockingObservable().forEach(new Action1<String>() {
+
+                    @Override
+                    public void call(String o) {
+                        System.out.println("o: " + o + " Thread: " + Thread.currentThread());
+                        threads.put(Thread.currentThread().getName(), Thread.currentThread().getName());
+                    }
+                });
+
+        // without injecting anything, the getStream() method uses Interval which runs on a default scheduler
+        assertEquals(Runtime.getRuntime().availableProcessors(), threads.keySet().size());
+
+        // clear
+        threads.clear();
+
+        // now we parallelMerge into 3 streams and observeOn for each
+        // we expect 3 threads in the output
+        OperationParallelMerge.parallelMerge(getStreams(), 3)
+                .flatMap(new Func1<Observable<String>, Observable<String>>() {
+
+                    @Override
+                    public Observable<String> call(Observable<String> o) {
+                        // for each of the parallel 
+                        return o.observeOn(Schedulers.newThread());
+                    }
+                })
+                .toBlockingObservable().forEach(new Action1<String>() {
+
+                    @Override
+                    public void call(String o) {
+                        System.out.println("o: " + o + " Thread: " + Thread.currentThread());
+                        threads.put(Thread.currentThread().getName(), Thread.currentThread().getName());
+                    }
+                });
+
+        assertEquals(3, threads.keySet().size());
+    }
+
+    @Test
+    public void testNumberOfThreadsOnScheduledMerge() {
+        final ConcurrentHashMap<String, String> threads = new ConcurrentHashMap<String, String>();
+
+        // now we parallelMerge into 3 streams and observeOn for each
+        // we expect 3 threads in the output
+        Observable.merge(OperationParallelMerge.parallelMerge(getStreams(), 3, Schedulers.newThread()))
+                .toBlockingObservable().forEach(new Action1<String>() {
+
+                    @Override
+                    public void call(String o) {
+                        System.out.println("o: " + o + " Thread: " + Thread.currentThread());
+                        threads.put(Thread.currentThread().getName(), Thread.currentThread().getName());
+                    }
+                });
+
+        assertEquals(3, threads.keySet().size());
+    }
+
+    private static Observable<Observable<String>> getStreams() {
+        return Observable.range(0, 10).map(new Func1<Integer, Observable<String>>() {
+
+            @Override
+            public Observable<String> call(final Integer i) {
+                return Observable.interval(10, TimeUnit.MILLISECONDS).map(new Func1<Long, String>() {
+
+                    @Override
+                    public String call(Long l) {
+                        return "Stream " + i + "  Value: " + l;
+                    }
+                }).take(5);
+            }
+        });
+    }
 }