RxJava - Maximizing Parallelization

Earlier I posted an entry about achieving parallelization with RxJava. But I have learned a few things from experiments and the RxJava devs to maximize parallelization and get better performance. This post is going to focus on computation parallelization and not really IO. Computation is where you really have to constrain the number of threads you use, and roughly you get the best performance accounting for the number of cores your computer has.

Maximizing Computation with Executors

In my previous article on parallelization, I had a short conversation with David Karnok, the current maintainer of RxJava, in the comments section. I asked him how many threads the computation() Scheduler creates and utilizes. This was his response:

If [your computer has] 8 cores, this will only use 5 of them. If you have 4 cores, two groups will share the same core.

Thinking about this, I began to have many questions. I am not an expert on multithreading, but I read most of Java Concurrency in Practice by Brian Goetz a year or so ago. I recall in Chapter 8, Brian talked about sizing thread pools for computation-intensive tasks. Pragmatically, you can follow this formula to roughly get an optimal number of threads. It's not perfect, but its a simple rule to get well-balanced computation performance.

Nthreads = NCPU + 1

So if you have 4 cores, that means you should have 5 threads. If you have 8 cores, you should have 9 threads. That extra thread is to fill up idle CPU cycles when the other threads get short moments of idleness. You can get this formula dynamically using a simple Java expression.

int threadCt = Runtime.getRuntime().availableProcessors() + 1;

I may be very naive and might be overlooking something. Or maybe the RxJava developers simply had different goals and wanted to keep the number of threads conservative. But I began to suspect Schedulers.computation() was not maximizing parallelization since it was not explicitly designed for it.

Observable.range(1,1000)
.flatMap(i -> Observable.just(i)
    .subscribeOn(Schedulers.computation())
    .map(i2 -> intenseCalculation(i2))
).subscribe(System.out::println)

I was running a process at work in a similar manner as shown above, and the process took 225 seconds to emit all the items and do an intensive calculation on each one in parallel.
I had a hunch and decided to try Goetz’s simple rough formula for calculation thread pool sizing. This means I need to create my own computation ExecutorService, which is basically a thread pool. Fortunately Schedulers has a factory to turn an Executor into a Scheduler easily.
We can create an ExecutorService with a fixed number of threads matching the optimization formula.

int threadCt = Runtime.getRuntime().availableProcessors() + 1;

ExecutorService executor = Executors.newFixedThreadPool(threadCt);

Then we can call Schedulers.from() and pass this executor to it.

int threadCt = Runtime.getRuntime().availableProcessors() + 1;

ExecutorService executor = Executors.newFixedThreadPool(threadCt);
Scheduler scheduler = Schedulers.from(executor);

Observable.range(1,1000)
.flatMap(i -> Observable.just(i)
    .subscribeOn(scheduler)
    .map(i2 -> intenseCalculation(i2))
).subscribe(System.out::println);

One thing we need to add: we need to shutdown() the ExecutorService when we are done with it so the threads are disposed and do not keep the application alive inadvertently. We can simply call this in a finallyDo() somewhere in the chain as shown below.

int threadCt = Runtime.getRuntime().availableProcessors() + 1;

ExecutorService executor = Executors.newFixedThreadPool(threadCt);
Scheduler scheduler = Schedulers.from(executor);

Observable.range(1,1000)
.flatMap(i -> Observable.just(i)
    .subscribeOn(scheduler)
    .map(i2 -> intenseCalculation(i2))
).finallyDo(() -> executor.shutdown())
.subscribe(System.out::println);

Running my process at work on a 4-core machine (so it is running 5 threads), it ran in 92 seconds, a huge difference from 225 seconds. This makes my custom ExecutorService 2.44 times faster than the Schedulers.computation()
Based on what David Karnok said, it seems RxJava maintains a conservative number of computation() threads and/or core utilization. Concluding from my observations and experiments, if you need to do some intensive parallel computation work you might be better off creating your own Scheduler off an ExecutorService with an appropriately fixed pool size.

Round-Robin Parallelization

David Karnok posted a neat trick: round-robin parallelization. If you are flatMapping to already existing Observables, then this does not really apply. But take our example where we are creating a single-emission Observables for each value.

Observable.range(1,1000)
.flatMap(i -> Observable.just(i)
    .subscribeOn(Schedulers.computation()
    .map(i2 -> intenseCalculation(i2))
).subscribe(System.out::println)

That Observable.just() factory can be expensive when used repeatedly because we are creating a new Observable for each emission. Is there not a way to consolidate 1000 single-emission Observables into a handful of Observables with several emissions?
We can do this by using groupBy() and a round-robin strategy. Create an AtomicInteger and use a modulus to break up the emissions into batches. Then you can use that modulus remainder to groupBy() on. For example, we can break up an Observable into 5 GroupedObservable items using a modulus operator, and we can flatMap() those.

final AtomicInteger batch = AtomicInteger(0);

Observable.range(1,1000)
    .groupBy(i -> batch.getAndIncrement() % 5)
    .flatMap(g -> g.observeOn(Schedulers.io())
        .map(i -> intenseCalculation(i))
    ).subscribe(System.out::println);

Note we have to use observeOn() since the GroupedObservable specifies its own subscribeOn() thread and can’t be overridden.
You can also use this in tandem with our ExecutorService strategy. Heck, you can throttle the number of GroupedObservables to be equal to the number of threads. If you have 4 cores, you would then have 5 cleanly divided streams of emissions on 5 threads.

int threadCt = Runtime.getRuntime().availableProcessors() + 1;

ExecutorService executor = Executors.newFixedThreadPool(threadCt);

Schedulers scheduler = Schedulers.from(executor);

Observable.range(1,1000)
    .groupBy(i -> batch.getAndIncrement() % threadCt )
    .flatMap(g -> g.observeOn(scheduler)
        .map(i -> intenseCalculation(i))
    ).subscribe(System.out::println);

If you have any other ideas, comments, questions, or your own discoveries with parallelization please comment below.