Introduction

Sometimes, one has several ordered sequences of events and would like to merge them into one single flow. Since one element from a sequence should come before another element in another sequence, we need a way to keep comparing elements with each other from different sequences.

Unfortunately, zip() doesn't work because it takes a row of available items and item #2 from sequence #2 may come before item #1 from stream #3. Plus, if one stream is shorter than the others, the end sequence stops. Similarly, flatMap() doesn't work because it takes the next item from any inner source sequence the moment it is available without any ordering considerations at that point. At least it emits all items from all sources (provided there are no errors of course).

Therefore, we need something between the two operators: one that collects up a row of items from the sources, decides which is the smallest/largest of them based on some comparison logic and only emits that. It then awaits a fresh item from that specific source (or completion) and repeats the picking of the smallest/largest item as long as there are requests for it.

Such operator, let's call it orderedMerge(), has an implication about the number of its inner source sequences: it has to be fixed. The reason for it is that it has to pick the smallest/largest of the available items in order for the output to be in order. If there is still a source missing, it can't know for sure the others are smaller/larger that any of the upcoming item from that missing source will produce.

The second implication is, what happens if the sources themselves are not ordered? The logic presented in this post still works, but the end output won't be totally ordered. It will act like some priority queue instead: picking important items first before turning to less important ones.

The inner consumer

Operators handling multiple sources often need a way to prefetch item from these sources and give out them on demand to some joining logic. This mainly happens by prefetching a fixed amount, putting items in a queue, calling the parent coordinator's drain() method, and batching out replenishing calls from the coordinator if the so-called stable-prefetch backpressure is required.

For this purpose, let's see how the inner consumer, OrderedMergeInnerSubscriber, of orderedMerge() could look like:

static final class OrderedMergeInnerSubscriber<T> 
extends AtomicReference<Flow.Subscription>
implements Flow.Subscriber<T>, Flow.Subscription {

    final OrderedMergeCoordinator<T> parent;

    final int prefetch;

    final int limit;

    final Queue<T> queue;

    int consumed;

    volatile boolean done;

    OrderedMergeInnerSubscriber(
        OrderedMergeCoordinator<T> parent,
        int prefetch
    ) {
        this.parent = parent;
        this.prefetch = prefetch;
        this.limit = prefetch - (prefetch >> 2);
        this.queue = new ConcurrentLinkedQueue<>()
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        // TODO implement
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        // TODO implement
    }

    @Override
    public void onComplete() {
        // TODO implement
    }

    @Override
    public void request(long n) {
        // TODO implement
    }

    @Override
    public void cancel() {
        // TODO implement
    }
}

We'll need a reference to the coordinator of the operator, the prefetch amount that we will use the 75% of for replenishing requests, a queue - ConcurrentLinkedQueue for simplicity, but bounded SpscArrayQueue from JCTools works here as well, and a counter of how many items have been consumed so far to know when to replenish.

    @Override
    public void onSubscribe(Flow.Subscription s) {
        if (compareAndSet(null, s)) {
            s.request(prefetch);
        } else {
            s.cancel();
        }
    }

    @Override
    public void onNext(T item) {
        queue.offer(item);
        parent.drain();
    }

    @Override
    public void onError(Throwable throwable) {
        parent.onInnerError(this, throwable);
    }

    @Override
    public void onComplete() {
        done = true;
        parent.drain();
    }

    @Override
    public void request(long n) {
        int c = consumed + 1;
        if (c == limit) {
            consumed = 0;
            Flow.Subscription s = get();
            if (s != this) {
                s.request(c);
            }
        } else {
            consumed = c;
        }
    }

    @Override
    public void cancel() {
        Flow.Subscription s = getAndSet(this);
        if (s != null && s != this) {
            s.cancel();
        }
    }

I'd say there is nothing too complicated here.

• onSubscribe() saves the upstream Flow.Subscription in the AtomicReference the operator extends if not already cancelled. If successful, the prefetch amount is requested.
• onNext() stores the item in the queue and calls drain() on parent to handle that case.
• onError() defers the error signal to be handled by the parent coordinator: the parent may save up the errors or cancel the whole flow at once.
• onComplete() sets the complete indicator, which tells the parent this particular source will not produce more values and thus can be skipped when looking for the next smallest/largest items to emit
• request() will only be called by the parent to replenish one item from its perspective once the previous item has been successfully chosen as the next item to be emitted towards downstream. Since replenishing one by one is costly, we batch up those via the consumed counter. Once that counter reaches the limit (75% of prefetch), a request is issued to the upstream. Since the AtomicReference will hold itself as a cancellation indicator, we don't want to call request on ourselves. It's important to state that request() will be guaranteed to be called from one thread at a time by the virtue of the queue-drain approach the coordinator implements below.
• cancel() atomically swaps in the this as a terminal indicator and cancels the non-null, non-this Flow.Subscription if present.

The coordinator

@SafeVarargs
public static <T> Flow.Publisher<T> orderedMerge(
        Comparator<? super T> comparator, 
        int prefetch,
        Flow.Publisher<? extends T>... sources) {
    return new OrderedMergePublisher<>(sources, prefetch, comparator);
}

final class OrderedMergePublisher<T> implements Flow.Publisher<T> {

    final Flow.Publisher<? extends T>[] sources;

    final int prefetch;

    final Comparator<? super T> comparator;

    OrderedMergePublisher(
            Flow.Publisher<? extends T>[] sources,
            int prefetch,
            Comparator<? super T> comparator) {
        this.sources = sources;
        this.prefetch = prefetch;
        this.comparator = comparator;
    }

    @Override
    public void subscribe(Flow.Subscriber<? super T> s) {
         // TODO implement
    }
}

The boilerplate of writing an operator is nothing special: save up on the parameters to be used by the implementation. We allow customization via a Comparator interface. If T is self-comparable, you can use Comparators.naturalOrder() from Java itself.

The coordinator implementation has to hold onto the inner OrderedMergeInnerSubscribers for mass cancellation support, subscribe them to the sources and work out the emissions from them. Let's see the non-exciting parts of it:

static final class OrderedMergeSubscription<T>
extends AtomicInteger implements Flow.Subscription {

    final Flow.Subscriber<? super T> downstream;

    final OrderedMergeInnerSubscriber<T>[] subscribers;

    final Comparator<? super T> comparator;

    final Object[] values;

    static final Object DONE = new Object();

    Throwable error;

    boolean cancelled;

    long requested;

    long emitted;

    // -------------------------------------------------

    static final VarHandle ERROR;

    static final VarHandle DONE;

    static final VarHandle CANCELLED;

    static final VarHandle REQUESTED;

    static {
        Lookup lk = MethodHandles.lookup();
        try {
            ERROR = lk.findVarHandle(
                OrderedMergeSubscription.class, "error", Throwable.class);
            CANCELLED = lk.findVarHandle(
                OrderedMergeSubscription.class, "cancelled", boolean.class);
            REQUESTED = lk.findVarHandle(
                OrderedMergeSubscription.class, "requested", long.class);
        } catch (Throwable ex) {
            throw new InternalError(ex);
        }
    }

    OrderedMergeSubscription(
            Flow.Subscriber<? super T> downstream,
            Comparator<? super T> comparator,
            int prefetch,
            int n) {
        this.downstream = downstream;
        this.comparator = comparator;
        this.subscribers = new OrderedMergeInnerSubscriber[n];
        for (int i = 0; i < n; i++) {
             this.subscriber[i] = new OrderedMergeInnerSubscriber<>(this, prefetch);
        }
        this.values = new Object[n];
    }

    void subscribe(Flow.Publisher<? extends T>[] sources) {
        // TODO implement
    }

    @Override
    public void request(long n) {
        // TODO implement
    }

    @Override
    public void cancel() {
        // TODO implement
    }

    void drain() {
        // TODO implement
    }

    void onInnerError(OrderedMergeInnerSubscriber<T> sender, Throwable ex) {
        // TODO implement
    }

    void updateError(Throwable ex) {
        // TODO implement
    }
}

We have a couple of fields and methods here, some should be familiar in its naming and intended purpose:

• downstream will receive the ordered sequence of items
• subscribers holds onto the fixed set of inner OrderedMergeInnerSubscribers, each will be subscribed to a particular Flow.Publisher and the total number of them won't ever change in this operator.
• comparator will compare elements from various sources
• values holds onto the next available value from each source. This allows the merging algorithm to work with queues that don't support peek() (such as RxJava 2's on queue implementations) and otherwise has nice properties such as locality, avoiding accessing internals of the inner subscribers' queues and the overhead of a peek()-poll() pair all the time.
• The DONE constant will indicate a particular source has no further elements and can be ignored (without looking at its subscriber).
• error will gather the errors signalled by the sources and emitted together once all sources terminated. There is an ERRORVarHandle for concurrent access to this field.
• cancelled indicates the downstream has issued a cancel() call to stop the flow. The CANCELLEDVarHandle will allow us to use compareAndSet() to cancel at most once.
• requested accumulates the requests done by the downstream via its REQUESTEDVarHandle.
• emitted counts how many items were emitted and will be compared against requested to detect when to pause emitting.

There is no separate done indicator field because we will deduce this state by detecting that all values items are marked as DONE.

Now let's see the shorter methods implemented:

    // ...

    void subscribe(Flow.Publisher<? extends T>[] sources) {
        for (int i = 0; i < sources.length; i++) {
            sources[i].subscribe(subscribers[i]);
        }
    }

    @Override
    public void request(long n) {
        if (n <= 0L) {
            updateError(new IllegalArgumentException("non-negative request expected"));
        } else {
            for (;;) {
                long current = (long)REQUESTED.getAcquire(this);
                long next = current + n;
                if (next < 0L) {
                    next = Long.MAX_VALUE;
                }
                if (REQUESTED.compareAndSet(this, current, next)) {
                    break;
                }
            }
        }
        drain();
    }

    @Override
    public void cancel() {
        if (CANCELLED.compareAndSet(this, false, true)) {
            for (OrderedMergeInnerSubscriber<T> inner : subscribers) {
                inner.cancel();
            }

            if (getAndIncrement() == 0) {
                Arrays.fill(values, null);

                for (OrderedMergeInnerSubscriber<T> inner : subscribers) {
                    inner.queue.clear();
                }
            }
        }
    }

    void onInnerError(OrderedMergeInnerSubscriber<T> sender, Throwable ex) {
        update(ex);
        sender.done = true;
        drain();
    }

    void updateError(Throwable ex) {
        for (;;) {
            Throwable current = (Throwable)ERROR.getAcquire(this);
            Throwable next;
            if (current == null) {
                next = throwable;
            } else {
                next = new Throwable();
                next.addSuppressed(current);
                next.addSuppressed(throwable);
            }
            if (ERROR.compareAndSet(this, current, next)) {
                break;
            }
        }
    }

    void drain() {
        // TODO implement
    }

}

The subscribe() method simply subscribes to all sources with the prepared array of OrderedMergeInnerSubscribers. The cancel() method cancels all inner subscribers and then enters a half-open drain mode where both the values array and each queue of the inner subscribers is cleared in order to help the GC. Both request() and updateError() should be familiar from the previous post of the series.

What's left is the drain() logic itself.

