Java 9 Flow API: Multicasting via a Processor

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 previous posts, 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 contributed both 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.

Java 9 Flow API: taking and skipping

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 VarHandle REQUEST_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.

Java 9 Flow API: arbitration and concatenation

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 CURRENT VarHandle 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.

Now that the infrastructure is ready, let's implement a couple of operators.

Concatenating an array of Flow.Publishers

Perhaps the simplest operator we could write on top of the SubscriptionArbiter is the concat() operator. It consumes one Flow.Publisher after another in a non-overlapping fashion until all of them have completed.

@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 source Flow.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.