Skip to content

concurrent-api: preserve demand across overlapping subscription switches #3677

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tuannx wants to merge 2 commits into
base: main
Choose a base branch
from
Open

concurrent-api: preserve demand across overlapping subscription switches #3677

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4419f3e
concurrent-api: reproduce lost and duplicated demand on resubscribe
tuannx Oct 6, 2026
506ce73
concurrent-api: preserve demand across overlapping subscription switches
tuannx Oct 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
210 changes: 87 additions & 123 deletions ...lk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/SequentialSubscription.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright © 2018-2019 Apple Inc. and the ServiceTalk project authors
* Copyright © 2018-2026 Apple Inc. and the ServiceTalk project authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand All @@ -19,10 +19,17 @@
import io.servicetalk.concurrent.PublisherSource.Subscription;
import io.servicetalk.concurrent.internal.FlowControlUtils;

import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import java.util.concurrent.atomic.AtomicLongFieldUpdater;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import javax.annotation.Nullable;

import static io.servicetalk.concurrent.internal.ConcurrentUtils.releaseLock;
import static io.servicetalk.concurrent.internal.ConcurrentUtils.tryAcquireLock;
import static io.servicetalk.concurrent.internal.EmptySubscriptions.EMPTY_SUBSCRIPTION_NO_THROW;
import static io.servicetalk.concurrent.internal.SubscriberUtils.isRequestNValid;
import static io.servicetalk.concurrent.internal.ThrowableUtils.catchUnexpected;
import static io.servicetalk.utils.internal.ThrowableUtils.throwException;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.atomic.AtomicLongFieldUpdater.newUpdater;