void drain() {
    if (getAndIncrement() != 0) {
        return;
    }

    int missed = 1;
    Flow.Subscriber<? super T> downstream = this.downstream;

    Comparator<? super T> comparator = comparator;

    OrderedMergeInnerSubscriber<T>[] subscribers = this.subscribers;
    int n = subscribers.length;

    Object[] values = this.values;

    long e = emitted;

    for (;;) {
         long r = (long)REQUESTED.getAcquire(this);

         for (;;) {
              // TODO implement
         }

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

We start out with the usual drain-exclusion logic: transitioning the work-in-progress counter of this (by extending AtomicInteger) from zero to one allows one thread to enter and perform the emissions. We load the frequently accessed components into local fields and do an almost-classical for-loop with missed accounting to determine when to leave the loop.

Note that in the loop, after reading the current request amount we don't have the usual while (e != r) and if (e == r) cases. The reason for this is that we can have one shared loop for both the cases when backpressure is applied and when there are no further source items to merge and can terminate the sequence without a request from downstream.

// inner for(;;) {

if ((boolean)CANCELLED.getAcquire(this)) {
    Arrays.fill(values, null);

    for (OrderedMergeInnerSubscriber<T> inner : subscribers) {
        inner.queue.clear();
    }
    return;
}

int done = 0;
int nonEmpty = 0;
for (int i = 0; i < n; i++) {
    Object o = values[i]
    if (o == DONE) {
        done++;
        nonEmpty++;
    } else
    if (o == null) {
        boolean innerDone = subscribers[i].done;
        o = subscribers[i].queue.poll();
        if (o != null) {
            values[i] = o;
            nonEmpty++;
        } else if (innerDone) {
            values[i] = DONE;
            done++;
            nonEmpty++;
        }
    } else {
        nonEmpty++;
    }
}

The first part is check if there was a cancellation from downstream. If so, we clear the internal state of the coordinator and each queue then quit. Next for each subscriber of a source, we have to poll the next available item into the common values array if there is not already an item available there. In addition, we account how many of those items indicate a completed source and how many of them has actual items.

Note that checking a source for completion has to happen before polling for the next item from its queue. As explained before on this blog, since Flow.Subscriber methods are invoked in a strict protocol order where an onComplete always happens after any previous onNext calls, if we detect done to be true and then get a null from the queue, we know there can't be any further items from that source. Otherwise, polling and seeing an empty queue first then checking done opens a window when the source quickly produces items and completes between these two checks.

Next, we handle the overall state of the operator:

if (done == n) {
    Throwable ex = (Throwable)ERROR.getAcquire(this);
    if (ex == null) {
        downstream.onComplete();
    } else {
        downstream.onError(ex);
    }
    return;
}

if (nonEmpty != n || e != r) {
    break;
}

If all of the elements turn out to be DONE, that means we exhausted all sources and can terminate the downstream accordingly (considering if there was any error along the way). If not all values slot have something or the downstream is not ready to receive an item, we break out this inner look and the outer loop will see if more work has to be done or not.

Finally, we find the smallest item from the available values:

    T min = null;
    int minIndex = -1;

    int i = 0;
    for (Object o : values) {
        if (o != DONE) {
            if (min == null || comparator.compare(min, (T)o) > 0) {
                min = (T)o;
                minIndex = i;
            }      
        }
        i++;
    }

    values[minIndex] = null;

    downstream.onNext(min);

    e++;
    subscribers[minIndex].request(1);

} // of the inner for (;;)

Once we know there we can emit an item, we'll find the smallest one among the non-DONE entries along with its index. Since we checked that not all entries are DONE before, min must end up non-null and minIndex non-negative. We clear the appropriate values entry indicating the next cycle should poll for more items from that particular source, we emit the found minimum item, increment the emission counter and signal the winning source to produce one more item.

Conclusion

The orderedMerge() operator shown in this post is perhaps one of the shortest and more comprehensible among the other ones, even if considering the lack of infrastructure with Java 9 Flows (i.e., cancelled indicator). The queue-drain approach present in many of the typical reactive operators can be observed here as well.

Since the operator collects you a row of values available, it can be relatively easily turned into a zip() operator:

• done != 0 indicates one or more sources run out of items thus the sequence can be completed. Note that the non-done inner subscribers have to be cancelled and cleared before the downstream gets the terminal event.
• instead of the loop that compares items, one copies the values array, clears the original one, applies a function to the copy and emit the result of that function call.

You can also turn it into a (rather clumsy) merge() operator that uses a round-robin collection strategy to pick the item to be emitted: an index (also saved into a field for subsequent drain rounds) that indicates which next slot to consider for emission if nonEmpty != 0, skipping over the DONE entries along the way.

However, there is one likely problem that troubles such value-collecting-and-emitting operators; it does - what Stephane Maldini, Project Reactor lead, once called our RxJava 2 / Reactor-Core 3 algorithms do - thread-stealing. Given all the sources, one of them will be doing the collecting and emitting the smallest item for all the other sources while that thread itself likely won't make any progress unless it finds a small pause in the onslaught of source items so the drain loop can quit.

This may be undesirable at times and there is a solution for it. To get there, we will investigate how thread switching in mid-flow can be implemented within the Java 9 Flow API in the next post.

Introduction

Writing imperative-looking code with Kotlin Coroutines is certainly an attractive property of it, but I'd think things can get quite convoluted pretty fast once, for example, Selectors are involved.

I haven't gotten there to look at what Selectors are, I only read that they can help you implement a flatMap like stream combiner. We are not goind to do that now, RxJava can do it for us after all.

However, the reasonable question arises: if I have a coroutine generator, a coroutine transformation or simply want to receive items from a Flowable, how can I make RxJava work with these coroutines?

Easily with the combined magic of Kotlin Coroutines and RxJava coroutines!

Suspendable Emitter

A generator is a source-like construct that emits items followed by a terminal signal. It should be familiar from RxJava as the Flowable.generate() operator. It gives you a FlowableEmitter and the usual onNext, onError and onComplete calls on it.

One limitation is that you can call onNext only once per invocation of your (Bi)Consumer lambda that receives the emitter. The reason is that we can't block a second call to onNext and we don't want to buffer it either; therefore, RxJava cooperates with the developer.

Compiler supported suspension and state machine built by it, however, allow us to prevent a second call from getting through by suspending it until there is a demand from the downstream, which then resumes the coroutine where it left off. Therefore, we can lift the single onNext requirement for our Coroutine-based generator.

So let's define the SuspendEmitter interface

interface SuspendEmitter<in T> : CoroutineScope {

    suspend fun onNext(t: T)

    suspend fun onError(t: Throwable)

    suspend fun onComplete()
}

By extending the CoroutineScope, we provide useful infrastructure (i.e., coroutineContext, isActive) to the block that will target our SuspendEmitter. One can argue that why use onError and onComplete since a coroutine can throw and simply end. The reason is that this way, a coroutine can terminate the sequence from a transformation we'll see later, just like our recent mapFilter operator allows it.

The flow-producer

Given our context providing interface for a generator coroutine, let's define the generator method the user will call:

fun <T> produceFlow(generator: suspend SuspendEmitter.() -> Unit) : Flowable<T> {
    return Produce(generator)
}

(For those unfamiliar with the Kotlin syntax, the SuspendEmitter.() -> Unit is practically a one parameter lambda of signature (param: SuspendEmitter) -> Unit where, when the lambda is implemented, accessing methods of param do not need to be qualified by it, thus you can write onNext(1) instead of param.onNext(1).)

We have to implement a Flowable that interacts with a suspendable generator function in some fashion. When implementing source-like operators, one usually has to write a Subscription instance and call Subscriber.onSubscribe() with it.

class Produce<T>(private val generator: suspend SuspendEmitter<T>.() -> Unit) : 
        Flowable<T>() {
    override fun subscribeActual(s: Subscriber<in T>) {
        launch(Unconfined) {
            val parent = ProduceSubscription(s)
            parent.setJob(coroutineContext[Job])
            s.onSubscribe(parent)
            generator(parent)
        }
    }
}

Since the generator is a suspendable coroutine, we need a context where it can run. The Unconfined context gives us a trampolined execution environment where resumptions of suspended coroutines are not confined to any particular thread, as if you'd run with the trampoline()Scheduler in RxJava.

We create our Subscription, attach the Job of the coroutine context itself to bridge the cancellation from a downstream Subscription.cancel(), signal the custom Subscription to the downstream and then execute the provided producer block by supplying it the parent which also implements SuspendEmitter.

So far, nothing is too hairy or convoluted, however, the interaction between regular trampolined coroutines of RxJava and the Kotlin Coroutine infrastructure is more involved.

Non-blocking await/notify

We will need a way to get the generator coroutine suspended if there are no downstream requests and we have to resume that coroutine when the downstream does request an amount. This resembles the wait-notify pair of a typical BlockingQueue implementation where a blocking emission due to a full queue gets unblocked by a notification by a concurrent take()/poll() invocation. Since we don't want to block and the coroutine infrastructure supports programmatic resuming of a coroutine, we'll use this feature in two helper methods establishing a non-blocking wait-notify exchange:

typealias Cont = Continuation<Unit>

fun notify(ref: AtomicReference<Cont?>) {
    while (true) {
        val cont = ref.get()
        val next : Cont?
        if (cont != null && cont != TOKEN) {
            if (ref.compareAndSet(cont, null)) {
                cont.resume(Unit)
                break
            }
        } else {
            if (ref.compareAndSet(cont, TOKEN)) {
                break;
            }
        }
    }
}

We will use a valueless Continuation<Unit>, Cont for short, and atomics to place an indicator or an actual continuation object in an AtomicReference. The notify() atomically performs the following logic: if there is a real continuation in the reference, we clear it and then call resume on it to trigger the resumption. Otherwise, we set it to the shared TOKEN object indicating that when the other side, await, wanted to get continued, it can do so immediately.

fun await(ref: AtomicReference<Cont?>, cont: Cont) {
    while (true) {
        val a = ref.get()
        if (a == TOKEN) {
            if (ref.compareAndSet(a, null)) {
                cont.resume(Unit)
                break
            }
        } else {
            if (ref.compareAndSet(a, cont)) {
                break;
            }
        }

    }
}

The await() method uses the same reference and the continuation instance provided by a suspendCoroutine in its code block.The method atomically checks if there is a TOKEN and if so, it calls resume on the continuation parameter after clearing the TOKEN from the reference. Otherwise, it stores the continuation in the reference and quits.

val TOKEN: Cont = object: Cont {
    override val context: CoroutineContext
        get() = throw UnsupportedOperationException()

    override fun resume(value: Unit) {
        throw UnsupportedOperationException()
    }

    override fun resumeWithException(exception: Throwable) {
        throw UnsupportedOperationException()
    }

}

Finally, the TOKEN is just an empty implementation of a Continuation - we should never call its methods as the object reference itself serves only a purpose of indicator for an immediate resumption.

The ProduceSubscription  

Now we can implement the ProduceSubscription class. First, let's see the skeleton with the relevant fields:

open class ProduceSubscription<T>(
        private val actual: Subscriber<in T>,
        private val ctx : CoroutineContext
) : Subscription, SuspendEmitter<T> {

    companion object {
        val CANCELLED = Object()
    }

    @Suppress("DEPRECATION")
    override val context: CoroutineContext
        get() = ctx!!

    override val isActive: Boolean
        get() = job.get() != CANCELLED

    private val job = AtomicReference<Any>()

    private val requested = AtomicLong()

    private val resume = AtomicReference<Cont?>()

    private var done: Boolean = false

    override suspend fun onNext(t: T) {
        // TODO implement
    }

    override suspend fun onError(t: Throwable) {
        // TODO implement
    }

    override suspend fun onComplete() {
        // TODO implement
    }

    override fun cancel() {
        // TODO implement
    }

    override fun request(n: Long) {
        // TODO implement
    }

    fun setJob(j: Job?) {
        // TODO implement
    }
}

We see the methods of both Subscription and SuspendEmitter along with a couple of fields/properties:

• It takes the downstream's Subscriber and the CoroutineContext it will provide to the produce callback in the operator.
• We will use the companion object's CANCELLED value to indicate the the parent job we get from the coroutineContext is cancelled exactly once.
• It considers being active when the job object is not the CANCELLED indicator
• Of which Job is then stored in the jobAtomicReference.
• We have to track the requested amount from downstream via an AtomicLong.
• The resumeAtomicReference stores the continuation to be used with the non-blocking await-notify shown in the previous section.
• Finally, we have the done flag indicating the generator coroutine called onError or onComplete at most once.

    override suspend fun onNext(t: T) {
        if (job.get() == CANCELLED) {
            suspendCoroutine { }
        }
        val r = requested.get()
        if (r == 0L) {
            suspendCoroutine { cont -> await(resume, cont)  }
        }

        actual.onNext(t)

        if (job.get() == CANCELLED) {
            suspendCoroutine { }
        }
        if (resume.get() == TOKEN) {
            resume.compareAndSet(TOKEN, null)
        }
        if (r != Long.MAX_VALUE) {
            requested.decrementAndGet()
        }
    }

First we check if the downstream has cancelled the generator in which case we should get out of the coroutine entirely. I'm not sure if there is a more appropriate way for doing this other than suspending indefinely.

Next, we check the request amount and if it is zero, we suspend the current coroutine by using our non-blocking await mechanism. Once notified, or there was at least one requested item, the code should continue with the emission of the item. This could trigger an in-sequence cancellation and we suspend the coroutine indefinitely again.

Since the downstream can immediately request some amount due to the s.onSubscribe(parent) call in the operator, before the generator can even run and call onNext, we may have a TOKEN in the resume field, that would otherwise incorrectly indicate the next call to await it can resume immediately, violating the backpressure we expect. I know this sounds convoluted, but I learned it the hard way...

Finally, we decrement the request amount if not unbounded.

The onError and onComplete look pretty much alike:

    override suspend fun onError(t: Throwable) {
        if (!done) {
            done = true
            actual.onError(t)
            cancel()
        }
        suspendCoroutine { }
    }

    override suspend fun onComplete() {
        if (!done) {
            done = true
            actual.onComplete()
            cancel()
        }
        suspendCoroutine { }
    }

We set the done flag to true, emit the relevant event to the downstream and then cancel the job/Subscription we are running with. I defensively suspend the coroutine afterwards.

Next we see how cancel() and setJob() works:

    override fun cancel() {
        val o = job.getAndSet(CANCELLED)
        if (o != CANCELLED) {
            (o as Job).cancel()
        }
    }

    fun setJob(j: Job?) {
        while (true) {
            val o = job.get()
            if (o == CANCELLED) {
                j?.cancel()
                break
            }
            if (job.compareAndSet(o, j)) {
                break
            }
        }
    }

They are pretty much implemented along RxJava's typical deferred cancellation mechanism. cancel() atomically swaps in the CANCELLED indicator and calls cancel on the Job it contained. setJob() atomically set the Job instance or cancels it if cancel() swapped in the CANCELLED indicator just before that.

Lastly, the request() implementation that is responsible for accounting downstream requests and resuming the suspended generator if inside onNext().

    override fun request(n: Long) {
        if (BackpressureHelper.add(requested, n) == 0L) {
            notify(resume)
        }
    }

In the RxJava world, a transition from 0 to n triggers the emission loop in a range() operator for example. Here, we notify a possibly suspended coroutine that will resume from the await() method we implemented.

Testing it is simple with RxJava:

val f = produceFlow {
    for (i in 0 until 10) {
         println("Generating $i")
         onNext(i)
    }
    onComplete()
}

f.test(0)
.assertEmpty()
.requestMore(5)
.assertValues(0, 1, 2, 3, 4)
.requestMore(5)
.assertResult(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)

Outstanding!

The flow-transformer

Now that we have a way to emit items, we would like to emit an item in response to an upstream value, like the map() operator but with a suspendable coroutine function. RxJava's map is confined to return one item in exchange for one upstream item.

With coroutines and the ProduceSubscription described in the previous section, we could emit any number of items without overflowing a Subscriber!

Let's define our API and a skeleton implementation for it first:

fun <T, R> Flowable<T>.transform(
        transformer: suspend SuspendEmitter.(T) -> Unit)
 : Flowable<R> {
    return Transform(this, transformer)
}

class Transform<T, R>(
        private val source: Flowable<T>, 
        private val transformer: suspend SuspendEmitter<R>.(T) -> Unit)
 : Flowable() {
    override fun subscribeActual(s: Subscriber) {
        // TODO implement
    }
}

We define a transform extension method on Flowable with a suspendable transformer that takes our SuspendEmitter, the upstream's value and returns nothing.

This time, we have an upstream we have to subscribe to via a regular FlowableSubscriber from RxJava, call the coroutine in some way and make sure we keep calling the upstream for more values as we have to deal with the backpressure of the coroutine itself transitively.

The first step into this direction is the handling of the upstream's own Subscription we get through Subscriber.onSubscribe. We have to attach that to the Subscription we show to the downstream Subscriber. Since we will use the ProduceSubscription anyway, we extend it and override its cancel() for this purpose:

class ProduceWithResource<T>(
        actual: Subscriber<in T>,
        ctx : CoroutineContext
) : ProduceSubscription<T>(actual, ctx) {
    private val resource = AtomicReference<Subscription>()

    fun setResource(s: Subscription) {
        SubscriptionHelper.replace(resource, s)
    }

    override fun cancel() {
        SubscriptionHelper.cancel(resource)
        super.cancel()
    }
}

We simply use the deferred cancellation helper for Subscriptions.

Now let's see how we can prepare the context for running the coroutine inside the transform operator's subscribeActual() method:

    val ctx = newCoroutineContext(Unconfined)
    val parent = ProduceWithResource(s, ctx)
    s.onSubscribe(parent)
    source.subscribe(object: FlowableSubscriber {

        var upstream : Subscription? = null

        val wip = AtomicInteger()
        var error: Throwable? = null

        override fun onSubscribe(s: Subscription) {
            // TODO implement
        }

        override fun onNext(t: T) {
            // TODO implement
        }

        override fun onError(t: Throwable) {
            // TODO implement
        }

        override fun onComplete() {
            // TODO implement
        }
    })

First we create an unconfinded context where each invocation of the transformer coroutine will execute and suspend in. We create the producer that can hold an additional Subscription and send it to the downstream Subscriber. Finally, we subscribe to the upstream with a FlowableSubscriber.

In this custom FlowableSubscriber, we will have request from upstream, thus we save the Subscription we'll get from it. The wip and error fields will be used to achieve something similar to a half-serialization. I'll explain it once the methods are implemented.

Handling onSubscribe() is straightforward and typical for an RxJava operator:

    override fun onSubscribe(s: Subscription) {
        upstream = s
        parent.setResource(s)
        s.request(1)
    }

We store the upstream's subscription locally and in the ProducerWithResource to link up the cancellation across the operator. Then we request one item; this is partly due to simplifying the interaction between a suspended coroutine and the upstream producer. Using larger prefetch would require the use of some intermediate queue - possible, but left for the reader as an exercise. (Finally, we found a use for request(1)!)

Next, onNext():

    override fun onNext(t: T) {
        launch(ctx) {
           parent.setJob(coroutineContext[Job])

           wip.getAndIncrement()

           transformer(parent, t)

           if (wip.decrementAndGet() == 0) {
               upstream!!.request(1)
           } else {
               val ex = error;
               if (ex == null) {
                   s.onComplete()
               } else {
                   s.onError(ex)
               }
               parent.cancel()
           }
       }
    }

First, the Job of the actual coroutineContext has to be stored so a downstream cancellation can can call its Job.cancel() method. We have to do this because we will go in and out of the launch() when the upstream sends an item.

Next, the wip counter is incremented, which may seem odd. The reason for this is that if the transformer coroutine gets suspended, the execution returns to the caller of onNext(), a regular RxJava producer of some sorts. If this producer has reached its end, it will call onError or onComplete as these can be issued without request. As we'll see a bit later, forwarding these signals cuts out any pending emission from the suspended coroutine, therefore, we use the pattern of a half-serializer to save this terminal indication.

The transformer is executed with the parent ProducerWithResource instance that handles the suspendable onNext emissions towards the downstream.

Once the transformer's job has been done, the execution (resumes) with the atomic decrement of the wip counter. If it successfully decrements to 0, there was no terminal event signalled from the upstream while the transformer was suspended, thus we can request the next item from the upstream RxJava source.

The onError and onComplete are much simpler fortunately:

    override fun onError(t: Throwable) {
        error = t
        if (wip.getAndIncrement() == 0) {
            s.onError(t)
            parent.cancel()
        }
    }

    override fun onComplete() {
        if (wip.getAndIncrement() == 0) {
            s.onComplete()
            parent.cancel()
        }
    }

We store the Throwable (in onError only), then atomically increment the wip counter. If there was no ongoing coroutine, we are safe to emit the terminal event and cleanup/cancel the contextual Job we may still be referencing. If the original wip value was 1, the increment bumps it to 2 and the decrement in onNext() will detect the terminal condition and act accordingly.

Let's test it (by reusing the generator for fun)!

    f.transform({
        if (it % 2 == 0) {
            onNext(it)
        }
    })
    .test()
    .assertResult(0, 2, 4, 6, 8)

    f.transform({
        onNext(it)
        onNext(it + 1)
    })
    .test()
    .assertResult(0, 1, 1, 2, 2, 3, 3, 4, 4,
            5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10)

    f.transform({
        launch(CommonPool) {
            onNext(it + 1)
        }
    })
    .test()
    .awaitDone(5, TimeUnit.SECONDS)
    .assertResult(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)

We can filter or amplify a source, synchronously or asynchronously if necessary with a single operator! Excellent!

The receiver

The last operation we'd do is, given a Flowable flow, we'd like to return to the coroutine world and consume the flow. For that, a ReceiverChannel seems to be appropriate output type as it can be for-each looped nicely.

Let's define the extension method toReceiver() with a skeleton as well:

suspend fun  Flowable<T>.toReceiver(capacityHint: Int = 128) : ReceiveChannel<T> {
    val queue = Channel<T>(capacityHint)

    val upstream = AtomicReference<Subscription>()
    val error = AtomicReference<Throwable>()
    val wip = AtomicInteger()

    subscribe(object: FlowableSubscriber<T> {

        override fun onSubscribe(s: Subscription) {
            // TODO implement
        }

        override fun onNext(t: T) {
            // TODO implement
        }

        override fun onComplete() {
            // TODO implement
        }

        override fun onError(t: Throwable) {
            // TODO implement
        }

    })

    return // TODO implement
}

First, a Channel of type T and the given capacity is created. It is followed by the AtomicReference that will hold the source Flowable's Subscription, which will have to be linked up with the consumer to propagate cancellation. Next, since the upstream may signal terminal events while the channel is suspended in a send() we'll use - similar to the ProducerWithResource.onNext() situation, we will use the same AtomicInteger-based technique. The error AtomicReference will serve as the intermediary when handing over the terminal event to the channel.

Let's see the FlowableSubscriber implementation first:

        override fun onSubscribe(s: Subscription) {
            if (SubscriptionHelper.setOnce(upstream, s)) {
                s.request(1)
            }
        }

        override fun onNext(t: T) {
            launch (Unconfined) {
                wip.getAndIncrement()

                queue.send(t);

                if (wip.decrementAndGet() == 0) {
                    upstream.get().request(1)
                } else {
                    queue.cancel(error.get());
                }
            }
        }

        override fun onComplete() {
            if (wip.getAndIncrement() == 0) {
                launch(Unconfined) {
                    queue.cancel();
                }
            }
        }

        override fun onError(t: Throwable) {
            error.lazySet(t)
            if (wip.getAndIncrement() == 0) {
                launch(Unconfined) {
                    queue.cancel(t);
                }
            }
        }

The FlowableSubscriber implementation, practically, performs the same bookeeping as the transformer() operator did, with the exception that the closing of the channel has to happen in a launch-provided context.

However, this is only the producer half of the channel, we still need the consumer part, more specifically, the consumer-reemitter. Luckily, the build in produce() operator of the Coroutines library help with it. Why not return the channel directly? Because we need a way to detect if the channel is closed from the consumer's end and Channel doesn't allow us to register a completion handler for it. However, the Job inside the coroutineContext of produce() does:

    return produce(Unconfined) {
        coroutineContext[Job]?.invokeOnCompletion { 
            SubscriptionHelper.cancel(upstream) 
        }

        for (v in queue) send(v)
    }

Let's test this last operator:

runBlocking {
    for (i in f.toReceiver()) {
         println(i)
    }
    println("Done")

    for (i in f.subscribeOn(Schedulers.single()).toReceiver()) {
         println("Async $i")
    }
    println("Async Done")
}

Well done!

Conclusion

In this blog post, I demonstrated how one can write three operators, produceFlow, transform and toReceiver, that can interoperate with RxJava's own, backpressure enabled Flowable type reasonably well.

This should prove that both technologies, at the end, can be combined by the developer as seen fit for the target domain or business requirements.

This was somewhat a heated week for me so for now, until something interesting comes up in this topic, me writing about Kotlin Coroutines will be ... suspended.

Introduction

Ensuring certain computations happen on the right thread, usually off the main thread, is a very common development task when dealing with reactive flows. When building up tools for Java 9's Flow API, one can decide to add this thread-switching support to each operator directly - see the range() operator from the start of the series -, or have a standalone stage for this purpose.

This is a tradeoff. Inlining thread switching avoids bogging down the source thread like the thread-stealing behavior of most of the queue-drain approach presented so far. A separate operator allows better composition and may even allow working with exotic asynchrony-providing components.

The observeOn operator

In Java, threading support is provided via the Executor, ExecutorService and ScheduledExecutorService-based API. Executor is is the most basic one of them which only provides a single execute(Runnable) method. This allows creating an Executor from a lambda:

Executor trampoline = Runnable::run;

Executor swing = SwingUtilities::invokeLater;

Executor pool = ForkJoinPool.commonPool();

As the least common denominator, we'll use Executor in defining our observeOn operator:

public static <T> Flow.Publisher<T> observeOn(
        Flow.Publisher<T> source, 
        Executor exec,
        int prefetch) {
    return new ObserveOnPublisher<>(source, exec, prefetch);
}

Crossing an asynchronous boundary requires the temporary storage of an event until the other side can pick it up. The queue-drain approach can provide a nice bounded queue we can size with prefetch. In addition, the so-called stable-prefetch request management (shown in the mapFilter operator before) allows minimizing the overhead of requesting more items.

First, let's see the skeleton of the operator's main Flow.Subscriber implementation:

static final class ObserveOnSubscriber<T> implements
Flow.Subscriber<T>, Flow.Subscription, Runnable {

    final Flow.Subscriber<? super T> downstream;

    final Executor exec;

    final int prefetch;

    final Queue<T> queue;

    Flow.Subscription upstream;

    int wip;
    static final VarHandle WIP = 
        VH.find(MethodHandles.lookup(), ObserveOnSubscriber.class, 
"wip", int.class);

    long requested;
    static final VarHandle REQUESTED =
        VH.find(MethodHandles.lookup(), ObserveOnSubscriber.class, 
"requested", long.class);

    boolean done;
    static final VarHandle DONE =
        VH.find(MethodHandles.lookup(), ObserveOnSubscriber.class, 
"done", boolean.class);

    boolean cancelled;
    static final VarHandle CANCELLED =
        VH.find(MethodHandles.lookup(), ObserveOnSubscriber.class, 
"cancelled", boolean.class);

    Throwable error;

    long emitted;

    int consumed;

    ObserveOnSubscriber(
            Flow.Subscriber<? super T> downstream,
            Executor exec,
            int prefetch
    ) {
        this.downstream = downstream;
        this.exec = exec;
        this.prefetch = prefetch;
        this.queue = new SpscArrayQueue<>(prefetch);
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        // TODO implement
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        // TODO implement
    }

    @Override
    public void onComplete() {
        // TODO implement
    }

    @Override
    public void request(long n) {
        // TODO implement
    }

    @Override
    public void cancel() {
        // TODO implement
    }

    void schedule() {
        // TODO implement
    }

    @Override
    public void run() {
        // TODO implement
    }
}

The SpscArrayQueue is available from the JCTools library. VH is a shortand utility class for getting a VarHandle for a particular field:

public final class VH {
    public static VarHandle find(
            MethodHandles.Lookup lookup, 
            Class clazz,
            String field,
            Class type) {
        try {
            return lookup.findVarHandle(clazz, field, type);
        } catch (Throwable ex) {
            throw new InternalError(ex);
        }
    }
}

If you are using IntelliJ IDEA, the latest version has nice parameter matching support for findVarHandle and highlights it if the field or type doesn't match the target class' member definition. Unfortunately, the method throws a checked exception which requires wrapping it in try-catch unavailable at field initialization time. By factoring the logic out, we lose that support but gain some other convenience. Needless to say, if the field names or types are off, we'll get a nice InternalError during unit tests anyway. Note that the MethodHandles.Lookup instance has to be provided because otherwise the find method could not access the fields of the target class due to visibility restrictions enforced by the JVM otherwise.

Now let's describe the fields quickly:

• downstream comes from the parent Flow.Publisher's subscribe method as usual.
• exec is the Executor the operator will use
• prefetch defines the number of items to request from the upstream when the connection is established. A specific proportion (75%) will be used for replenishing items.
• queue will hold up to the given prefetch amount of items. Since there is only one thread calling onNext and one thread draining the queue, we'll use a bounded single-producer single-consumer queue.
• upstream is received through onSubscribe from the upstream source that allows requesting and cancelling the flow.
• wip will ensure there is only one drain running at a time, executing on the given Executor. This is the same trampolining logic as with other concurrent operators: a transition from 0 to 1 will start the drain process and any further increments to wip will indicate additional work has to be done by the drain process.
• requested will track the items requested by the downstream as the requesting is decoupled by the operator. The reason for this is that the operator uses a bounded buffer and there is no way to expect the downstream to request only as much as the buffer can hold at a given moment.
• done indicates the upstream has finished emitting items. Calling onComplete on the downstream immediately is not an option as there could be items still queued up.
• cancelled indicates the downstream issued a cancel() call and both the drain loop and the upstream has to stop producing events.
• error holds the Throwable from the upstream. We'll only read it if done == true ensuring proper cross-thread visibility.
• emitted counts how many items have been emitted towards the downstream. Decrementing the requested field is possible but generally expensive due to unavoidable atomics (and may not auto-batch decrements at all, thus each item incurs the cost).
• consumed counts how many items have been taken from the queue. It can be reset to zero when a certain limit is reached (75% of prefetch for example) and issue a request() towards the upstream, asking for a fixed number of items periodically.

Implementing the Flow.Subscriber methods is straightforward as they have little to do; store values and try starting the drain process:

    @Override
    public void onSubscribe(Flow.Subscription s) {
        upstream = s;
        downstream.onSubscribe(this);
        s.request(prefetch);
    }

    @Override
    public void onNext(T item) {
        queue.offer(item);
        schedule();
    }

    @Override
    public void onError(Throwable throwable) {
        error = throwable;
        DONE.setRelease(this, true);
        schedule();
    }

    @Override
    public void onComplete() {
        DONE.setRelease(this, true);
        schedule();
    }

The implementation of Flow.Subscription part is similarly not that complicated. The handling of non-positive request amount is left out for brevity. We have to trigger the draining logic when the downstream requests in case the upstream has items queued up already due to possible speed difference.

    @Override
    public void request(long n) {
        for (;;) {
           long a = (long)REQUESTED.getAcquire(this);
           long b = a + b;
           if (b < 0L) {
               b = Long.MAX_VALUE;
           }
           if (REQUESTED.compareAndSet(this, a, b)) {
               schedule();
               break;
           }
    }

    @Override
    public void cancel() {
        if (CANCELLED.compareAndSet(this, false, true)) {
            upstream.cancel();
            if ((int)WIP.getAndAdd(this, 1) == 0) {
                queue.clear();
            }
        }
    }

The request() uses the typical, capped atomic addition of the requested amount. The cancel() method sets the cancelled flag once, cancels the upstream and if there is no draining happening at the moment, clears the queue to help the GC.

The schedule() method conceptually works like any drain() method we've implemented so far. The difference is that the drain loop inside it has to be run by the Executor (potentially) on another thread:

    void schedule() {
        if ((int)WIP.getAndAdd(this, 1) == 0) {
            exec.execute(this);
        }
    }

By implementing Runnable directly, it saves us creating the drain task all the time; here, all the state that need to be communicated between the caller of schedule() and the run() method is done through fields. Since the transition from 0 to 1 and later on, N back to 0 happens on a single thread, the underlying Executor can be a pool of any threads; this type of trampolining will make the Executor to pick one thread from this pool and while the drain lasts, prevent any other thread from the same pool to start another drain run (which would violate the Flow protocol and/or reorder events).

Finally, let's see the implementation of run() that holds the rest of a typical drain-loop:

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

        Flow.Subscriber<? super T> a = downstream;
        Queue<T> q = queue;
        long e = emitted;
        int c = consumed;
        int limit = prefetch - (prefetch >> 2);

        for (;;) {

             long r = (long)REQUESTED.getAcquire(this);

             while (e != r) {
                 // TODO implement
             }

             if (e == r) {
                 // TODO implement
             }

             emitted = e;
             consumed = c;
             missed = (int)WIP.getAndAdd(this, -missed) - missed;
             if (missed == 0) {
                 break;
             }
        }
    }

The pattern for the drain loop is pretty standard: the missing counter will detect if there is more work to be done, we load fields into local variables to avoid fetching them from last-level-cache due to all the atomics around. Lastly, we have the usual while loop that keeps running until we run out of requests or upstream events, then checking if by reaching a terminal state in the operator, the terminal events can be emitted without actual requests or not.

There is nothing unusual at this point in the inner while loop and the if statements:

    while (e != r) {
        if ((boolean)CANCELLED.getAcquire(this)) {
            q.clear();
            return;
        }

        boolean d = (boolean)DONE.getAcquire(this);
        T v = q.poll();
        boolean empty = v == null;

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

        if (empty) {
            break;
        }

        a.onNext(v);

        e++;
        if (++c == limit) {
            c = 0;
            upstream.request(limit);
        }
    }

    if (e == r) {
        if ((boolean)CANCELLED.getAcquire(this)) {
            q.clear();
            return;
        }

        boolean d = (boolean)DONE.getAcquire(this);
        boolean empty = q.isEmpty();

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

Refitting operators with async drain

As hinted during the conclusion of the orderedMerge() operator, sometimes we'd want to drain the events of a multi-source operator on another thread. If one thinks about the structure shown in the previous section, the solution to the problem should be clear: propagate an Executor instance into the operator, implement Runnable, keep the WIP increment in drain() and move the rest of the method into the run() method:

    final Executor exec;

    // ...

    void drain() {
        if (getAndIncrement() != 0) {
            return;
        }

        exec.execute(this);
    }

    @Override
    public void run() {

        int missed = 1;
        Flow.Subscriber downstream = this.downstream;

        Comparator comparator = comparator;

        OrderedMergeInnerSubscriber[] subscribers = this.subscribers;
        int n = subscribers.length;

        Object[] values = this.values;

        long e = emitted;

        for (;;) {
             long r = (long)REQUESTED.getAcquire(this);

             for (;;) {
                 /* unchanged, left out for brevity */
             }

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

Cancelling the drain task

So far, we used the void execute(Runnable) method and relied upon cancelled to stop the drain process itself. However, there are a couple of problems:

• What if the onNext call on downstream blocks or takes a long time?
• What if the Executor is so busy the drain doesn't execute within a timeout?

    final ExecutorService execService;

    // ...

    void schedule() {
        if ((int)WIP.getAndAdd(this, 1) == 0) {
            Future<?> future = exec.submit(this);
            // do something with "future"
        }
    }

At this point we have to make future available to the cancellation routine in some thread-safe fashion because cancel() can be called at any time from any thread, even while the schedule() method is executing. The solution is to use the deferred cancellation pattern. First we have to define a field to store the current Future of the running drain task and an associated VarHandle to perform atomic operations with the reference:

    Future<?> future;
    static final VarHandle FUTURE =
        VH.find(MethodHandles.lookup(), ObserveOnSubscriber.class,
"future", Future.class);
    static final Future<?> CANCELLED_TASK = new FutureTask<>(() -> { }, null);

Next, we have to update the schedule() method to store the Future obtained from the ExecutorService.submit call not just atomically, but also avoiding a peculiar race.

    void schedule() {
        if ((int)WIP.getAndAdd(this, 1) == 0) {

            Future<?> old = (Future<?>)FUTURE.getAcquire(this);
            if (old != CANCELLED_TASK) {
                Future<?> future = exec.submit(this);
                if (!FUTURE.compareAndSet(this, old, future)) {
                    old = (Future<?>)FUTURE.getAcquire(this);
                    if (old == CANCELLED_TASK) {
                        future.cancel(true);
                    }
                }
            }
        }
    }

Once the submit() returns, we might no longer be under the protection of the wip != 0 and thus a concurrent call to schedule() (via request() for example) may have triggered another drain run with another Future instance. If we'd update the future reference unconditionally, that newer Future instance would be knocked out and the task-cancellation would no longer work properly.

In addition, a concurrent cancellation may have swapped in the CANCELLED_TASK indicator in which case we should cancel the Future right away. There is no need to loop this construct because if the CAS fails, it is either due to a cancellation or a newer Future task. The former requires cancelling the returned Future and the latter can be ignored because the returned Future is practically finished anyway.

    @Override
    public void cancel() {
        if (CANCELLED.compareAndSet(this, false, true)) {
            upstream.cancel();

            Future<?> future = (Future<?>)FUTURE.getAndSet(this, CANCELLED_TASK);
            if (future != null && future != CANCELLED_TASK) {

                future.cancel(true);
            }

            if ((int)WIP.getAndAdd(this, 1) == 0) {
                queue.clear();
            }
        }
    }

The cancel() method requires less changes: we atomically replace the current future with the CANCELLED_TASK indicator and if otherwise the operator was not cancelled already, we cancel the non-nullFuture instance and carry on.

One can, of course, go extra lengths and make sure the future field is nulled out in between drain runs to avoid leaking references with certain executor service-alikes.

Conclusion

In this post I've shown how one can implement a thread-switching operator, observeOn in RxJava terms, with the Java 9 Flow API and the standard Java async-providing in options such as Executors. The presented algorithms which automatically trampoline the draining of the queued events allows using all sorts of Executors, including a completely synchronous Runnable::runExecutor that executes the task on the caller thread. Since the operator's logic is largely independent of how the asynchrony is established, it can be tested quite easily in a synchronous unit test by itself or as part of a complicated chain of potentially asynchronous flows.

In practice, one may have consider the handling of a RejectedExecutionException thrown by the execute/submit methods on certain ExecutorServices. The trouble is that even though downstream.onError can be called - as no drain() is/can run at that time -, the downstream may really expect errors also delivered on the same background thread as the onNext items. In addition, one has to suppress other onXXX signals from the upstream similar to how a failed user-provided function in an onNext has to suppress any subsequent onError or onComplete call by the upstream because these terminal events are not required to stop emitting events. This part is left to the library developer to decide on his/her own.

In the next post, we'll see how one can stop a flow if items didn't arrive in a timely manner.

Introduction

One of the main properties of reactive programming is that the events may arrive over time instead of immediately available to a consumer. In traditional Future-based programming, one could wait for the result in a blocking manner via Future.get(long, TimeUnit). Other data sources, such as network InputStream have either their own built-in timeout facility or one has to use external means to close the stream after certain period of time to unblock the reader to it. Java 8 Streams have also no direct timeout support.

In the reactive mindset, one can consider timing out events (items) as requesting an element and racing its arrival against the clock. If the item arrives in time, we should ignore the clock. If the clock fires first, we should stop the sender of the items and somehow notify the consumer of the situation. Perhaps the simplest way is to signal onError with a TimeoutException. Since there could be multiple items from a flow, we have to do this racing for each potential items over and over until the flow terminates.

The timeout operator

Since there is "time" in timeout, we'll need a source of time that can be started and stopped at will. The first tool that comes into mind is the java.util.Timer class, however, even its Javadoc suggest one uses a ScheduledExecutorService instead. If one has to deal with a lot of timed operations, besides of timing out flows, having the control over such signals via a (set of) ScheduledExecutorServices is desirable. Therefore, let's define our timeout API with it:

public static <T> Flow.Publisher<T> timeout(
        Flow.Publisher<T> source,
        long timeout, TimeUnit unit,
        ScheduledExecutorService timer) {

    return new TimeoutPublisher<>(source, timeout, unit, timer);
}

(Note that if one uses the Executors.newScheduledExecutorService(), it has to be shutdown at some point, otherwise it's non-daemon thread by default would prevent the JVM from quitting.)

One primary responsibility of this type of operator is to make sure the downstream's onXXX methods are called in sequence and non-overlapping. However, a timeout may happen at the same time the main source signals an event which could lead to call to onNext and onError concurrently - violating the Flow protocol.

The natural non-blocking solution could be the use of a serialization primitive described some time ago in this blog, but we could go more lean on such behavior due to the special nature of this operator.

Thinking about the possible states of such operator, there are two state transitions we should consider: receiving and emitting the next upstream signal, receiving and emitting an error for the timeout signal. In concurrency land, this means using a state variable and compare-and-set to transition from state to state in an atomic fashion.

There is a small problem though: how does a timeout know it happened after the right item didn't arrive? Relying on the accuracy of Future.cancel() for this might not be the best option. Luckily, there is a way to have both proper mutual exclusion and serialization: by using a long field index to track the index of the upstream item as well as indicate a terminal state via Long.MAX_VALUE - an index unlikely reached.

The idea is as follows: a regular onNext signal from the upstream increments this long field unless it is Long.MAX_VALUE at which point the event is ignored. An upstream onError or onComplete will atomically swap in the Long.MAX_VALUE and if not already at Long.MAX_VALUE, emit the respective signal to the downstream. At the same time, a thread executing the Runnable.run() for on the timeout side will try to swap in the last event's index with a Long.MAX_VALUE too. If the current index hasn't changed during the wait (no onXXX signal from the main source), the atomic swap will ensure that any subsequent onXXX call will ignore its event, thus only one thread at a time will emit a signal to the downstream. This may sound complicated when written in plaintext but the code should look much clearer in a couple of paragraphs.

When doing the parent Flow.Publisher implementation of TimeoutPublisher, one has to consider a peculiar case. When should the timeout for the very first item start? Remember, we have to call onSubscribe on the downstream to provide it with the ability to request and cancel the flow. If the timeout would start after downstream.onSubscribe(), a synchronous and blocking source responding to the downstream's request at that point may not give back control at all, rendering the timeout operator nonoperational. If the timeout started before the call to downstream.onSubscribe(), we risk emitting the TimeoutException before or while the downstream.onSubscribe() is executing. Since we also have to intercept a cancel() call from downstream to stop an ongoing timeout, we have to call downstream.onSubscribe() from within the Flow.Publisher.subscribe() method before we subscribe to the upstream or even start the timeout for the very first item:

    @Override
    public void subscribe(Subscriber<? super T> s) {
        TimeoutSubscriber<T> parent = new TimeoutSubscriber<>(
            s, timeout, unit, timer);

        s.onSubscribe(parent);

        parent.startTimeout(0L);

        source.subscribe(parent);
    }

We create the in-sequence Flow.Subscriber first (which implements Flow.Subscription), send it to the downstream, start the timeout for the first item (index is zero at this point) and we subscribe to the upstream source.

The next step is to write up the skeleton of the TimeoutSubscriber:

static final class TimeoutSubscriber<T> implements
Flow.Subscriber<T>, Flow.Subscription {

    final Flow.Subscriber<? super T> downstream;

    final long timeout;

    final TimeUnit unit;

    final ScheduledExecutorService timer;

    Flow.Subscription upstream;
    static final VarHandle UPSTREAM =
        VH.find(MethodHandles.lookup(), TimeoutSubscriber.class,
"upstream", Flow.Subscription.class);

    long requested;
    static final VarHandle REQUESTED =
        VH.find(MethodHandles.lookup(), TimeoutSubscriber.class,
"requested", long.class);

    long index;
    static final VarHandle INDEX =
        VH.find(MethodHandles.lookup(), TimeoutSubscriber.class,
"index", long.class);

    Future<?> task;
    static final VarHandle TASK =
        VH.find(MethodHandles.lookup(), TimeoutSubscriber.class,
"task", Future.class);

    static final Future<Object> CANCELLED_TASK = 
        new FutureTask<>(() -> { }, null);

    TimeoutSubscriber(
            Flow.Subscriber<? super T> downstream,
            long timeout,
            TimeUnit unit,
            ScheduledExecutorService timer) {
        this.downstream = downstream;
        this.timeout = timeout;
        this.unit = unit;
        this.timer = timer;
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        // TODO implement
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        // TODO implement
    }

    @Override
    public void onComplete() {
        // TODO implement
    }

    @Override
    public void request(long n) {
        // TODO implement
    }

    @Override
    public void cancel() {
        // TODO implement
    }

    void run(long idx) {
        // TODO implement
    }

    void startTimeout(long idx) {
        // TODO implement
    }
}

The fields can be described as follows:

• downstream is the Flow.Subscriber we'd like to signal events to.
• timeout and unit have to be stored as a new timeout task will have to be started after each in-time event.
• timer is the ScheduledExecutorService we will submit tasks delayed by the timeout and unit.
• upstream will hold the source's Flow.Subscription. Its VarHandleUPSTREAM will allow us to defer any cancellation coming from the downstream until there is an instance of it from the upstream. In addition, we have to defer requesting from downstream due to the same reason as with cancellation: since we call the downstream.onSubscribe before subscribing to the upstream, the downstream may issue cancel and/or request to a not-yet available Flow.Subscription.
• requested will hold onto any downstream requests until there is an upstream Flow.Subscription available to the operator. This is part of the so-called deferred requesting pattern.
• index holds the state that indicates if the operator has reached its terminal state naturally or due to timeout.
• task holds the current Future for the timeout task submitted to the timer ScheduledExecutorService or the CANCELLED_TASK indicator if the operator has been cancelled concurrently to its operation and no further scheduling of any tasks should happen. Since only one such task has to run at a time, the actual parallelism level of the ScheduledExecutorService doesn't matter. The value has to be changed atomically, hence the TASKVarHandle after it.

    @Override
    public void onSubscribe(Flow.Subscription s) {
        if (UPSTREAM.compareAndSet(null, s)) {
            long r = (long)REQUESTED.getAndSet(this, 0L);
            if (r != 0L) {
                s.request(r);
            }
        } else {
            s.cancel();
        }
    }

We atomically try to set in the upstream's Flow.Subscription and if successful, we atomically take the accumulated request amount from downstream. If there was any, we request it from upstream. Otherwise, a non-null upstream just means the operator has been cancelled and we cancel the upstream as well.

Before we look into the other onXXX methods, let's see the partners of onSubscribe to show the other sides of the deferred requesting and cancellation logic:

    @Override
    public void request(long n) {
        Flow.Subscription current = (Flow.Subscription)UPSTREAM.getAcquire(this);
        if (current != null) {
            current.request(n);
        } else {
            for (;;) {
                long r = (long)REQUESTED.getAcquire(this);
                long u = r + n;
                if (u < 0L) {
                    u = Long.MAX_VALUE;
                }
                if (REQUESTED.compareAndSet(this, r, u)) {
                    break;
                }
            }

            current = (Flow.Subscription)UPSTREAM.getAcquire(this);
            if (current != null && current != this) {
                long r = (long)REQUESTED.getAndSet(this, 0L);
                if (r != 0L) {
                    current.request(r);
                }
            }
        }
    }

When the downstream calls request(), we could be in two states: the upstream Flow.Subscription is available, in which case we simply forward the request amount to it; or the upstream Flow.Subscription is not available and we have to temporarily accumulate downstream requests (which could happen on any thread and any time). This happens via our standard bounded atomic add operation. Once this succeeds, the upstream may have just called onSubscribe with a valid Flow.Subscription (the value of this indicates cancel() has been invoked, see down below). In that case, we atomically take all the accumulated requests, swapping in a zero in exchange, and if that value was non-zero, we issue the request to this fresh upstream. Concurrent interleaving will find requested zero or non-zero and issue any excess request amount accordingly.

    @Override
    public void cancel() {
        if ((long)INDEX.getAndSet(this, Long.MAX_VALUE) != Long.MAX_VALUE) {

            Flow.Subscription current = (Flow.Subscription)UPSTREAM.getAndSet(this, this);
            if (current != null && current != this) {
                current.cancel();
            }

            Future<?> future = (Future<?>)TASK.getAndSet(this, CANCELLED_TASK);
            if (future != null) {
                future.cancel(false);
            }
        }
    }

First, we atomically swap in the terminal index Long.MAX_VALUE, locking out both the onXXX methods and the timeout task. Then, we atomically swap in the cancelled Flow.Subscription indicator (this) and cancel any available upstream. The same thing has to happen to the timeout task.

Now let's see the onNext implementation:

    @Override
    public void onNext(T item) {

        long idx = (long)INDEX.getAcquire(this);
        if (idx == Long.MAX_VALUE) {
            return;
        }
        if (!INDEX.compareAndSet(this, idx, idx + 1)) {
            return;
        }

        Future<?> future = (Future<?>)TASK.getAcqurie(this);
        if (future == CANCELLED_TASK) {
            return;
        }
        if (future != null) {
            future.cancel(false);
        }

        downstream.onNext(item);

        startTimeout(idx + 1);
    }

First we get the current index and if already at Long.MAX_VALUE, we quit. Next, we try to atomically update the index to the next value and if that fails (due to cancellation or a racy timeout), we quit as well. Once the index has been updated successfully, we cancel the ongoing timeout task, emit the current item and start a new task with the subsequent index value.

The onError and onComplete methods are relatively similar by cancelling the timeout task and signalling the appropriate terminal event:

    @Override
    public void onError(Throwable throwable) {
        if ((long)INDEX.getAndSet(this, Long.MAX_VALUE) != Long.MAX_VALUE) {

            Future<?> future = (Future<?>)TASK.getAndSet(this, CANCELLED_TASK);
            if (future != null) {
                future.cancel(false);
            }

            downstream.onError(throwable);
        }
    }

    @Override
    public void onComplete() {
        if ((long)INDEX.getAndSet(this, Long.MAX_VALUE) != Long.MAX_VALUE) {

            Future<?> future = (Future<?>)TASK.getAndSet(this, CANCELLED_TASK);
            if (future != null) {
                future.cancel(false);
            }

            downstream.onComplete();
        }
    }

The are pretty similar to how cancel() is implemented: we atomically swap in the terminal index Long.MAX_VALUE and cancel the timeout task. Note that cancelling the upstream is not necessary at this point as it is considered cancelled in accordance with the Flow (Reactive-Streams) specification.

    void run(long idx) {
        if (INDEX.compareAndSet(this, idx, Long.MAX_VALUE)) {

            Flow.Subscription current = (Flow.Subscription)UPSTREAM.getAndSet(this, this);
            if (current != null && current != this) {
                current.cancel();
            }

            downstream.onError(new TimeoutException());
        }
    }

Given the last known item index, we atomically try to swap in the terminal index indicator and if successful, we cancel the upstream (directly or in a deferred manner) followed by the signal of the TimeoutException. Since the task executing will end either way, there is no need or reason to cancel the Future tracking the timeout task itself. If the index has been changed (either due to the arrival of an onNext item or cancellation), the timeout task will do nothing.

Finally, let's see how the timeout tasks are scheduled:

    void startTimeout(long idx) {
        Future<?> old = (Future<?>)TASK.getAcquire(this);
        if (old != CANCELLED_TASK) { 
            Future<?> future = timer.schedule(() -> run(idx), timeout, unit);
            if (!TASK.compareAndSet(this, old, future)) {
                future.cancel(false);
            }
        }
    }

First, we get the previous, old task instance so we can conditionally swap in the new task if there was no cancellation in between the two. Then, we schedule the execution of the run() method with the current item index (provided by the subscribe() or the onNext() method). If the actual index doesn't change until the run() executes, it will trigger the timeout logic. Then, we atomically try to replace the old Future with the new one and if that fails (due to cancellation), we cancel the new Future task too.

Conclusion

As demonstrated in this post, writing a basic timeout operator - which just signals a TimeoutException, can be done with a relatively few lines. The complication comes from undestanding why the atomic index changes are actually correct in various race scenarios and how it also ensures proper event serialization towards the downstream.

The same pattern can be used for writing a set of other operators which uses one-shot signal to "interrupt" a main sequence, for example, takeUntil(Flow.Publisher), where in the latter, the Future tracking is replaced by tracking the Flow.Subscription of the other Flow.Publisher.

There is, however, the natural need for doing something else than signalling a TimeoutException: switching to another Flow.Publisher for example. One would think, let's subscribe the TimeoutSubscriber to that fallback Flow.Publisher! Unfortunately, it doesn't work, partly because we'd probably didn't want to timeout the fallback Flow.Publisher's consumption (and even if), partly because we have to make sure the switch over preserves the number of outstanding downstream requests.

The deferred requesting logic in the operator shown in the post is not suitable for such transitions and it requires a peculiar Flow.Subscription management logic I call Subscription Arbitration. It enables a whole suit of other operators to work across sources, such as typical RxJava operators concat(Map), repeat(When), retry(When), timeout with fallback and onErrorResumeNext. In the next blog post, we'll see how the arbitration logic and most of these operators can be implemented.

Introduction

A very common task is to combine multiple sources, or more generally, start consuming a source once the previous source has terminated. The naive approach would be to simply call otherSource.subscribe(nextSubscriber) from onError or onComplete. Unfortunately, this doesn't work for two reasons: 1) it may end up with deep stacks due to a "tail" subscription from onError/onComplete and 2) we should request the remaining, unfulfilled amount from the new source that hasn't be provided by the previous source to not overflow the downstream.

The first issue can be solved by applying a heavyweight observeOn in general and implementing a basic trampolining loop only for certain concrete cases such as flow concatenation to be described in this post.

The second issue requires a more involved source: not only do we have to switch between Flow.Subscriptions from different sources, we have to make sure concurrent request() invocations are not lost and are routed to the proper Flow.Subscription along with any concurrent cancel() calls. Perhaps the difficulty is lessened by the fact that switching sources happens on a terminal event boundary only, thus we don't have to worry about the old source calling onNext while the logic switches to the new source and complicating the accounting of requested/emitted item counts. Enter SubscriptionArbiter.

Subscription arbitration

We have to deal with 4 types of potentially concurrent signals when arbitrating Flow.Subscriptions:

1. A request(long) call from downstream that has to be routed to the current Flow.Subscription
2. A cancel() call from downstream that has to be routed to the current Flow.Subscription and cancel any future Flow.Subscription.
3. A setSubscription(Flow.Subscription) that is called by the current Flow.Subscriber after subscribing to any Flow.Publisher which is not guaranteed to happen on the same thread subscribe() is called (i.e., as with the standard SubmissionPublisher or our range() operator).
4. A setProduced(long n) that is called when the previous source terminates and we want to make sure the new source will be requested the right amount; i.e., we'll have to deduce this amount from the current requested amount so setSubscription will issue the request for the remainder to the new Flow.Subscription.

Let's start with the skeleton of the SubscriptionArbiter class providing these methods:

public class SubscriptionArbiter implements Flow.Subscription {

    Flow.Subscription current;
    static final VarHandle CURRENT =
        VH.find(MethodHandles.lookup(), SubscriptionArbiter.class,
"current", Flow.Subscription.class);

    Flow.Subscription next;
    static final VarHandle NEXT =
        VH.find(MethodHandles.lookup(), SubscriptionArbiter.class,
"next", Flow.Subscription.class);

    long requested;

    long downstreamRequested;
    static final VarHandle DOWNSTREAM_REQUESTED =
        VH.find(MethodHandles.lookup(), SubscriptionArbiter.class,
"downstreamRequested", Flow.Subscription.class);

    long produced;
    static final VarHandle PRODUCED =
        VH.find(MethodHandles.lookup(), SubscriptionArbiter.class,
"produced", Flow.Subscription.class);

    long wip;
    static final VarHandle WIP =
        VH.find(MethodHandles.lookup(), SubscriptionArbiter.class,
"wip", int.class);

    @Override
    public final void request(long n) {
        // TODO implement
    }

    @Override
    public void cancel() {
        // TODO implement
    }

    public final boolean isCancelled() {
        // TODO implement
        return false;
    }

    public final void setSubscription(Flow.Subscription s) {
        // TODO implement
    }

    public final void setProduced(long n) {
        // TODO implement
    }

    final void arbiterDrain() {
        // TODO implement
    }
}

We intend the class to be extended to save on allocation and object headers, however, some methods should not be overridden by any subclass as it would likely break the internal logic. The only relatively safe overridable method is cancel(): the subclass will likely have its own resources that have to be released upon a cancel() call from the downstream which will receive an instance of this class via onSubscribe. The meaning of each field is as follows:

• current holds the current Flow.Subscription. Its companion CURRENTVarHandle is there to support cancellation.
• next temporarily holds the next Flow.Subscription to replace current instance. Direct replacement can't work due to the required request accounting.
• requested holds the current outstanding request count. It doesn't have any VarHandle because it will be only accessed from within a drain-loop.
• downstreamRequested accumulates the downstream's requests in case there the drain loop is executing.
• produced holds the number of items produced by the previous source which has to be deduced from requested before switching to the next source happens. It is accompanied by a VarHandle to ensure proper visibility of its value from within the drain loop.
• wip is our standard work-in-progress counter to support the queue-drain like lock-free serialization we use almost everywhere now.

The first method we implement is request() that will be called by the downstream from an arbitrary thread at any time:

    @Override
    public final void request(long n) {
        for (;;) {
            long r = (long)DOWNSTREAM_REQUESTED.getAcquire(this);
            long u = r + n;
            if (u < 0L) {
                u = Long.MAX_VALUE;
            }
            if (DOWNSTREAM_REQUESTED.compareAndSet(this, r, u)) {
                arbiterDrain();
                break;
            }
        }
    }

We perform the usual atomic addition capped at Long.MAX_VALUE and call arbiterDrain().

    @Override
    public void cancel() {
        Flow.Subscription s = (Flow.Subscription)CURRENT.getAndSet(this, this);
        if (s != null && s != this) {
            s.cancel();
        }
        s = (Flow.Subscription)NEXT.getAndSet(this, this);
        if (s != null && s != this) {
            s.cancel();
        }
    }

    public final boolean isCancelled() {
        return CURRENT.getAcquire(this) == this;
    }

We atomically swap in both the current and the next Flow.Subscription instances with the cancelled indicator of this. To support some eagerness in cancellation, the isCancelled can be called by the subclass (i.e., concat an array of Flow.Publishers) to quit its trampolined looping.

Next, we "queue up" the next Flow.Subscription:

    public final void setSubscription(Flow.Subscription subscription) {
        if (NEXT.compareAndSet(this, null, subscription)) {
            arbiterDrain();
        } else {
            subscription.cancel();
        }
    }

We expect there will be only one thread calling setSubscription and that call happens before the termination of the associated source, thus a simple CAS from null to subscription should be enough. In this scenario, a failed CAS can only mean the arbiter was cancelled in the meantime and we cancel the subscription accordingly. We'll still have to relay the unfulfilled request amount to this new subscription which will be done in arbiterDrain().

The setProducer will have to "queue up" the fulfilled amount in a similar fashion:

    public final void setProduced(long n) {
        PRODUCED.setRelease(this, n);
        arbiterDrain();
    }

As with the setSubscription, we expect this to happen once per a terminated source before the subscription to the next source happens, thus there is no real need to atomically accumulate the item count.

Finally, let's see the heavy lifting in arbiterDrain() itself now:

    final void arbiterDrain() {
        if ((int)WIP.getAndAdd(this, 1) != 0) {
            return;
        }

        Flow.Subscription requestFrom = null;
        long requestAmount = 0L;

        for (;;) {

            // TODO implement

            if ((int)WIP.getAndAdd(this, -1) - 1 == 0) {
                break;
            }
        }

        if (requestFrom != null && requestFrom != this
&& requestAmount != 0L) {
            requestFrom.request(requestAmount);
        }
    }

The arbiterDrain(), whose name was chosen to avoid clashing with a subclass'drain() method if any, method starts out as most typical trampolined drain loop did: the atomic increment to wip from 0 to 1 enters the loop and the decrement to zero leaves the loop.

One oddity may come from the requestFrom and requestAmount local variables. Unlike a traditional stable-prefetch queue-drain, requesting from within the loop can bring back the reentrancy issue, the tail-subscription problem and may prevent other actions from happening with the arbiter until the request() call returns. Therefore, once the loop decided what the current target Flow.Subscription is, we'll issue a request to it outside the loop. It is possible by the time the drain method reaches the last if statement that the current requestFrom is outdated or the arbiter was cancelled. This is not a problem because request() and cancel() in general are expected to race and an outdated Flow.Subscription means it has already terminated and a request() call is a no-op to it.

The last part inside the loop has to "dequeue" the deferred changes and apply them to the state of the arbiter:

    for (;;) {

        // (1) ----------------------------------------------
        Flow.Subscription currentSub = (Flow.Subscription)CURRENT.getAcquire(this);
        if (currentSub != this) {

            // (2) ------------------------------------------
            long req = requested;

            long downstreamReq = (long)DOWNSTREAM_REQUESTED.getAndSet(this, 0L);

            long prod = (long)PRODUCED.getAndSet(this, 0L);

            Flow.Subscription nextSub = (Flow.Subscription)NEXT.getAcquire(this, null);
            if (nexSub != null && nextSub != this) {
                NEXT.compareAndSet(this, nextSub, null);
            }

            // (3) ------------------------------------------
            if (downstreamReq != 0L) {
                req += downstreamReq;
                if (req < 0L) {
                    req = Long.MAX_VALUE;
                }
            }

            // (4) ------------------------------------------
            if (prod != 0L && req != Long.MAX_VALUE) {
                req -= prod;
            }

            requested = req;

            // (5) ------------------------------------------
            if (nextSub != null && nextSub != this) {
                requestFrom = nextSub;
                requestAmount = req;
                CURRENT.compareAndSet(currentSub, nextSub);
            } else {
                // (6) --------------------------------------
                requestFrom = currentSub;
                requestAmount += downstreamReq;
                if (requestAmount < 0L) {
                    requestAmount = Long.MAX_VALUE;
                }
            }
        }

        if ((int)WIP.getAndAdd(this, -1) - 1 == 0) {
            break;
        }
    }

1. First we check if the current instance holds the cancelled indicator (this). If so, we don't have to execute any of the logic as the arbiter has been cancelled by the downstream.
2. We read out the current and queued state: the current outstanding requested amount, the request amount from the downstream if any, the produced item count by the previous source and the potential next Flow.Subscription instance. While it is safe to atomically swap in 0 for both the downstreamRequested and produced values, swapping in null unconditionally may overwrite the cancelled indicator and the setSubscription won't cancel its argument.
3. If there was an asynchronous request() call, we add the downstreamReq amount to the current requested amount, capped at Long.MAX_VALUE (unbounded indicator).
4. If there was a non-zero produced amount and the requested amount isn't Long.MAX_VALUE, we subtract the two. The new requested amount is then saved.
5. If there was a new Flow.Subscription set via setSubscription, we indicate where to request from outside the loop and we indicate the whole current requested amount (now including any async downstream request and upstream produced count) should be used. The CAS will make sure the next Flow.Subscription only becomes the current one if there was no cancellation in the meantime.
6. Otherwise, we target the current Flow.Subscription, add up the downstream's extra requests capped at Long.MAX_VALUE. The reason for this is that the downstream may issue multiple requests (r1, r2) in a quick succession which makes the logic to loop back again, now having r1 + r2 items outstanding from the downstream's perspective.

Concatenating an array of Flow.Publishers

@SafeVarargs
public static <T> Flow.Publisher<T> concat(Flow.Publisher<? extends T>... sources) {
    return new ConcatPublisher<>(sources);
}

The ConcatPublisher itself is straightforward: create a coordinator, send it to the downstream and trigger the consumption of the first source:

     @Override
     public void subscribe(Flow.Subscriber<? super T> s) {
         ConcatCoordinator<T> parent = new ConcatCoordinator<>(s, sources);
         s.onSubscribe(parent);
         parent.drain();
     }

The ConcatCoordinator can be implemented as follows:

static final class ConcatCoordinator<T> extends SubscriptionArbiter
implements Flow.Subscriber<T> {

    final Flow.Subscription<? super T> downstream;

    final Flow.Publisher<? extends T>[] sources;

    int index;

    int trampoline;
    static final VarHandle TRAMPOLINE =
        VH.find(MethodHandles.lookup(), ConcatCoordinator.class,
"trampoline", int.class);

    long consumed;

    ConcatCoordinator(
           Flow.Subscription<? super T> downstream,
           Flow.Publisher<? extends T>[] sources
    ) {
        this.downstream = downstream;
        this.sources = sources;
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        // TODO implement
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        // TODO implement
    }

    @Override
    public void onComplete() {
        // TODO implement
    }

    void drain() {
        // TODO implement
    }
}

The ConcatCoordinator extends the SubscriptionArbiter, thus it is a Flow.Subscription as well and as such will be used as the connection object towards the downstream. It also implements Flow.Subscription because we'll use the same instance to subscribe to all of the Flow.Publishers one after the other.

One may come up with the objection that reusing the same Flow.Subscriber instance is not allowed by the Reactive-Streams specification the Flow API inherited. However, the specification actually just discourages the reuse and otherwise mandates external synchronization so that the onXXX methods are invoked in a serialized manner. We'll see that the trampolining in the operator will just ensure that property along with the arbiter itself. Of course, we could just new up a Flow.Subscriber for the next source but that Flow.Subscriber would be itself nothing more than a delegator for the coordinator instance (no need for a per-source state in it); combining the two just saves on allocation and indirection.

The fields are interpreted as follows:

• downstream is the Flow.Subscriber that receives the signals.
• sources is the array of Flow.Publishers that will be consumed one after the other
• index points at the current Flow.Publisher and gets incremented once one completes.
• trampoline is the work-in-progress indicator for the drain loop; chosen to avoid clashing with the arbiter's own wip field in this blog for better readability. In practice, since they are in different classes, one can name them both wip.
• consumed tracks how many items the current source has produced (and has the coordinator consumed). We'll use this to update the arbiter at the completion of the current source instead of doing it for each item received because that saves a lot of overhead and we don't really care about each individual item's reception.

The coordinator's onXXX methods are relatively trivial at this point:

   @Override
    public void onSubscribe(Flow.Subscription s) {
        setSubscription(s);
    }

    @Override
    public void onNext(T item) {
        consumed++;
        downstream.onNext(item);
    }

    @Override
    public void onError(Throwable throwable) {
        downstream.onError(throwable);
    }

    @Override
    public void onComplete() {
        drain();
    }

We save the Flow.Subscription into the arbiter, write through the item or throwable and call the drain method upon normal completion.

What's left is the drain() method itself:

    void drain() {
        // (1) -------------------------------------------------------
        if ((int)TRAMPOLINE.getAndAdd(this, 1) == 0) {
            do {

                // (2) -----------------------------------------------
                if (isCancelled()) {
                    return;
                }

                // (3) -----------------------------------------------
                if (index == sources.length) {
                    downstream.onComplete();
                    return;
                }

                // (4) -----------------------------------------------
                long c = consumed;
                if (c != 0L) {
                    consumed = 0L;
                    setProduced(c);
                }

                // (5) -----------------------------------------------
                sources[index++].subscribe(this);

            // (6) ---------------------------------------------------
            } while ((int)TRAMPOLINE.getAndAdd(this, -1) - 1 != 0);
        }
    }

Again, not really a complicated method, but as usual, the difficulty may come from understanding why such short code is actually providing the required behavior and safeguards:

1. We know that drain() is only invoked from the subscribe() or onComplete() methods. This standard lock-free trampolining check ensures only one thread is busy setting up the consumption of the next (or the first) source. In addition, since only a guaranteed one-time per source onComplete() can trigger the consumption of the next, we don't have to worry about racing with onNext in this operator. (However, an in-flow concatMap is a different scenario.) This setup also defends against increasing the stack depth due to tail-subscription: a trampoline> 1 indicates we can immediately subscribe to the next source.
2. In case the downstream cancelled the operator, we simply quit the loop.
3. In case the index is equal to the number of sources, it means we reached the end of the concatenation and can complete the downstream via onComplete().
4. Otherwise, we indicate to the arbiter the number of items consumed from the previous source so it can update its outstanding (current) request amount. Note that consumed is not concurrently updated because onNext and onComplete (and thus drain) on the same source can't be executed concurrently.
5. We then subscribe to the next source, move the index forward by one to point to the next-next source and subscribe with this.
6. Finally if there was no synchronous or racy onComplete, we quit the loop, otherwise we resume with the subsequent sources.

One can add a few features and safeguards to this coordinator, such as delaying errors till the very end and ensuring the indexth sources entry is not null. These are left as exercise to the reader.

Repeat

How can we turn this into a repeat operator where the source is resubscribed on successful completion? Easy: drop the index and have only a single sourceFlow.Publisher to be worked on!

public static <T> Flow.Publisher<T> repeat(
        Flow.Publisher<T> source, long max) {
    return new RepeatPublisher<>(source, max);
}

// ... subscribe() has the same pattern.

static final class RepeatCoordinator<T> extends SubscriptionArbiter
implements Flow.Subscriber<T> {

    final Flow.Publisher<T> source;

    long max;

    // ... the onXXX methods are the same

    final void drain() {
        if ((int)TRAMPOLINE.getAndAdd(this, 1) == 0) {
            do {

                if (isCancelled()) {
                    return;
                }

                if (--max < 0L) {
                    downstream.onComplete();
                    return;
                }

                long c = consumed;
                if (c != 0L) {
                    consumed = 0L;
                    setProduced(c);
                }

                source.subscribe(this);

            } while ((int)TRAMPOLINE.getAndAdd(this, -1) - 1 != 0);
        }
    }
}

Given that repeating indefinitely is usually not desired, we limit the resubscriptions to a number of times specified by the user. Since there is only one source Flow.Publisher, no indexing into an array is needed and we only have to decrement the counter to detect the condition for completing the downstream.

Retry

How about retrying a Flow.Publisher in case it failed with an onError? Easy: have onError call drain() and onComplete call downstream.onComplete() straight:

public static <T> Flow.Publisher<T> retry(
        Flow.Publisher<T> source, long max) {
    return new RepeatPublisher<>(source, max);
}

// ... subscribe() has the same pattern.

static final class RetryCoordinator<T> extends SubscriptionArbiter
implements Flow.Subscriber<T> {

    final Flow.Publisher<T> source;

    long max;

    // ... the onSubscribe and onNext methods are the same

    @Override
    public void onError(Throwable throwable) {
        if (--max < 0L) {
            downstream.onError(throwable);
        } else {
            drain();
        }
    }

    @Override
    public void onComplete() {
        downstream.onComplete();
    }

    final void drain() {
        if ((int)TRAMPOLINE.getAndAdd(this, 1) == 0) {
            do {

                if (isCancelled()) {
                    return;
                }

                long c = consumed;
                if (c != 0L) {
                    consumed = 0L;
                    setProduced(c);
                }

                source.subscribe(this);

            } while ((int)TRAMPOLINE.getAndAdd(this, -1) - 1 != 0);
        }
    }
}

There are two slight changes in retry(). First, in case we run out of the retry count, the latest Flow.Publisher's error is delivered to the downstream from within onError and no further retry can happen. Consequently, the drain logic no longer should call onComplete when the number of allowed retries have been used up.

On error, resume with another Flow.Publisher

Now that we've seen multi-source switchover and single-source continuation, switching to an alternative or "fallback"Flow.Publisher should be straightforward to set up: have a 2 element array with the main and fallback Flow.Publishers, then make sure onError triggers the switch.

public static <T> Flow.Publisher<T> onErrorResumeNext(
        Flow.Publisher<T> source, Flow.Publisher<T> fallback) {
    return new RepeatPublisher<>(source, fallback);
}

// ... subscribe() has the same pattern.

static final class RetryCoordinator<T> extends SubscriptionArbiter
implements Flow.Subscriber<T> {

    final Flow.Publisher<T> source;

    final Flow.Publisher<T> fallback;

    boolean switched;

    // ... the onSubscribe and onNext methods are the same

    @Override
    public void onError(Throwable throwable) {
        if (switched) {
            downstream.onError(throwable);
        } else {
            switched = true;
            drain();
        }
    }

    @Override
    public void onComplete() {
        downstream.onComplete();
    }

    final void drain() {
        if ((int)TRAMPOLINE.getAndAdd(this, 1) == 0) {
            do {

                if (isCancelled()) {
                    return;
                }

                long c = consumed;
                if (c != 0L) {
                    consumed = 0L;
                    setProduced(c);
                }

                if (switched) {
                    fallback.subscribe(this);
                } else {
                    source.subscribe(this);
                }

            } while ((int)TRAMPOLINE.getAndAdd(this, -1) - 1 != 0);
        }
    }
}

Here, we have two states, switched == false indicates we are consuming the main source. If that fails, we set switched = true and the drain loop will subscribe to the fallback Flow.Publisher. However, if the fallback fails, the onError also checks for switched == true and instead of draining (and thus retrying) the fallback Flow.Publisher again, it just terminates with the Throwable the fallback emitted.

Conclusion

In this post, the subscription arbitration concept was presented which allows us to switch between non-overlapping Flow.Publisher sources when one terminates (completes or fails with an error) while maintaining the link of cancellation between the individual Flow.Subscriptions as well as making sure backpressure is properly transmitted and preserved between them.

When combining with a trampolining logic, such arbitration allowed us to implement a couple of standard ReactiveX operators such as concat, repeat, retry and onErrorResumeNext while only applying small changes to the methods and algorithms in them.

Note however, that even if the arbiter can be reused for in-flow operators such as concatMap (concatenate Flow.Publishers generated from upstream values), repeatWhen (repeat if a companion Flow.Publisher signals an item) and retryWhen, one can no longer use a single Flow.Subscriber to subscribe to both the main flow and the inner/companion flows at the same time. We will explore these types of operators in a later post.

The arbitration has its own limit: it can't support live switching between sources, i.e., when onNext may be in progress while the switch has to happen. If you are familiar with the switchMap operator, this is what can happen during its execution. We'll look into this type of operator in a subsequent post.

But for now, we'll investigate a much lighter set of operators in the next post: limiting the number of items the downstream can receive and skipping certain number of items from the upstream; both based on counting items and based on a per-item predicate checks, i.e., the take() and skip() operators.

Introduction

Limiting or skipping over parts of a flow is a very common task: either we are only interested in the first N items or we don't care about the first N items. Sometimes, N is unknown but we can decide, based on the current item, when to stop relaying items or, in contrast, when to start relaying items.

Take(N)

In concept, limiting a flow to a certain size should be straightforward: count the number of items received via onNext and when the limit is reached, issue a cancel() towards the upstream and onComplete() towards the downstream.

public static <T> Flow.Publisher<T> take(Flow.Publisher<T> source, long n) {
    return new TakePublisher<>(source, n);
}

The operator's implementation requires little state:

static final class TakeSubscriber<T> implements Flow.Subscriber<T> {

    final Flow.Subscriber<? super T> downstream;

    Flow.Subscription upstream;

    long remaining;

    TakeSubscriber(
            Flow.Subscriber<? super> downstream,
            long n) {
        this.downstream = downstream;
        this.remaining = n;
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        // TODO implement
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        // TODO implement
    }

    @Override
    public void onComplete() {
        // TODO implement
    }
}

In its simplest form, there is no need for intercepting the request() and cancel() calls from the downstream: these can be passthrought, however, since the operator has to stop the sequence upon reaching the limit (remaining == 0), the upstream's Flow.Subscriber has to be stored.

    @Override
    public void onSubscribe(Flow.Subscription s) {
        this.upstream = s;
        downstream.onSubscribe(s);
    }

In onSubscribe, we only have to store the Flow.Subscription and forward it to the downstream.

    @Override
    public void onNext(T item) {
        long r = remaining;
        if (r > 0L) {
            remaining = --r;
            downstream.onNext(item);

            if (r == 0) {
                upstream.cancel();
                downstream.onComplete();
            }
        }
    }

While remaining is positive, we decrement it and save it into its field followed by an emission of the current item. If the remaining count reached zero, the upstream is cancelled and the downstream is completed. Any further items will be ignored (in case cancel() doesn't immediately stop the upstream).

    @Override
    public void onError(Throwable throwable) {
        if (remaining != 0L) {
            downstream.onError(throwable);
        }
    }

    @Override
    public void onComplete() {
        if (remaining != 0L) {
            downstream.onComplete();
        }
    }

The onError and onComplete check the remaining count and if it's positive, the respective terminal event is reached. The reason for this is that if the operator runs out of items before its limit, the terminal event is relayed as well. However, if the operator happens to run with a limit equal to the actual length of the flow, the last item in onNext will trigger an onComplete which could be followed by a terminal event from the upstream that has to be suppressed. This behavior is allowed by the Reactive-Streams specification (i.e., cancel() may not have immediate effect) and we have dealt with it via a done field in other operators. Here, the fact remaining == 0 is the done indicator.

Being a pass-through for backpressure (the downstream calls request() directly on the upstream's Flow.Subscription), the upstream is not limited by this operator and may attempt to produce more items than the limit. In other times, knowing the downstream requested more than the limit, the upstream can actually be consumed in an unbounded fashion, utilizing fast-paths may result in improved performance.

Deciding which of the three behaviors should be default is up to the operator writer, but it is interesting to look at the remaining two modes: limiting the request amount and unbounding it.

Limiting the request amount

In this mode, if the downstream requests more or equal to the limit, we can issue a single request with the limit amount. However, if the downstream requests less, we have to perform some intricate request accounting.

We'll need a new field, requestRemaining with a VarHandleREQUEST_REMAINING companion. We also have to intercept request() from downstream and figure out atomically what number to request from the upstream so the total requested amount doesn't go above the limit.

static final class TakeSubscriber<T> implements 
Flow.Subscriber<T>, Flow.Subscription {

    final Flow.Subscriber<? super T> downstream;

    Flow.Subscription upstream;

    long remaining;

    long requestRemaining;
    static final VarHandle REQUEST_REMAINING =
        VH.find(MethodHandles.lookup(), TakeSubscriber.class,
"requestRemaining", long.class);

    TakeSubscriber(
            Flow.Subscriber<? super> downstream,
            long n) {
        this.downstream = downstream;
        this.remaining = n;
        this.requestRemaining = n;
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        upstream = s;
        downstream.onSubscribe(this);
    }

    // onNext, onError and onComplete are the same

    @Override
    public void request(long n) {
        // TODO implement
    }

    @Override
    public void cancel() {
        upstream.cancel();
    }
}

The onSubscribe now forwards this as the Flow.Subscription and cancel() just delegates to the upstream.cancel().

    @Override
    public void request(long n) {
        for (;;) {
            long r = (long)REQUEST_REMAINING.getAquire(this);
            long newRequest;
            if (r <= n) {
                newRequest = r;
            } else {
                newRequest = n;
            }
            long u = r - newRequest;
            if (REQUEST_REMAINING.compareAndSet(this, r, u)) {
                upstream.request(newRequest);
                break;
            }
        }
    }

First we read the remaining request amount. If it is smaller or equal to the downstream's request, we'll request the remaining amount from the upstream. If the downstream requested less than the remaining amount, we'll request that amount instead. The new remaining amount is then the current minus the new request amount decided. After the successful CAS, we request the new amount from the upstream and quit the loop.

Unbounding the request amount

The other direction, namely, unbounding the upstream if the downstream requested at least the limit of the operator, can be achieved through the same requestRemaining field but a different request() logic:

    @Override
    public void request(long n) {
        long r = (long)REQUEST_REMAINING.getAcquire(this);
        if (r != 0L) {
             r = (long)REQUEST_REMAINING.getAndSet(this, 0L);
             if (r != 0L && r <= n) {
                 upstream.request(Long.MAX_VALUE);
                 return;
             }
        }
        upstream.request(n);
    }

If the remaining request amount is non-zero, we atomically replace it with zero. This happens exactly once, for the first time request() is invoked by the downstream, and we'll check if the remaining amount (same as the limit) is less or requal to the downstream's request amount. If so, we request from the upstream an unbounded amount (Long.MAX_VALUE). If it was not the first request or the amount was less than the limit, we revert back to the pass-through behavior.

Take with predicate

Sometimes, we'd like to stop the flow when an item matches a certain condition. This can happen before the item is emitted (i.e., takeWhile) or after (i.e., takeUntil). When working with a predicate, we no longer know how many items we'll let pass thus manipulating the request amount is not really an option here.

public static <T> Flow.Publisher<T> takeWhile(
        Flow.Publisher<T> source, Predicate<? super T> predicate) {

    return new TakeWhilePublisher<>(source, predicate);
}

// Flow.Publisher boilerplate omitted

static final class TakeWhileSubscriber<T> implements Flow.Subscriber<T> {

    final Flow.Subscriber<? super T> downstream;

    final Predicate<? super T> predicate;

    Flow.Subscription upstream;

    boolean done;

    TakeWhileSubscriber(
            Flow.Subscriber<? super> downstream,
            Predicate<? super T> predicate) {
        this.downstream = downstream;
        this.predicate = predicate;
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        upstream = s;
        downstream.onSubscribe(s);
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        if (!done) {
            downstream.onError(throwable);
        }
    }

    @Override
    public void onComplete() {
        if (!done) {
            downstream.onComplete();
        }
    }
}

The onSubscribe simply forwards the Flow.Subscription and the terminal onError/onComplete methods forward their respective event if the done flag is false. The flag will be set in onNext if the flow has been stopped due to the condition and is there to prevent the emission of multiple terminal events in case the flow would have ended anyway.

    @Override
    public void onNext(T item) {
        if (done) {
            return;
        }

        boolean b;

        try {
            b = predicate.test(item);
        } catch (Throwable ex) {
            upstream.cancel();
            done = true;
            downstream.onError(ex);
            return;
        }

        if (b) {
            downstream.onNext(item);
        } else {
            upstream.cancel();
            done = true;
            downstream.onComplete();
        }  
    }

The first step is to protect against sources that would produce a few items after a cancel() call. Next, we don't trust the user-provided Predicate, which when crashes, will have to cancel the upstream, lock out further events and signal the Throwable to the downstream. A true result from the predicate will allow us to emit the item to the downstream. A false will stop the source, lock out further upstream events and complete the downstream.

If we'd still want to receive the item before the predicate indicates the flow should stop, aka the takeUntil operator, only the onNext logic should be changed a bit:

public static <T> Flow.Publisher<T> takeUntil(
        Flow.Publisher<T> source, Predicate<? super T> stopPredicate) {

    return new TakeUntilPublisher<>(source, stopPredicate);
}

    @Override
    public void onNext(T item) {
        if (done) {
            return;
        }

        downstream.onNext(item);

        boolean b;

        try {
            b = stopPredicate.test(item);
        } catch (Throwable ex) {
            upstream.cancel();
            done = true;
            downstream.onError(ex);
            return;
        }

        if (b) {
            upstream.cancel();
            done = true;
            downstream.onComplete();
        }  
    }

Here, the stopPredicate of the API entry point should indicate when to stop by returning true, whereas the previous takeWhile operator indicated a stop via false. It is a matter of taste I guess. The RxJava convetion is that takeWhile(item -> item < 5) will take items while each of them is less than five, never emitting five itself whereas takeUntil(item -> item == 5) will stop after emitting five.

Skip(N)

The dual of take(N), in some sense, is the operator which skips the first N items then lets the rest through. Again, counting the items is a crucial part of the operator's implementation.

public static <T> Flow.Publisher<T> skip(
        Flow.Publisher<T> source, long n) {

    return new SkipPublisher<>(source, n);
}

The operator's Flow.Subscriber uses similar counting method as take(), decrementing a remaining field and once it reaches zero, all subsequent events are forwarded.

static final class SkipSubscriber<T> implements Flow.Subscriber<T> {

    final Flow.Subscriber<? super T> downstream;

    Flow.Subscription upstream;

    long remaining;

    SkipSubscriber(
            Flow.Subscriber<? super> downstream,
            long n) {
        this.downstream = downstream;
        this.remaining = n;
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        // TODO implement
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        downstream.onError(throwable);
    }

    @Override
    public void onComplete() {
        downstream.onComplete();
    }
}

The onError/onComplete methods can now emit the respective event directly without the need to check the remaining count. May the sequence be shorter, equal or longer than the provided limit, the terminal events can be emitted as onNext is there to skip items, not stop the flow.

In general, when an operator drops items, it has the responsibility to ask for more from the upstream as the downstream has no way of knowing it didn't receive an item (other than timing out). Thus, it doesn't know it has to request more and has to be requested on behalf by the operator itself. However, requesting one by one is expensive - 1 CAS + 0..2 atomic increments per invocation - which can be avoided by requesting in batches.

The skip() operator is in a particular position where we know the first N items will be dropped, thus we can simply request N on top of what the downstream requested. onNext will drop the first N and even if the downstream hasn't requested, it won't get overflown.

    @Override
    public void onSubscribe(Flow.Subscription s) {
        upstream = s;

        long n = remaining;
        downstream.onSubscribe(s);
        s.request(n);
    }

The reason remaining is read before sending the Flow.Subscription to the downstream is that this call may result in value emission which decrements remaining and we'd end up with less than the amount to be skipped. This also saves a field storing the logically immutable skip amount.

The role of onNext is to drop items and then let the rest pass through:

    @Override
    public void onNext(T item) {
        long r = remaining;
        if (r == 0L) {
            downstream.onNext(item);
        } else {
            remaining = r - 1;
        }
    }

Skipping while a condition holds

Similar to a conditional take, we can skip unwanted items that match a condition (predicate) and the let the rest through unconditionally. Unfortunately, this case requires per item replenishment from the upstream if the condition doesn't hold as we can't be sure which upstream item will yield a true value and switch the operator into a pass-through mode.

public static <T> Flow.Publisher<T> skipWhile(
        Flow.Publisher<T> source, Predicate<? super T> predicate) {

    return new SkipWhilePublisher<>(source, predicate);
}

static final class SkipWhileSubscriber<T> implements Flow.Subscriber<T> {

    final Flow.Subscriber<? super T> downstream;

    final Predicate<? super T> predicate;

    Flow.Subscription upstream;

    boolean passthrough;

    boolean done;

    SkipWhileSubscriber(
            Flow.Subscriber<? super> downstream,
            Predicate<? super T> predicate) {
        this.downstream = downstream;
        this.predicate = predicate;
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        upstream = s;
        downstream.onSubscribe(s);
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable throwable) {
        if (!done) {
            downstream.onError(throwable);
        }
    }

    @Override
    public void onComplete() {
        if (!done) {
            downstream.onComplete();
        }
    }
}

The onSubscribe becomes a direct pass-through for the Flow.Subscription and the onError/onComplete methods get their done checks back. This is required because the predicate may fail for the very last item and the cancellation may not stop the terminal event. There is also the need for a flag that tells the onNext to let all items through and skip the predicate altogether.

    @Override
    public void onNext(T item) {
        if (done) {
            return;
        }

        if (!passthrough) {
            boolean b;
            try {
                b = predicate.test(item);
            } catch (Throwable ex) {
                upstream.cancel();
                done = true;
                downstream.onError(ex);
                return;
            }
            if (b) {
                upstream.request(1);
                return;
            }
            passthrough = true;           
        }

        downstream.onNext(item);
    }

First we prevent an excess onNext if there was a crash in the predicate before. Next, if the operator is not in the pass-through mode, we test with the predicate. If this turns out to be true, that indicates the item has to be dropped and a fresh item has to be requested as replenishment. Otherwise the operator enters its pass-through mode and the current and subsequent items will be emitted directly without invoking the predicate again. Implementing an "until" variant, where the predicate returning true still drops the current item, is left to the reader as an exercise.

Conclusion

The post demonstrated a way to implement the take, takeWhile, takeUntil, skip and skipWhile operators. The logic behind them doesn't require a fully fledged queue-drain approach and for being synchronous in-flow operators, don't really have to bother with concurrency and request management (except the two take alternatives which do have to manage requests).

The fact that the onXXX methods of a Flow.Subscriber are not allowed to be invoked concurrently greatly helps reducing the complexity and the need for being defensive all the time.

In the next post, we'll see how a multicasting Processor can be implemented and how a fan-out-fan-in operator can be implemented with its help.

Introduction

Threading and lifecycle are one of the top concerns when developing applications for the Android platform. UI has to be interacted with on a dedicated thread (main thread) but in order to keep the UI responsible to user input and rendering, blocking or CPU intensive calculations should be kept off it.

In addition, views can get destroyed and recreated in a way that is outside of a given application's control unlike a desktop Swing application. This means background tasks must be stopped and listeners removed to prevent leaking references to now logically dead objects.

RxJava and RxAndroid can help with threading concerns and there are other libraries that tap into the various lifecycle events; in general, this means someone will call dispose() on a particular flow or clear() on a CompositeDisposable to mass-cancel multiples of them.

Having a rich set of transformative and coordinating operators along with support for normal values, errors and finite sequences may be overwhelming compared to a classical Listener-based API. Google's LiveData is one of such classical Listener style APIs but unlike Swing's ActionListener for example, there are explicit requirements that interaction with the LiveData object itself happens on the main thread and signals will be dispatched from the main thread to Observers to it.

LiveData API

Unfortunately, I wasn't able to locate a public repository for the LiveData sources and had to rely on the sources downloaded from Google's Maven repository:

compile "android.arch.lifecycle:reactivestreams:1+"

There is an interoperation library associated with LiveData that allows presenting and consuming events from any Reactive-Streams Publisher. This will transitively import the actual LiveData library. Note that LiveData is currently considered beta and may change arbitrarily before release. That said, I don't think the core structure and premise will actually change.

The main consumer type is the android.arch.lifecycle.Observer with its single onChanged(T) method. Given an arbitrary LiveData instance, one can use observe(LifecycleOwner, Observer) or observeForever(Observer) methods to attach the listener to a source. The onChanged(T) will be called on the main thread. Cancelling an observation requires one to call removeObserver() for which one has to remember both the original LiveData object and the Observer instance itself. All the methods return void. The LifecycleOwner will be listened for destroy events, at which point the Observers are removed from the LiveData instance that depended on that specific source. This is similar to RxBinding's bindToLifecycle() approach composed onto a flow.

The LiveData abstract class doesn't feature any transformative or coordinating operators on itself, but the library features an utility class, Transformations, currently with a map and switchMap operations. Both will execute the transformations and emissions on the main thread, but otherwise they match the behavior of the same RxJava operators.

In theory, one could interpret LiveData as a cold source, but its usage patterns point to a hot source, just like a BehaviorSubject in RxJava; LiveData remembers the last value and emits it first to new consumers. To signal to consumers, LiveData has protected methods setValue and postValue. setValue requires the value to be set on the main thread and postValue will schedule a task to the main thread to call setValue there. To manually invoke these methods, a MutableLiveData object is available where the two methods have been made public.

Possible shortcomings

The main shortcoming I see with LiveData is the ubiquitous "return to main thread" mentality. For example, querying data from a LiveData-enabled datasource and preprocessing it should happen in a background thread as well before the aggregated data is sent back to the main thread for displaying them. Unfortunately, the datasource will notify the preprocessing Observer on the main thread for each resulting item, which Observer then has to schedule tasks on a background thread for all those items. The main thread is involved too early in such flows which wastes time for the main thread.

A second shortcoming could be the lack of terminal event notifications included in the Observer API. Of course, one can define an union-class with the value, error and complete subtypes, just like RxJava's Notification object, to work around this problem. Needless to say, such wrapper will cost allocation and indirection.

Lastly, addListener/removeListener pairs compose relatively poorly and requires the operator to remember both the original LiveData and the Observer that was used when the link was established. An onSubscribe(Disposable) style cancellation support worked out much nicer in RxJava 2. Of course, having a similar method in Observer would break its lambda-friendly Single Abstract Method structure.

Reactive-Streams interoperation

Consuming a LiveData as Publisher

RxJava has infrastructure support to work with such classical Listener-style sources. Specifically, the create() operator on Observable and Flowable show such an example usage in its JavaDoc. Consequently, talking to LiveData is a straightforward routine:

LiveData<String> source = new MutableLiveData<String>();
Flowable<String> f = Flowable.<String>create(emitter -> {
   Observer<String> o = v -> {
       emitter.onNext(v);
   };

   source.observeForever(o);
   emitter.setCancellable(() -> 
       AndroidSchedulers.mainThread()
       .scheduleDirect(() -> 
            mld.removeObserver(o)
       )
   );
}, BackpressureStrategy.LATEST);

The Reactive-Streams bridge in LiveDataReactiveStreams, of course, has no access to Flowable and has to work out the Subscription management manually. Unfortunately, the toPublisher() implementation is wrong. If there was a public repository, I could contribute a correct one (without using RxJava of course).

Let's see the implementation piece by piece to demonstrate how not to write Publishers. (The file is marked as Apache 2.0 open source, should be fine to quote it here, right?)

    public static <T> Publisher<T> toPublisher(
            final LifecycleOwner lifecycle, final LiveData<T> liveData) {

        return new Publisher<T>() {
            boolean mObserving;
            boolean mCanceled;
            long mRequested;
            @Nullable
            T mLatest;

            @Override
            public void subscribe(final Subscriber<T> subscriber) {
                // continued below
            }
        };
    }

The first red flag comes from the mutable, non thread-safe state of the anonymous Publisher. Hot sources may have mutable fields but those are protected with proper synchronization primitives (see PublishProcessor for example). However, these fields should be part of the state of the Subscription that is sent to a Subscriber as Publishers are expected to be consumed by any number of Subscribers and interaction with one Subscriber should not affect the interactions with the others.

Of course, a Publisher can chose to only service a single Subscriber (as with UnicastProcessor), but for that, there is no field to let the subscribe() know there is already someone consuming the Publisher.

Next, the Observer instance is defined with a backpressure strategy of keeping the latest if the downstream is not ready.

    final Observer observer = new Observer() {
        @Override
        public void onChanged(@Nullable T t) {
            if (mCanceled) {
                return;
            }
            if (mRequested > 0) {
                mLatest = null;
                subscriber.onNext(t);
                if (mRequested != Long.MAX_VALUE) {
                    mRequested--;
                }
            } else {
                mLatest = t;
            }
        }
    };

Apart from the wrong shared state of the mCanceled, mRequested and mLatest, the code relies on the fact that both the onChanged invocation and the request() call happens on the main thread (see below). Having a per Subscriber atomic mRequested instead would save a trip to the main thread for an increment.

The next segment deals with the request() call from the downstream.

    subscriber.onSubscribe(new Subscription() {
        @Override
        public void request(final long n) {
            if (n < 0 || mCanceled) {
                return;
            }
            ArchTaskExecutor.getInstance().executeOnMainThread(new Runnable() {
                @Override
                public void run() {
                    if (mCanceled) {
                        return;
                    }
                    // Prevent overflowage.
                    mRequested = mRequested + n >= mRequested
                               ? mRequested + n : Long.MAX_VALUE;
                    if (!mObserving) {
                         mObserving = true;
                         liveData.observe(lifecycle, observer);
                    } else if (mLatest != null) {
                         observer.onChanged(mLatest);
                         mLatest = null;
                    }
                }
            });
        }

The first (smaller) problem is that non-positive requests should be rewarded with an onError(IllegalArgumentException) properly serialized with any onNext calls. Here, it is ignored. In addition, request() may be called from any thread and such call may not see mCanceled properly. The mObserving adds an odd delay to start observing the actual LiveData object. (Since most sequences will issue a request() almost immediately, I'm not sure what benefit this can give, however, there were voices pretty keen on such behavior recently who couldn't understand that to control the first invocation of request(), one would have to consume a Publisher directly with a Subscriber and not have any intermediate operators between the two.)

Lastly, let's see cancel():

    @Override
    public void cancel() {
        if (mCanceled) {
            return;
        }
        ArchTaskExecutor.getInstance().executeOnMainThread(new Runnable() {
            @Override
            public void run() {
                if (mCanceled) {
                    return;
                }
                if (mObserving) {
                    liveData.removeObserver(observer);
                    mObserving = false;
                }
                mLatest = null;
                mCanceled = true;
            }
       });
   }

Since removeObserver has to be called from the main thread (which by itself could be inconvenient), a task is scheduled which will also update mCanceled to true. Since cancel() can be invoked from any thread, a non-volatile mCanceled may not be visible. In addition, since the mCanceled is set to true on the main thread, there could be a significant delay between a cancel() and its effects, letting onChanged() calls still through. Setting a volatile mCanceled true right in the cancel() method could have the stopping effects much earlier.

(In fact, this was a problem with our early unsubscribeOn() where the Reactive-Streams TCK would complain about still receiving too events after it has cancelled the flow. The solution was to set a volatile cancelled flag that would be read by onNext and drop events until the upstream.cancel() got executed asynchronously.)

It is relatively simple to fix the mistakes here: combine Observer and Subscription into one class and move the mXXX fields into this class.

Exposing a Publisher as LiveData

The unfortunate lack of terminal events in LiveData forces us to deal with only onNext calls. There is no way to tell an Observer that no more events will come (unless agreeing upon an event container type) and then get rid of them en-masse (there is no clearObservers()).

With RxJava, this could be done relatively easily:

    MutableLiveData<String> mld2 = new MutableLiveData<>();
    Disposable d = Flowable.range(1, 5).map(Object::toString)
            .subscribe(mld2::postValue);

Of course, one has to manage the Disposable as usual and possibly decide, when to subscribe() to a cold source with respect to available Observers of the (Mutable)LiveData. A refCount-like behavior is possible but for that, one has to implement a custom LiveData object.

Apparently, there exist events on the LiveData class, onActive and onInactive, that will be called when the first Observer arrives and the last Observer leaves respectively. The ReactiveStreamsLiveData.fromPublisher() uses them in its implementation:

    private static class PublisherLiveData<T> extends LiveData<T> {
        private WeakReference<Subscription> mSubscriptionRef;
        private final Publisher mPublisher;
        private final Object mLock = new Object();

        PublisherLiveData(@NonNull final Publisher publisher) {
            mPublisher = publisher;
        }

        // the rest is below
}

The mSubscriptionRef holds the current active connection's Subscription and the mLock protects it so that a concurrent call to onSubscribe may not clash with the deactivation of the LiveData object. As we'll see in a short, this is not enough because the race can leave the upstream's Subscription from an old activation not cancelled and thus still posting to the LiveData object.

The onActivate event should subscribe to the source and relay its onNext events:

    @Override
    protected void onActive() {
        super.onActive();

        mPublisher.subscribe(new Subscriber<T>() {
            @Override
            public void onSubscribe(Subscription s) {
                // Don't worry about backpressure. If the stream is too noisy then
                // backpressure can be handled upstream.
                synchronized (mLock) {
                    s.request(Long.MAX_VALUE);
                    mSubscriptionRef = new WeakReference<>(s);
                }
            }

            @Override
            public void onNext(final T t) {
                postValue(t);
            }

            @Override
            public void onError(Throwable t) {
                synchronized (mLock) {
                    mSubscriptionRef = null;
                }
                // Errors should be handled upstream, so propagate as a crash.
                throw new RuntimeException(t);
            }

            @Override
            public void onComplete() {
                synchronized (mLock) {
                    mSubscriptionRef = null;
                }
            }
        });
    }

Here, apart from the mSubscriptionRef management, the only small mistake is the call to request() from within the lock itself. A bigger mistake is that onError should not throw.

The onInactive should cancel any existing Subscription if the Subscription representing that connection has been received:

    @Override
    protected void onInactive() {
        super.onInactive();
        synchronized (mLock) {
            WeakReference<Subscription> subscriptionRef = mSubscriptionRef;
            if (subscriptionRef != null) {
                Subscription subscription = subscriptionRef.get();
                if (subscription != null) {
                    subscription.cancel();
                }
                mSubscriptionRef = null;
            }
        }
    }

Unfortunately, the call to onSubscribe is not guaranteed to happen on the same thread the Publisher.subscribe() is invoked. Even though RxJava's Flowable almost always does so, other Reactive-Streams implementations are free to delay the call to onSubscribe and issue it from any thread. (We deal with such situations in Flowable via the deferred requesting/cancellation mechanism.)

Consider the following scenario. We know onActive and onInactive execute on the main thread thus they are serialized in respect to each other. Let's assume a Publisher will signal onSubscribe in a delayed manner when subscribed to. The main thread executes onActive that calls subscribe() on the Publisher. Then the main thread executes onInactive() after. Since the Publisher hasn't called onSubscribe yet, there is no mSubscriptionRef and thus there is nothing to cancel. The onSubscribe() is then invoked by the Publisher and it starts streaming events (to nobody practically). Now the main thread executes onActive again, triggering yet another subscription to the Publisher. After onSubscribe() is called again, the previous Subscription is overwritten and there are now twice as many onNext events posted to the LiveData (these could be duplicates in case the Publisher is hot/multicasting, or arbitrary items from various stages of a cold Publisher). Since the previous Subscription is lost, it is not possible to stop the first subscription other than letting it terminate on its own pace.

One possible solution is to host the Subscriber in an AtomicReference inside the custom LiveData, plus have that Subscriber store the incoming Subscription in an AtomicReference too. Therefore, an onInactive can get rid of the previous Subscriber and at the same time, instruct it to cancel the Subscription incoming through onSubscribe without requesting from it:

    static final class PublisherLiveData<T> extends LiveData<T> {
        final Publisher<T> mPublisher;
        final AtomicReference<SubscriberLiveData> mSubscriber;

        PublisherLiveData(@NonNull final Publisher<T> publisher) {
            mPublisher = publisher;
            mSubscriber = new AtomicReference<>();
        }

        @Override
        protected void onActive() {
            super.onActive();
            SubscriberLiveData s = new SubscriberLiveData();
            mSubscriber.set(s);
            mPublisher.subscribe(s);
        }

        @Override
        protected void onInactive() {
            super.onInactive();
            SubscriberLiveData s = mSubscriber.getAndSet(null);
            if (s != null) {
                s.cancel();
            }
        }

        final class SubscriberLiveData 
        extends AtomicReference<Subscription>
        implements Subscriber<T>, Subscription {
            @Override
            public void onSubscribe(Subscription s) {
                if (compareAndSet(null, s)) {
                    s.request(Long.MAX_VALUE);
                } else {
                    s.cancel();
                }
            }

            @Override
            public void onNext(T item) {
                postValue(item);
            }

            @Override
            public void onError(Throwable ex) {
                lazySet(this);
                mSubscriber.compareAndSet(this, null);
                ex.printStackTrace();
            }

            @Override
            public void onComplete() {
                lazySet(this);
                mSubscriber.compareAndSet(this, null);
            }

            @Override
            public void cancel() {
                Subscription s = getAndSet(this);
                if (s != null && s != this) {
                    s.cancel();
                }
            }

            @Override
            public void request(long n) {
                // never called
            }
        }
    }

Although onActive and onInactive would update mSubscriber from the main thread, the atomics is preferable since we want to set the current Subscriber to null once the Publisher terminated, thus playing nice with the GC. (For comparison, the refCount in RxJava is more involved because connection and disconnection may be triggered from any thread at any time.)

Conclusion

The LiveData class can be considered a classical (gen 0) reactive-push component which has practically one purpose on Android: notify its Observers of data (changes) on the main thread. If the typical thread crossing in the app is mainly from a background thread to the main thread, LiveData may be just enough to suit one's needs. It's one step above Agera after all since LiveData actually transmits data to be acted upon, not just "something happened" signals. However, if possibly failing and finite sources are present, which also have to be coordinated, transformed and kept on background thread(s) as long as possible, I believe RxJava is a better choice.

Nevertheless, there are standard interoperation bridges provided with LiveData which would allow it to work with or pose as a Reactive-Streams Publisher, but unfortunately, the implementation of this bridge has concurrency flaws that could be easily fixed by a PR, if there was an open repository for the library (because you'd want to post unit tests as well).

Introduction

In many tutorials and explanations, it has been said that having multiple subscribeOn()s has no effect and only the one closest to the source wins. I often tell this with the wording "no practical effect". However, it is possible to demonstrate the effects of multiple subscibeOn()s that have some actual effects.

What is subscribeOn again?

The most precise definition of this operator I can formulate is as follows:

• subscribeOn changes where (on what thread) the (side) effects of calling subscribe() on the parent/upstream Observable (Flowable, Single, etc.) happen.

So what are these subscription (side) effects look like in code?

Observable.create(emitter -> {
    for (int i = 0; i < 10; i++) {
        emitter.onNext(i + ": " + Thread.currentThread().getName());
    }
    emitter.onComplete();
})
.subscribeOn(Schedulers.io())
.blockingSubscribe(System.out::println);

// Prints:
// -------
// 0: RxCachedThreadScheduler-1
// 1: RxCachedThreadScheduler-1
// 2: RxCachedThreadScheduler-1
// 3: RxCachedThreadScheduler-1
// 4: RxCachedThreadScheduler-1
// 5: RxCachedThreadScheduler-1
// 6: RxCachedThreadScheduler-1
// 7: RxCachedThreadScheduler-1
// 8: RxCachedThreadScheduler-1
// 9: RxCachedThreadScheduler-1

In this example, the effect of subscribing is that the body of the ObservableOnSubscribe starts running on the thread provided via the io()Scheduler.

Applying yet another subscribeOn after the first one won't change what is printed to the console.

Most source-like operators, such as create(), fromCallable(), fromIterable(), do have subscription side-effects as they often start emitting event(s) immediately.

Most instance operators, such as map(), filter(), take(), don't have subscription side-effects on their own and just subscribe() to their upstream.

Instance operators with subscription side-effects

However, there are a couple of instance operators that do have subscription side-effects. Specifically, any operator that offers a way to specify a per subscriber initial state via some callback.

Observable.create(emitter -> {
    for (int i = 0; i < 10; i++) {
        emitter.onNext(i + ": " + Thread.currentThread().getName());
    }
    emitter.onComplete();
})
.subscribeOn(Schedulers.io())
.collect(() -> {
       List<String> list = new ArrayList<>();
       list.add(Thread.currentThread().getName());
       return list;
    }, 
    (a, b) -> a.add(b)
)
.blockingSubscribe(list -> list.forEach(System.out::println));

This will print the following:

main
0: RxCachedThreadScheduler-1
1: RxCachedThreadScheduler-1
2: RxCachedThreadScheduler-1
3: RxCachedThreadScheduler-1
4: RxCachedThreadScheduler-1
5: RxCachedThreadScheduler-1
6: RxCachedThreadScheduler-1
7: RxCachedThreadScheduler-1
8: RxCachedThreadScheduler-1
9: RxCachedThreadScheduler-1

The first printout, main, is due to the fact that the blockingSubscribe() is subscribing on the main thread and the collect() operator performs its initialization side-effect on the main thread as well.

Now let's see what happens if we add subscribeOn() after collect:

Observable.create(emitter -> {
    for (int i = 0; i < 10; i++) {
        emitter.onNext(i + ": " + Thread.currentThread().getName());
    }
    emitter.onComplete();
})
.subscribeOn(Schedulers.io())
.collect(() -> {
       List<String> list = new ArrayList<>();
       list.add(Thread.currentThread().getName());
       return list;
    }, 
    (a, b) -> a.add(b)
)
.subscribeOn(Schedulers.computation())  // <---------------------------------------
.blockingSubscribe(list -> list.forEach(System.out::println));

will print the following:

RxComputationThreadPool-1
0: RxCachedThreadScheduler-1
1: RxCachedThreadScheduler-1
2: RxCachedThreadScheduler-1
3: RxCachedThreadScheduler-1
4: RxCachedThreadScheduler-1
5: RxCachedThreadScheduler-1
6: RxCachedThreadScheduler-1
7: RxCachedThreadScheduler-1
8: RxCachedThreadScheduler-1
9: RxCachedThreadScheduler-1

Because subscribeOn(Schedulers.computation()) performs its subscribe() call on one of the computation thread, the subscription side-effect of creating the list (for collecting Strings from upstream) is happening on that thread. Then, subscribeOn(Schedulers.io()) again switches threads to an io thread where the subscription effect of emitting more Strings begins.

Conclusion

Introduction

There are situations when the same elements of a source should be dispatched to multiple consumers. Certainly, if the source supports multiple subscribers and is deterministic (such as our previous async range), one can just instantiate the flow multiple times. However, if the source doesn't support multiple subscribers or each subscription ends up being unique and/or non-deterministic, that simple approach doesn't work anymore.

We need a way to have a single realization of the (upstream) source yet allow multiple consumers. Since we are dealing with Flow.Publishers that require backpressure management, such intermediate solution has to coordinate requests from its Flow.Subscribers in addition to handling the dynamic subscription and unsubscription (cancellation) of said Flow.Subscribers while the flow is active. Enter, MulticastProcessor.

Flow.Processor recap

What is a Processor? By definition, it is a combination of a Flow.Publisher and a Flow.Subscriber, i.e., it can act as a source and can be subscribed to via subscribe() as well as the processor itself can be used with somebody else's Flow.Publisher.subscribe(). It has a mixed history as the idea comes from the original Observer pattern (i.e., java.util.Observable) and Rx.NET's Subject that allows dispatching signals to multiple Observers in an imperative (and synchronous) fashion.

The Flow.Processor in Java 9 defines two type arguments, one for its input side (Flow.Subscriber) and one for its output side (Flow.Publisher). The idea behind it was that a Flow.Processor can act as a transformation step between an upstream and a downstream. However, such transformation often mandates the Flow.Processor implementation only accepts a single Flow.Subscriber during its entire lifetime. Since the implementation has to follow the Reactive Streams specification nonetheless, this adds a lot of overhead to the flow. As demonstrated in previousposts, when a flow is realized, there are only one subscriber per stage in it and thus the Flow.Processor overhead can be completely avoided.

Still, restricting the number of Flow.Subscribers has some use in operators such as groupBy() and window() where the inner Flow.Publishers are actually Flow.Processors that work in an unicast mode: holding onto elements until the consumer is ready to receive them while letting the main source progress with other groups/windows.

Of course, when the items in these inner Flow.Publishers are required by multiple processing stages, a separate multicasting stage, in the form of a Flow.Processor can come in quite handy.

MulticastProcessor

The first step is to define the skeleton of the class.

public final class MulticastProcessor<T> implements Flow.Processor<T, T> {

    final int prefetch;

    final SpscArrayQueue<T> queue;    

    /* Some other fields, see below. */

    public MulticastProcessor(int prefetch) {
        this.prefetch = prefetch;
        this.queue = new SpscArrayQueue<>(prefetch);
    }

    @Override
    public void subscribe(Flow.Subscriber<? super T> subsciber) {
        // TODO implement
    }

    @Override
    public void onSubscribe(Flow.Subscription s) {
        // TODO implement
    }

    @Override
    public void onNext(T item) {
        // TODO implement
    }

    @Override
    public void onError(Throwable t) {
        // TODO implement
    }

    @Override
    public void onComplete() {
        // TODO implement
    }     
}

The class has to implement 5 public methods from the Flow.Processor interface. Since the intended usage is to have multiple Flow.Subscribers and be able to subscribe the MulticastProcessor itself to some other Flow.Publisher, both the upstream's Flow.Subscription and any Flow.Subscribers have to be tracked. The input and output types are the same as there won't be any item transformation happening in this particular Flow.Processor implementation.

We need a queue to store the upstream values because some of the downstream Flow.Subscribers may not be ready to receive items and we don't want any discontinuity in item delivery for them. A JCTools single-producer single-consumer queue suffices because the items coming from the upstream are guaranteed to be single thread at a time and the dispatch logic consuming the queue will be serialized as well by MulticastProcessor.

The two sides can operate asynchronously to each other, we will need some atomics around the state of the MulticastProcessor:

    Flow.Subscription upstream;

    static final VarHandle UPSTREAM = 
        VH.find(MethodHandles.lookup(), MulticastProcessor.class, 
"upstream", Flow.Subscription.class);

    MulticastSubscription<T>[] subscribers = EMPTY;
    static final VarHandle SUBSCRIBERS = 
        VH.find(MethodHandles.lookup(), MulticastProcessor.class, 
"subscribers", MulticastSubscription[].class);

    int wip;
    static final VarHandle WIP = 
        VH.find(MethodHandles.lookup(), MulticastProcessor.class, 
"wip", int.class);

    static final MulticastSubscription[] EMPTY = new MulticastSubscription[0];

    static final MulticastSubscription[] TERMINATED = new MulticastSubscription[0];

    volatile boolean done;
    Throwable error;

    int consumed;

The upstream stores the Flow.Subscription received from the source to be multicasted. Flow.Subscribers can come and go independently and as such, a typical copy-on-write array structure with terminal state can be employed. If the terminal state has been reached, any subsequent subscribe() attempt will either call onError() or onComplete() on the Flow.Subscriber, depending on how the MulticastProcessor itself has been terminated (i.e., the error field is null or not). The VH utility class has been introduced in a previous post.

One could ask, why do we need the wip field? The upstream will invoke onNext synchronously on the MulticastProcessor so a simple loop over the current array of Flow.Subscribers should suffice, right? Not exactly. Becase we want every Flow.Subscriber to receive the same items in the same order without skipping, they will move in a so-called lockstep-mode. If a Flow.Subscriber is not ready to receive, the others won't receive data even if they could. Once this single Flow.Subscriber is ready to receive (it called request()), we have to try and drain the queue for the available items, which can happen from a different thread than the upstream is calling onNext from. Second, if a Flow.Subscriber decides to cancel and was preventing the others from making progress because of not requesting, its leave may enable the others to receive available items right on.

Due to the prefetching nature of MulticastProcessor, we won't be able to use the TERMINAL indicator as a completion indicator because the queue may hold non-dispatched elements when an onComplete() or onError() arrives. Therefore, a separate done flag is employed which tells the drain logic later on to not expect any further items from the queue.

Each individual Flow.Subscriber needs its own state so that its request amounts can be tracked, the MulticastSubscription class is employed:

static final class MulticastSubscription<T> 
extends AtomicLong
implements Flow.Subscription {

    final Flow.Subscriber<? super T> downstream;

    final MulticastProcessor<T> parent;

    long emitted;

    MulticastSubscription(Flow.Subscriber<? super T> downstream, 
            MulticastProcessor<T> parent) {
        this.downstream = downstream;
        this.parent = parent;
    }

    @Override
    public void request(long n) {
        if (n <= 0L) {
            cancel();
            // for brevity, serialization is omitted here, see HalfSerializer.
            downstream.onError(new IllegalArgumentException());
        } else {
            for (;;) {
                long current = (long)getAcquire();
                if (current == Long.MIN_VALUE || current == Long.MAX_VALUE) {
                    break;
                }
                long next = current + n;
                if (next < 0L) {
                    next = Long.MAX_VALUE;
                }
                if (compareAndSet(current, next)) {
                    parent.drain();
                    break;
                }
            }
        }
    }

    @Override
    public void cancel() {
        if ((long)getAndSet(Long.MIN_VALUE) != Long.MIN_VALUE) {
            parent.remove(this);
            parent.drain();
        }
    }
}

The MulticastSubscription takes a reference to the actual downstream Flow.Subscriber and the parent MulticastProcessor to trigger removal and item draining. The emitted field counts how many times downstream.onNext() was invoked to know the remaining items this particular subscription can take. The AtomicLong tracks the total downstream requested amounts so the remaining can be calculated with getAcquire() - emitted (should happen on a single thread). The Long.MIN_VALUE indicates the downstream has cancelled the subscription.

The parent MulticastProcessor requires a few services to handle the various MulticastSubscription instances: drain() to start emitting items if possible, remove() to let a particular Flow.Subscriber go. We'll need its pair, add(), to register a new Flow.Subscriber to work with. First, let's see add() and remove():

    boolean add(MulticastSubscription<T> ms) {
        for (;;) {
            MulticastSubscription<T>[] current = 
                (MulticastSubscription<T>[])SUBSCRIBERS.getAcquire(this);
            if (current == TERMINATED) {
                return false;
            }
            int n = current.length;
            MulticastSubscription<T>[] next = new MulticastSubscription[n + 1];
            System.arraycopy(current, 0, next, 0, n);
            next[n] = ms;
            if (SUBSCRIBERS.compareAndSet(this, current, next)) {
                return true;
            }
        }
    }

The add() method should look familiar. We read the current array of MulticastSubscriptions, see if it is the terminal indicator (in which case we return false indicating the MulticastProcessor is in a terminal state an no new Flow.Subscribers are accepted) and if not, we copy and extend the array to incorporate the incoming new MulticastSubscription. The compareAndSet then ensures the changes are applied atomically and in case of a concurrent change, the process is retried.

    void remove(MulticastSubscription<T> ms) {
        for (;;) {
            MulticastSubscription<T>[] current = 
                (MulticastSubscription<T>[])SUBSCRIBERS.getAcquire(this);
            int n = current.length;
            if (n == 0) {
                break;
            }

            int j = -1;
            for (int i = 0; i < n; i++) {
                 if (ms == current[i]) {
                     j = i;
                     break;
                 }
            }

            if (j < 0) {
                 break;
            }

            MulticastSubscription<T>[] next;
            if (n == 1) {
                next = EMPTY;
            } else {
                next = new MulticastSubscription[n + 1];
                System.arraycopy(current, 0, next, 0, j);
                System.arraycopy(current, j + 1, next, j, n - j - 1);
            }

            if (SUBSCRIBERS.compareAndSet(this, current, next)) {
                break;
            }
        }
    }

The remove() method is fairly typical again. We read the current array and if it is empty (due to having no subscribers or being terminated) we quit. Otherwise we search for the index of the MulticastSubscription and if found, a new array is created and items copied around the index into it. The compareAndSet again ensures proper visibility and retry in case of concurrent state changes. If the current array happens to be one element long, we avoid creating a new empty array and simply reuse the shared static EMPTY instance.

We can now implement the subscribe() method:

    @Override
    public void subscribe(Flow.Subscriber<? super T> subscriber) {
        Objects.requireNonNull(subscriber, "subscriber == null");

        MulticastSubscription<T> ms = new MulticastSubscription<T>(subscriber, this);
        subscriber.onSubscribe(ms);

        if (add(ms)) {
            if ((long)ms.getAcquire() == Long.MIN_VALUE) {
                remove(ms);
            }
            drain();
        } else {
            if ((long)ms.getAcquire() != Long.MIN_VALUE) {
               Throwable ex = error;
               if (ex == null) {
                   ms.downstream.onComplete();
               } else {
                   ms.downstream.onError(ex);
               }
            }
        }
    }

First, we verify subscriber is not null. Then, an unique instance of MulticastSubscription is associated with the Flow.Subscriber. We try to register this with the MulticastProcessor via add(). This may fail in case the MulticastProcessor reached its terminal state in which case we emit the terminal event, unless the Flow.Subscriber has cancelled in the meantime (the request accounting shows Long.MIN_VALUE). This can happen concurrently (remember, we sent out the Flow.Subscription via onSubscribe before) with the add() call so the remove() triggered by MulticastSubscription.cancel() may not find itself in the array: calling it again ensures that the reference is not retained in this corner case. In either case, the additional call to drain() will ensure items are dispatched as soon as possible to the (remaining) parties.

One could ask, why not call add() first and then call onSubscribe() since as long as there is no request() from downstream, the visibility of the MulticastSubscription shouldn't be a problem, right? Not exactly. The Reactive Streams specification allows emitting onError or onComplete without prior request; in this case, since the reference to the MulticastSubscription would be visible sooner, the onError/onComplete could be invoked before onSubscribe which is a violation of the specification (and will likely lead to NullPointerException down the line). Reversing the call order avoids this violation but requires the aforementioned add()-if cancelled-remove() logic to avoid leaking cancelled subscriptions.

Let's see the implementation of the remaining 4 public methods:

    @Override
    public void onSubscribe(Flow.Subscription s) {
        Objects.requireNonNull(s, "s == null");

        if (UPSTREAM.compareAndSet(this, null, s)) {
            s.request(prefetch);
        } else {
            s.cancel();
        }
    }

    @Override
    public void onNext(T item) {
        Objects.requireNonNull(item, "item == null");

        queue.offer(item);
        drain();
    }

    @Override
    public void onError(Throwable t) {
        Objects.requireNonNull(t, "t == null");

        error = t;
        done = true;
        drain();
    }

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

The onSubscribe atomically sets the Flow.Subscription if not already set, cancels it otherwise. The other onXXX methods queue the item/error/done indicator and call drain() to dispatch events if possible. The Reactive Streams specification requires throwing NullPointerException if the input parameters are null.

Lastly, we can now implement the drain() method, piece by piece for better clarity. It is implemented as our typical queue-drain loop:

    void drain() {
        if ((int)WIP.getAndAdd(this, 1) != 0) {
            return;
        }

        int missed = 1;

        for (;;) {

            // TODO implement

            missed = (int)WIP.getAndAdd(this, -missed) - missed;
            if (missed == 0) {
                break;
            }
        }
    }

Usually, the inner loop starts with reading the requested amount, however, we don't have a single requested amount in this case: we have to find out the smallest number of items that can be emitted to all current Flow.Subscribers:

    // for (;;) {

    MulticastSubscription<T>[] current = 
        (MulticastSubscription<T>[])SUBSCRIBERS.getAcquire(this);
    int n = current.length;

    if (n != 0) {
        long requested = Long.MAX_VALUE;
        for (MulticastSubscription<T> ms : current) {
            long r = ms.getAcquire();
            if (r != Long.MIN_VALUE) {
                requested = Math.min(requested, r - ms.emitted);
            }
        }

        // TODO try emitting handle
    }

    // }

Remember, the amount of items a particular MulticastSubscription can receive is its total requested amount minus the emitted items to it so far.

Next, we try to emit any available item from the queue:

    // if (n != 0) { ...

    while (requested != 0) {

        boolean d = done;
        T v = queue.poll();
        boolean empty = v == null;

        if (d && empty) {
            // TODO terminal case
            return;
        }

        if (empty) {
            break;
        }

        for (MulticastSubscription<T> ms : current) {
            if ((long)ms.getAcquire() != Long.MIN_VALUE) {
                ms.downstream.onNext(v);
                ms.emitted++;
            }
        }

        requested--;

        if (++consumed == (prefetch >> 1)) {
            upstream.request(consumed);
            consumed = 0;
        }
    }

    // TODO no request - no more items case

    // }

In this inner loop, which is quit if requested reaches zero, we try to get the next available item from the queue and if found, we emit it to non-cancelled Flow.Subscribers inside the MulticastSubscriptions. Since the MulticastProcessor works in a stable-prefetch mode, we have to track the consumption via consumed and request more after a certain amount of items have been taken. In this case, we request half of the prefetch amount if we reach that amount.

One case to be implemented is when the upstream has terminated and we emitted all queued items:

    // if (d && empty) {
    Throwable ex = error;
    current = (MulticastSubscription<T>[])SUBSCRIBERS.getAndSet(this, TERMINATED);
    if (ex == null) {
        for (MulticastSubscription<T> ms : current) {
             if ((long)ms.getAcquire() != Long.MIN_VALUE) {
                ms.downstream.onComplete();
             }
        }
    } else {
        for (MulticastSubscription<T> ms : current) {
             if ((long)ms.getAcquire() != Long.MIN_VALUE) {
                ms.downstream.onError(ex);
             }
        }
    }
    return;
    // }

We atomically swap in the terminal array which gives us the last array of MulticastSubscriptions. We have to do this because more subscribers could have arrived since the previous SUBSCRIBERS.getAcquire() and would be left non-terminated and hanging forever. The non-cancelled Flow.Subscribers are terminated with onError or onComplete based on the error field's contents.

Lastly, we have to deal with the case when none of the MulticastSubscribers have requested yet (or run out of requests) and the upstream has terminated without any further values stored in the queue:

    // while (requested != 0) {
    // ...
    // }

    if (requested == 0) {
        boolean d = done;
        boolean empty = queue.isEmpty();

        if (d && empty) {
            Throwable ex = error;
            current = (MulticastSubscription<T>[])SUBSCRIBERS.getAndSet(this, TERMINATED);
            if (ex == null) {
                for (MulticastSubscription<T> ms : current) {
                     if ((long)ms.getAcquire() != Long.MIN_VALUE) {
                        ms.downstream.onComplete();
                     }
                }
            } else {
                for (MulticastSubscription<T> ms : current) {
                     if ((long)ms.getAcquire() != Long.MIN_VALUE) {
                        ms.downstream.onError(ex);
                     }
                }
            }
            return;
        }
    }

Unsurprisingly, this almost looks like as the fist part of the while (requested != 0) loop, except we don't poll the item if there is one as we can't put it back and would lose it.

Testing with the Reactive Streams TCK

Now that the code is complete, we want to make sure the implementation properly implements the Reactive Streams specification. We can do this via the Test Compatibility Kit of it, however, since Reactive Streams is a separate library predating the Java 9 Flow API, we need some adapters to utilize the TCK. Fortunately, the upcoming 1.0.2 both features such adapters and a Java 9 Flow-based port of the TCK itself. For that, one only needs to import the tck-flow (and its dependencies if doing it manually!):

    compile "org.reactivestreams:reactive-streams-tck-flow:1.0.2-RC2"

When writing this blog post, the release candidate version is only available, hence the -RC2 postfix.

The TCK contains the IdentityFlowProcessorVerification we must implement in our tests. Note that the TCK uses TestNG instead of JUnit.

@Test
public class MulticastProcessorTest extends IdentityFlowProcessorVerification<Integer> {

    public MulticastProcessorTest() {
        super(new TestEnvironment(50));
    }

    @Override
    public ExecutorService publisherExecutorService() {
        return Executors.newCachedThreadPool();
    }

    @Override
    public Integer createElement(int element) {
        return element;
    }

    @Override
    protected Flow.Publisher<Integer> createFailedFlowPublisher() {
        MulticastProcessor mp = new MulticastProcessor<>(128);
        mp.onError(new IOException());
        return mp;
    }

    @Override
    protected Flow.Processor<Integer, Integer> createIdentityFlowProcessor(int bufferSize) {
        return new MulticastProcessor<>(bufferSize);
    }
}

At first glance, there is no complication in writing the test. We have to specify a TestEnvironment with a default timeout (50 milliseconds should be enough), we have t specify a helper ExecutorService from which test items will be emitted. Those test items will be created via createElement so the target type of the Flow.Processor to be tested can be matched (plain integer in this case). The createFailedFlowPublisher verifies an immediately failing Flow.Publisher (remember Flow.Processor is one) and the createIdentityFlowProcessor should return a fresh and empty processor instance.

If we run the test, 3 of the 68 tests will fail (plus 24 gets ignored as they don't apply or can't be properly verified as of now):

1. required_spec104_mustCallOnErrorOnAllItsSubscribersIfItEncountersANonRecoverableError
2. required_mustRequestFromUpstreamForElementsThatHaveBeenRequestedLongAgo
3. required_exerciseWhiteboxHappyPath

From these, failures 1) and 2) are due to what the TCK expects the Flow.Processor behaves by default: independent requests from independend Flow.Subscribers must reach the upstream and produce elements independently. Unfortunately, our MulticastProcessor uses a prefetch and lockstep mechanisms, thus if only one Flow.Subscriber requested and the others didn't, nobody gets items. Fortunately, the 1.0.2 has been enhanced in this regard and one only has to override a method in the IdentityFlowProcessorVerification class:

    @Override
    public boolean doesCoordinatedEmission() {
        return true;
    }

This will instruct the methods 1) and 2) to try a different request-emission pattern ant they should now pass.

Now lets look at 3) which fails with a different exception: "Did not receive expected cancelling of upstream subscription within 50 ms". Indeed, looking at the code, there upstream's Flow.Subscription.cancel() is never called on a normal path (only when onSubscribe is called twice).

What's happening? This is another default expectation of the TCK: when all Flow.Subscribers cancel the Flow.Processor implementation is expected to cancel its upstream promptly. This may not be the desired business behavior as one may want to keep dispatching even if Flow.Subscribers come and go and temporarily, there are no downstream to send events to. Unfortunately, there is no switch for this in the TCK and we have to modify the MulticastProcessor for anticipate this expectation. This is similar to the concept and behavior of refCount() in RxJava.

Luckily, it requires only a small set of changes. We can detect when all the Flow.Subscribers have cancelled conveniently in the remove() method and instead of setting the subscribers array to EMPTY, we set it to TERMINATED and cancel() the upstream instance:

    void remove(MulticastSubscription<T> ms) {
        for (;;) {

            // the find index is omitted for brevity

            MulticastSubscription<T>[] next;
            if (n == 1) {
                if (UPSTREAM.getAcquire(this) == null) {
                    next = EMPTY;
                } else {
                    next = TERMINATED;
                }
            } else {
                next = new MulticastSubscription[n + 1];
                System.arraycopy(current, 0, next, 0, j);
                System.arraycopy(current, j + 1, next, j, n - j - 1);
            }

            if (SUBSCRIBERS.compareAndSet(this, current, next)) {
                if (next == TERMINATED) {
                    ((Flow.Subscription)UPSTREAM.getAcquire(this)).cancel();
                }
                break;
            }
        }
    }

This change will allow Flow.Subscribers to come and go until the MulticastPublisher is subscribed to another Flow.Publisher without terminating the MulticastPublisher itself and otherwise will cancel the upstream on a live one.

Conclusion

In this post, I showed how a multicasting Flow.Processor can be implemented and then tested with the Reactive Streams TCK. In fact, one reason I was holding off the demonstration of this MulticastProcessor is because of the limits of the TCK itself: prior to 1.0.2-RC2, now available from Maven, there was no standard adapter and no support for a lockstepping Flow.Publisher implementation. Now that I contributedboth and the maintainers released it after some time, we can now use it and experiment with various Java 9-based Flow.Processor and Flow.Publisher implementations.

We have applied the usual patterns: copy-on-write with terminal state, queue-drain and request accounting that should be quite familiar for veteran readers of this blog. Once the test indicated issues, we could simply adjust the logic to match its expectations by tweaking the nice thread-safe and lock-free logic of it.