2.x: Fix cancel/dispose upon upstream switch for some operators (#6258)

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

* Restore star imports

Author: akarnokd

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