Expand All @@ -39,18 +46,24 @@
* {@link #switchTo(Subscription)} will be cancelled.
*/
final class SequentialSubscription implements Subscription {
private static final long SWITCHING = -1;
private static final long REQUESTED = -2;
private static final long CANCELLED = -3;
private static final AtomicLongFieldUpdater<SequentialSubscription> requestedUpdater =
newUpdater(SequentialSubscription.class, "requested");
private static final AtomicLongFieldUpdater<SequentialSubscription> sourceRequestedUpdater =
newUpdater(SequentialSubscription.class, "sourceRequested");
private static final AtomicIntegerFieldUpdater<SequentialSubscription> emittingUpdater =
AtomicIntegerFieldUpdater.newUpdater(SequentialSubscription.class, "emitting");
private static final AtomicReferenceFieldUpdater<SequentialSubscription, Subscription> pendingSubscriptionUpdater =
AtomicReferenceFieldUpdater.newUpdater(SequentialSubscription.class, Subscription.class,
"pendingSubscription");

// Only the drain owner changes the active subscription and its demand accounting.
private Subscription subscription;
private long sourceRequested;
// Subscriber signals are serialized. Publishing the next subscription also publishes this count to the drain.
private long sourceEmitted;
private volatile long requested;
private volatile long sourceRequested;
private volatile int emitting;
private volatile boolean cancelled;
@Nullable
private volatile Subscription pendingSubscription;

/**
* New instance.
Expand All @@ -70,142 +83,93 @@ final class SequentialSubscription implements Subscription {

@Override
public void request(long n) {
final long currRequested;
if (isRequestNValid(n)) {
currRequested = requestedUpdater.accumulateAndGet(this, n,
requestedUpdater.accumulateAndGet(this, n,
FlowControlUtils::addWithOverflowProtectionIfNotNegative);
} else {
currRequested = sanitizeInvalidRequestN(n);
requested = currRequested;
}

for (;;) {
final long currSourceRequested = sourceRequested;
if (currSourceRequested == CANCELLED) {
break;
} else if (currSourceRequested < 0) {
assert currSourceRequested == SWITCHING || currSourceRequested == REQUESTED;
if (sourceRequestedUpdater.compareAndSet(this, currSourceRequested, REQUESTED)) {
break;
}
} else {
// We must read the subscription BEFORE the CAS (which involves a read barrier). This ensures if the
// sourceRequested value is visible then the subscription (which may have been switched by
// another thread) is also visible.
final Subscription currSubscription = subscription;
if (isRequestNValid(currRequested)) {
// sourceRequested ...[delta]... requested
final long delta = currRequested - currSourceRequested;
if (sourceRequestedUpdater.compareAndSet(this, currSourceRequested, currSourceRequested + delta)) {
// sourceRequested is either monotonically increasing, or set to an invalid value
// (e.g. negative) if a Subscription switch is on going and atomically set to requestN to
// preserve the monotonic increasing property. If the CAS worked that means the value of
// subscription before will be visible if there was previously a switch. We also know there is
// no concurrent interaction on the subscription because currSourceRequested is known not to be
// SWITCHING, and the value would have increased.
if (delta != 0) {
currSubscription.request(delta);
}
break;
}
} else if (sourceRequestedUpdater.compareAndSet(this, currSourceRequested, CANCELLED)) {
currSubscription.request(currRequested);
break;
}
}
requested = sanitizeInvalidRequestN(n);
}
drain();
}

@Override
public void cancel() {
final Subscription currSubscription = subscription;
final long currSourceRequested = sourceRequestedUpdater.getAndSet(this, CANCELLED);
// To avoid concurrent invocation with the switch thread we defer to that thread to cancel.
if (currSourceRequested >= 0) {
currSubscription.cancel();
}
cancelled = true;
drain();
}
Comment on lines 96 to 99

@bryce-anderson bryce-anderson Oct 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I worry a little bit about this not being able to trigger cancel on a reentrant call. Eg, we're draining and that triggers a cancel call synchronously. Could we detect if this thread is the one draining, and if so, fire the cancel signal? In my minds eye, this can be done by setting a private Thread drainingThread field on this class that gets set when we successfully enter the drain() loop, and if cancel() fails to enter the drain loop (drain can return a boolean signaling whether we were denied by contention or not) we can check if this thread is the current drainingThread, and if so it's a reentrant call and safe to call subscription.cancel().

@idelpivnitskiy idelpivnitskiy Oct 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1, I simple reproducer would be fromIterable(it).concat(empty()) with a subscriber that cancels in onNext, it keeps draining the iterator (passes on main). Note that drainingThread proposal would only fix the same-thread case. A cancel() from another thread still waits behind the drain owner. If request() blocks, e.g. fromBlockingIterable in hasNext(timeout), the cancel that would unblock it never reaches it. Could cancel() bypass the lock and call subscription.cancel() directly, like the old implementation?

Here are examples of tests that fail on this PR and pass on main:

  1. ConcatPublisherTest (end-to-end, re-entrant cancel)
@Test
void cancelFromOnNextStopsSynchronousFirstSource() {
    final AtomicInteger pulled = new AtomicInteger();
    // Bounded so a regression fails instead of hanging: an infinite source would never stop emitting.
    final Publisher<Integer> p = fromIterable(() -> new Iterator<Integer>() {
        @Override
        public boolean hasNext() {
            return pulled.get() < 1000;
        }

        @Override
        public Integer next() {
            return pulled.incrementAndGet();
        }
    }).concat(empty());
    final List<Integer> received = new ArrayList<>();
    toSource(p).subscribe(new Subscriber<Integer>() {
        @Nullable
        private Subscription subscription;

        @Override
        public void onSubscribe(final Subscription s) {
            subscription = s;
            s.request(Long.MAX_VALUE);
        }

        @Override
        public void onNext(@Nullable final Integer item) {
            received.add(item);
            assert subscription != null;
            subscription.cancel();
        }

        @Override
        public void onError(final Throwable t) {
        }

        @Override
        public void onComplete() {
        }
    });
    assertThat(received, contains(1));
    assertThat("Items pulled from the source after cancel", pulled.get(), is(1));
}
  1. SequentialSubscriptionTest (re-entrant and cross-thread cancel)
@Test
void reentrantCancelDuringRequestCancelsActiveSubscription() {
    final AtomicBoolean cancelled = new AtomicBoolean();
    doAnswer(invocation -> {
        cancelled.set(true);
        return null;
    }).when(s1).cancel();
    doAnswer(invocation -> {
        s.cancel();
        // A synchronous source, e.g. Publisher.fromIterable, stops its emission loop only when it observes cancel.
        assertThat("Cancel did not reach the active subscription", cancelled.get(), is(true));
        return null;
    }).when(s1).request(anyLong());
    s.request(MAX_VALUE);
}

@Test
void cancelWhileRequestBlocksCancelsActiveSubscription() throws Exception {
    final CountDownLatch requested = new CountDownLatch(1);
    final CountDownLatch cancelled = new CountDownLatch(1);
    doAnswer(invocation -> {
        cancelled.countDown();
        return null;
    }).when(s1).cancel();
    doAnswer(invocation -> {
        requested.countDown();
        // A blocking source, e.g. Publisher.fromBlockingIterable, is unblocked only by cancel.
        cancelled.await(DEFAULT_TIMEOUT_SECONDS, SECONDS);
        return null;
    }).when(s1).request(anyLong());
    final Future<?> requesting = executor.submit(() -> s.request(1));
    try {
        assertThat("The request was not reached", requested.await(DEFAULT_TIMEOUT_SECONDS, SECONDS), is(true));
        s.cancel();
        assertThat("Cancel did not reach the blocked subscription", cancelled.getCount(), is(0L));
    } finally {
        cancelled.countDown();
        requesting.get();
    }
}

@bryce-anderson bryce-anderson Oct 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think technically allowing cancel() and request(n) calls to happen concurrently from different threads is a RS spec violation, part 2.7. That said, it looks like we already break these rules on main. It feels like the right thing to do is fix the blocking iterable (and potentially others) to not block in those calls.

Unfortunately, that puts us deeper down a rabbit hole. 😞
What do you think @idelpivnitskiy, should we allow it now and try to fix it later after we can fix the BlockingIterable (and maybe others, we'd need to do an audit), or should we try to cull any blocking in request(n) calls first?

@idelpivnitskiy idelpivnitskiy Oct 10, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. To make incremental progress, would be nice to preserve the current behavior of main and then fix later after other changes


/**
* Switches <strong>current</strong> {@link Subscription} to {@code next}. It is assumed the {@link Subscriber}
* associated with the previous {@link Subscription} will no longer call {@link #itemReceived()}.
* <p>
* Only can be called in the {@link Subscriber} thread!
* Switches <strong>current</strong> {@link Subscription} to {@code next}. The previous source must no longer call
* {@link #itemReceived()}. Subscriber signals must be serialized, but a previous request callback may still be
* unwinding on another thread.
* @param next {@link Subscription} that should now be <strong>current</strong>.
*/
void switchTo(Subscription next) {
requireNonNull(next);
// No special concurrency considerations for sourceEmitted access is required in this method because we are
// on the Subscriber thread in this method. We want to track the effective source requested for the purposes of
// how much more request(n) is necessary below.
long effectiveSourceRequested = sourceEmitted;
for (;;) {
final long currSourceRequested = sourceRequested;
if (currSourceRequested == CANCELLED) {
final long currRequested = requested;
if (currRequested >= 0) {
next.cancel();
} else { // invalid requestN is pending, deliver to each subscription.
next.request(currRequested);
}
break;
} else if (currSourceRequested == SWITCHING) {
// concurrency is not allowed on this method, but reentry is allowed. save next into subscription so
// we can use it when the stack unwinds to deliver demand to it.
subscription = next;
break;
} else if (sourceRequestedUpdater.compareAndSet(this, currSourceRequested, SWITCHING)) {
assert currSourceRequested >= 0 || currSourceRequested == REQUESTED;
final long currRequested = requested;
if (currRequested < 0) { // invalid requestN is pending.
sourceRequested = CANCELLED;
next.request(currRequested);
break;
}

// sourceEmitted is stable here because we are on the Subscriber thread. We want to request the
// difference between total requested and what has been emitted from the new subscription. We also
// need to set the value of total requested below to make sure it is monotonically increasing.
// effectiveSourceRequested ...[delta]... requested
final long delta = currRequested - effectiveSourceRequested;
assert delta >= 0;
final Subscription beforeSubscription = subscription;
if (delta != 0) {
// There maybe concurrency with the Subscription thread, or synchronous delivery of data from
// request(n). In these cases we want to avoid "double request" from requested, so we track how much
// we have already requested and decrement it on future loop iterations.
effectiveSourceRequested = currRequested;
next.request(delta);
}

final boolean reentry = beforeSubscription != subscription;
if (reentry) {
// subscription was overwritten higher in the stack to track the more recent value. overwrite next
// in our current stack frame to use the more recent value on the next loop iteration.
next = subscription;

// There is a new subscription so we need to reset state for how much has been emitted, so we
// deliver demand to the more recent subscription on the next loop iteration.
effectiveSourceRequested = sourceEmitted;
final Subscription previous = pendingSubscriptionUpdater.getAndSet(this, requireNonNull(next));
try {
// A source may terminate without demand before the drain reaches it. Only the latest source needs demand,
// but a displaced subscription must still observe a pending terminal action.
if (previous != null) {
if (cancelled) {
previous.cancel();
} else {
// Make the subscription visible before restoring the state of sourceRequested. If the Subscription
// thread observes the sourceRequested change it will also observe the subscription change. The
// Subscription thread also uses sourceRequested to make sure there is no concurrent invocation of
// the switched Subscription.
subscription = next;
final long n = requested;
if (n < 0) {
previous.request(n);
}
}
}
} finally {
drain();
}
}

// We want to set sourceRequested to currRequested because we have already requested the delta between
// the two above, and we want to ensure sourceRequested is always monotonically increasing
// (besides control values) to prevent the Subscription thread from requesting from an old subscription.
if (sourceRequestedUpdater.compareAndSet(this, SWITCHING, currRequested) && !reentry) {
break;
private void drain() {
Throwable delayedCause = null;
boolean tryAcquire = true;
while (tryAcquire && tryAcquireLock(emittingUpdater, this)) {
try {
final Subscription next = pendingSubscriptionUpdater.getAndSet(this, null);
if (cancelled) {
final Subscription current = subscription;
subscription = EMPTY_SUBSCRIPTION_NO_THROW;
try {
current.cancel();
} finally {
if (next != null) {
next.cancel();
}
}
} else {
if (next != null) {
subscription = next;
sourceRequested = sourceEmitted;
}
final long n = requested;
if (sourceRequested >= 0) {
if (n < 0) {
sourceRequested = n;
subscription.request(n);
} else {
final long delta = n - sourceRequested;

@idelpivnitskiy idelpivnitskiy Oct 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider preserving the old assert delta >= 0 here. In case source over-emits, sourceEmitted > requested after a switch, and the next source gets request(<negative>)

if (delta != 0) {
// Commit before the callback: reentrant request/switch calls only enqueue more work.
sourceRequested = n;
subscription.request(delta);
}
}
}
}
// else the Subscription thread was active in the mean time, we need to process the pending event(s).
// if the state is cancelled the Subscription thread defers to this thread to do the cancel on the
// next loop invocation. or this method was invoked in a re-entry fashion and we need to loop again.
} catch (Throwable cause) {
delayedCause = catchUnexpected(delayedCause, cause);
} finally {
// A publisher racing with release either marks pending work or becomes the next drain owner.
tryAcquire = !releaseLock(emittingUpdater, this);
}
}
if (delayedCause != null) {
throwException(delayedCause);
}
}

/**
Expand Down
Loading