Rx Java 2 pre-pull next item on separate thread

Question

Scenario: I have a stream of data I am reading from the database. What I would like to do is read a chunk of data, process it and stream it using rx-java 2. But while I am processing and streaming it I would like to load the next chunk of data on a separate thread (pre-pull the next chunk).

I have tried:

Flowable.generate(...)
  .subscribeOn(Schedulers.io())
  .observeOn(Schedulers.computation())
  .map(...)
  .subscribe(...)

Unfortunately this causes the generate method to continually run on an io thread. I just want one pre-pull. I have tried using buffer, but that really just ends up creating lists of chunks.

So basically while I am streaming the current chunk on a separate thread I want to read the next chunk and have it ready.

Not sure if this is possible. I need to use generate because there is no concept of when the data will end.

I have tried using subscribe(new FlowableSubscriber(){...}) using Subscription::request but that did not seem to work.

Answer 1

There are no standard operators in RxJava that would have this type of request-response pattern. You'd need a custom observeOn that requests before it sends the current item to its downstream.

import java.util.concurrent.atomic.*;

import org.junit.Test;
import org.reactivestreams.*;

import io.reactivex.*;
import io.reactivex.Scheduler.Worker;
import io.reactivex.internal.util.BackpressureHelper;
import io.reactivex.schedulers.Schedulers;

public class LockstepObserveOnTest {

    @Test
    public void test() {
        Flowable.generate(() -> 0, (s, e) -> {
            System.out.println("Generating " + s);
            Thread.sleep(500);
            e.onNext(s);
            return s + 1;
        })
        .subscribeOn(Schedulers.io())
        .compose(new LockstepObserveOn<>(Schedulers.computation()))
        .map(v -> {
            Thread.sleep(250);
            System.out.println("Processing " + v);
            Thread.sleep(250);
            return v;
        })
        .take(50)
        .blockingSubscribe();
    }

    static final class LockstepObserveOn<T> extends Flowable<T> 
    implements FlowableTransformer<T, T> {

        final Flowable<T> source;

        final Scheduler scheduler;

        LockstepObserveOn(Scheduler scheduler) {
            this(null, scheduler);
        }

        LockstepObserveOn(Flowable<T> source, Scheduler scheduler) {
            this.source = source;
            this.scheduler = scheduler;
        }

        @Override
        protected void subscribeActual(Subscriber<? super T> subscriber) {
            source.subscribe(new LockstepObserveOnSubscriber<>(
                subscriber, scheduler.createWorker()));
        }

        @Override
        public Publisher<T> apply(Flowable<T> upstream) {
            return new LockstepObserveOn<>(upstream, scheduler);
        }

        static final class LockstepObserveOnSubscriber<T>
        implements FlowableSubscriber<T>, Subscription, Runnable {

            final Subscriber<? super T> actual;

            final Worker worker;

            final AtomicReference<T> item;

            final AtomicLong requested;

            final AtomicInteger wip;

            Subscription upstream;

            volatile boolean cancelled;

            volatile boolean done;
            Throwable error;

            long emitted;

            LockstepObserveOnSubscriber(Subscriber<? super T> actual, Worker worker) {
                this.actual = actual;
                this.worker = worker;
                this.item = new AtomicReference<>();
                this.requested = new AtomicLong();
                this.wip = new AtomicInteger();
            }

            @Override
            public void onSubscribe(Subscription s) {
                upstream = s;
                actual.onSubscribe(this);
                s.request(1);
            }

            @Override
            public void onNext(T t) {
                item.lazySet(t);
                schedule();
            }

            @Override
            public void onError(Throwable t) {
                error = t;
                done = true;
                schedule();
            }

            @Override
            public void onComplete() {
                done = true;
                schedule();
            }

            @Override
            public void request(long n) {
                BackpressureHelper.add(requested, n);
                schedule();
            }

            @Override
            public void cancel() {
                cancelled = true;
                upstream.cancel();
                worker.dispose();
                if (wip.getAndIncrement() == 0) {
                    item.lazySet(null);
                }
            }

            void schedule() {
                if (wip.getAndIncrement() == 0) {
                    worker.schedule(this);
                }
            }

            @Override
            public void run() {
                int missed = 1;
                long e = emitted;

                for (;;) {

                    long r = requested.get();

                    while (e != r) {
                        if (cancelled) {
                            item.lazySet(null);
                            return;
                        }

                        boolean d = done;
                        T v = item.get();
                        boolean empty = v == null;

                        if (d && empty) {
                            Throwable ex = error;
                            if (ex == null) {
                                actual.onComplete();
                            } else {
                                actual.onError(ex);
                            }
                            worker.dispose();
                            return;
                        }

                        if (empty) {
                            break;
                        }

                        item.lazySet(null);

                        upstream.request(1);

                        actual.onNext(v);

                        e++;
                    }

                    if (e == r) {
                        if (cancelled) {
                            item.lazySet(null);
                            return;
                        }
                        if (done && item.get() == null) {
                            Throwable ex = error;
                            if (ex == null) {
                                actual.onComplete();
                            } else {
                                actual.onError(ex);
                            }
                            worker.dispose();
                            return;
                        }
                    }

                    emitted = e;
                    missed = wip.addAndGet(-missed);
                    if (missed == 0) {
                        break;
                    }
                }
            }
        }
    }
}