2.x: Fix cancel/dispose upon upstream switch for some operators

src/main/java/io/reactivex/internal/operators/flowable/FlowableConcatArray.java

@@ -59,6 +59,7 @@ static final class ConcatArraySubscriber<T> extends SubscriptionArbiter implemen
         long produced;
 
         ConcatArraySubscriber(Publisher<? extends T>[] sources, boolean delayError, Subscriber<? super T> downstream) {
+            super(false);
             this.downstream = downstream;
             this.sources = sources;
             this.delayError = delayError;

src/main/java/io/reactivex/internal/operators/flowable/FlowableConcatMap.java

@@ -571,6 +571,7 @@ static final class ConcatMapInner<R>
         long produced;
 
         ConcatMapInner(ConcatMapSupport<R> parent) {
+            super(false);
             this.parent = parent;
         }
 

src/main/java/io/reactivex/internal/operators/flowable/FlowableDelaySubscriptionOther.java

@@ -12,10 +12,12 @@
  */
 package io.reactivex.internal.operators.flowable;
 
+import java.util.concurrent.atomic.*;
+
 import org.reactivestreams.*;
 
 import io.reactivex.*;
-import io.reactivex.internal.subscriptions.SubscriptionArbiter;
+import io.reactivex.internal.subscriptions.SubscriptionHelper;
 import io.reactivex.plugins.RxJavaPlugins;
 
 /**
@@ -35,94 +37,105 @@ public FlowableDelaySubscriptionOther(Publisher<? extends T> main, Publisher<U>
 
     @Override
     public void subscribeActual(final Subscriber<? super T> child) {
-        final SubscriptionArbiter serial = new SubscriptionArbiter();
-        child.onSubscribe(serial);
+        MainSubscriber<T> parent = new MainSubscriber<T>(child, main);
+        child.onSubscribe(parent);
+        other.subscribe(parent.other);
+    }
 
-        FlowableSubscriber<U> otherSubscriber = new DelaySubscriber(serial, child);
+    static final class MainSubscriber<T> extends AtomicLong implements FlowableSubscriber<T>, Subscription {
 
-        other.subscribe(otherSubscriber);
-    }
+        private static final long serialVersionUID = 2259811067697317255L;
+
+        final Subscriber<? super T> downstream;
 
-    final class DelaySubscriber implements FlowableSubscriber<U> {
-        final SubscriptionArbiter serial;
-        final Subscriber<? super T> child;
-        boolean done;
+        final Publisher<? extends T> main;
 
-        DelaySubscriber(SubscriptionArbiter serial, Subscriber<? super T> child) {
-            this.serial = serial;
-            this.child = child;
+        final OtherSubscriber other;
+
+        final AtomicReference<Subscription> upstream;
+
+        MainSubscriber(Subscriber<? super T> downstream, Publisher<? extends T> main) {
+            this.downstream = downstream;
+            this.main = main;
+            this.other = new OtherSubscriber();
+            this.upstream = new AtomicReference<Subscription>();
         }
 
-        @Override
-        public void onSubscribe(final Subscription s) {
-            serial.setSubscription(new DelaySubscription(s));
-            s.request(Long.MAX_VALUE);
+        void next() {
+            main.subscribe(this);
         }
 
         @Override
-        public void onNext(U t) {
-            onComplete();
+        public void onNext(T t) {
+            downstream.onNext(t);
         }
 
         @Override
-        public void onError(Throwable e) {
-            if (done) {
-                RxJavaPlugins.onError(e);
-                return;
-            }
-            done = true;
-            child.onError(e);
+        public void onError(Throwable t) {
+            downstream.onError(t);
         }
 
         @Override
         public void onComplete() {
-            if (done) {
-                return;
-            }
-            done = true;
-
-            main.subscribe(new OnCompleteSubscriber());
+            downstream.onComplete();
         }
 
-        final class DelaySubscription implements Subscription {
-
-            final Subscription upstream;
-
-            DelaySubscription(Subscription s) {
-                this.upstream = s;
+        @Override
+        public void request(long n) {
+            if (SubscriptionHelper.validate(n)) {
+                SubscriptionHelper.deferredRequest(upstream, this, n);
             }
+        }
 
-            @Override
-            public void request(long n) {
-                // ignored
-            }
+        @Override
+        public void cancel() {
+            SubscriptionHelper.cancel(other);
+            SubscriptionHelper.cancel(upstream);
+        }
 
-            @Override
-            public void cancel() {
-                upstream.cancel();
-            }
+        @Override
+        public void onSubscribe(Subscription s) {
+            SubscriptionHelper.deferredSetOnce(upstream, this, s);
         }
 
-        final class OnCompleteSubscriber implements FlowableSubscriber<T> {
+        final class OtherSubscriber extends AtomicReference<Subscription> implements FlowableSubscriber<Object> {
+
+            private static final long serialVersionUID = -3892798459447644106L;
+
             @Override
             public void onSubscribe(Subscription s) {
-                serial.setSubscription(s);
+                if (SubscriptionHelper.setOnce(this, s)) {
+                    s.request(Long.MAX_VALUE);
+                }
             }
 
             @Override
-            public void onNext(T t) {
-                child.onNext(t);
+            public void onNext(Object t) {
+                Subscription s = get();
+                if (s != SubscriptionHelper.CANCELLED) {
+                    lazySet(SubscriptionHelper.CANCELLED);
+                    s.cancel();
+                    next();
+                }
             }
 
             @Override
             public void onError(Throwable t) {
-                child.onError(t);
+                Subscription s = get();
+                if (s != SubscriptionHelper.CANCELLED) {
+                    downstream.onError(t);
+                } else {
+                    RxJavaPlugins.onError(t);
+                }
             }
 
             @Override
             public void onComplete() {
-                child.onComplete();
+                Subscription s = get();
+                if (s != SubscriptionHelper.CANCELLED) {
+                    next();
+                }
             }
         }
-    }
+}
 }

src/main/java/io/reactivex/internal/operators/flowable/FlowableOnErrorNext.java

@@ -58,6 +58,7 @@ static final class OnErrorNextSubscriber<T>
         long produced;
 
         OnErrorNextSubscriber(Subscriber<? super T> actual, Function<? super Throwable, ? extends Publisher<? extends T>> nextSupplier, boolean allowFatal) {
+            super(false);
             this.downstream = actual;
             this.nextSupplier = nextSupplier;
             this.allowFatal = allowFatal;

src/main/java/io/reactivex/internal/operators/flowable/FlowableRepeat.java

@@ -29,7 +29,7 @@ public FlowableRepeat(Flowable<T> source, long count) {
 
     @Override
     public void subscribeActual(Subscriber<? super T> s) {
-        SubscriptionArbiter sa = new SubscriptionArbiter();
+        SubscriptionArbiter sa = new SubscriptionArbiter(false);
         s.onSubscribe(sa);
 
         RepeatSubscriber<T> rs = new RepeatSubscriber<T>(s, count != Long.MAX_VALUE ? count - 1 : Long.MAX_VALUE, sa, source);

src/main/java/io/reactivex/internal/operators/flowable/FlowableRepeatUntil.java

@@ -31,7 +31,7 @@ public FlowableRepeatUntil(Flowable<T> source, BooleanSupplier until) {
 
     @Override
     public void subscribeActual(Subscriber<? super T> s) {
-        SubscriptionArbiter sa = new SubscriptionArbiter();
+        SubscriptionArbiter sa = new SubscriptionArbiter(false);
         s.onSubscribe(sa);
 
         RepeatSubscriber<T> rs = new RepeatSubscriber<T>(s, until, sa, source);

src/main/java/io/reactivex/internal/operators/flowable/FlowableRepeatWhen.java

@@ -143,6 +143,7 @@ abstract static class WhenSourceSubscriber<T, U> extends SubscriptionArbiter imp
 
         WhenSourceSubscriber(Subscriber<? super T> actual, FlowableProcessor<U> processor,
                 Subscription receiver) {
+            super(false);
             this.downstream = actual;
             this.processor = processor;
             this.receiver = receiver;
@@ -160,6 +161,7 @@ public final void onNext(T t) {
         }
 
         protected final void again(U signal) {
+            setSubscription(EmptySubscription.INSTANCE);
             long p = produced;
             if (p != 0L) {
                 produced = 0L;

src/main/java/io/reactivex/internal/operators/flowable/FlowableRetryBiPredicate.java

@@ -33,7 +33,7 @@ public FlowableRetryBiPredicate(
 
     @Override
     public void subscribeActual(Subscriber<? super T> s) {
-        SubscriptionArbiter sa = new SubscriptionArbiter();
+        SubscriptionArbiter sa = new SubscriptionArbiter(false);
         s.onSubscribe(sa);
 
         RetryBiSubscriber<T> rs = new RetryBiSubscriber<T>(s, predicate, sa, source);

src/main/java/io/reactivex/internal/operators/flowable/FlowableRetryPredicate.java

@@ -35,7 +35,7 @@ public FlowableRetryPredicate(Flowable<T> source,
 
     @Override
     public void subscribeActual(Subscriber<? super T> s) {
-        SubscriptionArbiter sa = new SubscriptionArbiter();
+        SubscriptionArbiter sa = new SubscriptionArbiter(false);
         s.onSubscribe(sa);
 
         RetrySubscriber<T> rs = new RetrySubscriber<T>(s, count, predicate, sa, source);

src/main/java/io/reactivex/internal/operators/flowable/FlowableSwitchIfEmpty.java

@@ -43,7 +43,7 @@ static final class SwitchIfEmptySubscriber<T> implements FlowableSubscriber<T> {
             this.downstream = actual;
             this.other = other;
             this.empty = true;
-            this.arbiter = new SubscriptionArbiter();
+            this.arbiter = new SubscriptionArbiter(false);
         }
 
         @Override

src/main/java/io/reactivex/internal/operators/flowable/FlowableTimeout.java

@@ -208,6 +208,7 @@ static final class TimeoutFallbackSubscriber<T> extends SubscriptionArbiter
         TimeoutFallbackSubscriber(Subscriber<? super T> actual,
                 Function<? super T, ? extends Publisher<?>> itemTimeoutIndicator,
                         Publisher<? extends T> fallback) {
+            super(true);
             this.downstream = actual;
             this.itemTimeoutIndicator = itemTimeoutIndicator;
             this.task = new SequentialDisposable();

src/main/java/io/reactivex/internal/operators/flowable/FlowableTimeoutTimed.java

@@ -196,6 +196,7 @@ static final class TimeoutFallbackSubscriber<T> extends SubscriptionArbiter
 
         TimeoutFallbackSubscriber(Subscriber<? super T> actual, long timeout, TimeUnit unit,
                 Scheduler.Worker worker, Publisher<? extends T> fallback) {
+            super(true);
             this.downstream = actual;
             this.timeout = timeout;
             this.unit = unit;

src/main/java/io/reactivex/internal/operators/observable/ObservableConcatMap.java

@@ -244,7 +244,7 @@ static final class InnerObserver<U> extends AtomicReference<Disposable> implemen
 
             @Override
             public void onSubscribe(Disposable d) {
-                DisposableHelper.set(this, d);
+                DisposableHelper.replace(this, d);
             }
 
             @Override

src/main/java/io/reactivex/internal/operators/observable/ObservableRepeatWhen.java

@@ -92,7 +92,7 @@ static final class RepeatWhenObserver<T> extends AtomicInteger implements Observ
 
         @Override
         public void onSubscribe(Disposable d) {
-            DisposableHelper.replace(this.upstream, d);
+            DisposableHelper.setOnce(this.upstream, d);
         }
 
         @Override
@@ -109,6 +109,7 @@ public void onError(Throwable e) {
         @Override
         public void onComplete() {
             active = false;
+            DisposableHelper.replace(upstream, null);
             signaller.onNext(0);
         }
 

src/main/java/io/reactivex/internal/operators/observable/ObservableRetryBiPredicate.java

@@ -58,7 +58,7 @@ static final class RetryBiObserver<T> extends AtomicInteger implements Observer<
 
         @Override
         public void onSubscribe(Disposable d) {
-            upstream.update(d);
+            upstream.replace(d);
         }
 
         @Override

src/main/java/io/reactivex/internal/operators/observable/ObservableRetryPredicate.java

@@ -61,7 +61,7 @@ static final class RepeatObserver<T> extends AtomicInteger implements Observer<T
 
         @Override
         public void onSubscribe(Disposable d) {
-            upstream.update(d);
+            upstream.replace(d);
         }
 
         @Override

src/main/java/io/reactivex/internal/operators/observable/ObservableRetryWhen.java

@@ -102,6 +102,7 @@ public void onNext(T t) {
 
         @Override
         public void onError(Throwable e) {
+            DisposableHelper.replace(upstream, null);
             active = false;
             signaller.onNext(e);
         }

src/main/java/io/reactivex/internal/subscriptions/SubscriptionArbiter.java

@@ -55,11 +55,14 @@ public class SubscriptionArbiter extends AtomicInteger implements Subscription {
 
     final AtomicLong missedProduced;
 
+    final boolean cancelOnReplace;
+
     volatile boolean cancelled;
 
     protected boolean unbounded;
 
-    public SubscriptionArbiter() {
+    public SubscriptionArbiter(boolean cancelOnReplace) {
+        this.cancelOnReplace = cancelOnReplace;
         missedSubscription = new AtomicReference<Subscription>();
         missedRequested = new AtomicLong();
         missedProduced = new AtomicLong();
@@ -80,7 +83,7 @@ public final void setSubscription(Subscription s) {
         if (get() == 0 && compareAndSet(0, 1)) {
             Subscription a = actual;
 
-            if (a != null) {
+            if (a != null && cancelOnReplace) {
                 a.cancel();
             }
 
@@ -100,7 +103,7 @@ public final void setSubscription(Subscription s) {
         }
 
         Subscription a = missedSubscription.getAndSet(s);
-        if (a != null) {
+        if (a != null && cancelOnReplace) {
             a.cancel();
         }
         drain();
@@ -240,7 +243,7 @@ final void drainLoop() {
                 }
 
                 if (ms != null) {
-                    if (a != null) {
+                    if (a != null && cancelOnReplace) {
                         a.cancel();
                     }
                     actual = ms;