3.x: Fix switchMap not canceling properly during onNext-cancel races (ReactiveX#6917)

Author: akarnokd

src/main/java/io/reactivex/rxjava3/internal/operators/flowable/FlowableSwitchMap.java

@@ -198,7 +198,6 @@ void drain() {
             for (;;) {
 
                 if (cancelled) {
-                    active.lazySet(null);
                     return;
                 }
 

src/test/java/io/reactivex/rxjava3/internal/operators/flowable/FlowableSwitchTest.java

@@ -19,7 +19,7 @@
 
 import java.util.*;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.*;
 
 import org.junit.*;
 import org.mockito.InOrder;
@@ -1337,4 +1337,44 @@ public void innerCompletedBackpressureBoundaryDelayError() {
 
         ts.assertValuesOnly(1);
     }
+
+    @Test
+    public void cancellationShouldTriggerInnerCancellationRace() throws Throwable {
+        AtomicInteger outer = new AtomicInteger();
+        AtomicInteger inner = new AtomicInteger();
+
+        int n = 10_000;
+        for (int i = 0; i < n; i++) {
+            Flowable.<Integer>create(it -> {
+                it.onNext(0);
+            }, BackpressureStrategy.MISSING)
+            .switchMap(v -> createFlowable(inner))
+            .observeOn(Schedulers.computation())
+            .doFinally(() -> {
+                outer.incrementAndGet();
+            })
+            .take(1)
+            .blockingSubscribe(v -> { }, Throwable::printStackTrace);
+        }
+
+        Thread.sleep(100);
+        assertEquals(inner.get(), outer.get());
+        assertEquals(n, inner.get());
+    }
+
+    Flowable<Integer> createFlowable(AtomicInteger inner) {
+        return Flowable.<Integer>unsafeCreate(s -> {
+            SerializedSubscriber<Integer> it = new SerializedSubscriber<>(s);
+            it.onSubscribe(new BooleanSubscription());
+            Schedulers.io().scheduleDirect(() -> {
+                it.onNext(1);
+            }, 0, TimeUnit.MILLISECONDS);
+            Schedulers.io().scheduleDirect(() -> {
+                it.onNext(2);
+            }, 0, TimeUnit.MILLISECONDS);
+        })
+        .doFinally(() -> {
+            inner.incrementAndGet();
+        });
+    }
 }

src/test/java/io/reactivex/rxjava3/internal/operators/observable/ObservableSwitchTest.java

@@ -34,7 +34,7 @@
 import io.reactivex.rxjava3.internal.functions.Functions;
 import io.reactivex.rxjava3.internal.schedulers.ImmediateThinScheduler;
 import io.reactivex.rxjava3.internal.util.ExceptionHelper;
-import io.reactivex.rxjava3.observers.TestObserver;
+import io.reactivex.rxjava3.observers.*;
 import io.reactivex.rxjava3.plugins.RxJavaPlugins;
 import io.reactivex.rxjava3.schedulers.*;
 import io.reactivex.rxjava3.subjects.PublishSubject;
@@ -1398,4 +1398,44 @@ public void innerIgnoresCancelAndErrors() throws Throwable {
             TestHelper.assertUndeliverable(errors, 0, TestException.class);
         });
     }
+
+    @Test
+    public void cancellationShouldTriggerInnerCancellationRace() throws Throwable {
+        AtomicInteger outer = new AtomicInteger();
+        AtomicInteger inner = new AtomicInteger();
+
+        int n = 10_000;
+        for (int i = 0; i < n; i++) {
+            Observable.<Integer>create(it -> {
+                it.onNext(0);
+            })
+            .switchMap(v -> createObservable(inner))
+            .observeOn(Schedulers.computation())
+            .doFinally(() -> {
+                outer.incrementAndGet();
+            })
+            .take(1)
+            .blockingSubscribe(v -> { }, Throwable::printStackTrace);
+        }
+
+        Thread.sleep(100);
+        assertEquals(inner.get(), outer.get());
+        assertEquals(n, inner.get());
+    }
+
+    Observable<Integer> createObservable(AtomicInteger inner) {
+        return Observable.<Integer>unsafeCreate(s -> {
+            SerializedObserver<Integer> it = new SerializedObserver<>(s);
+            it.onSubscribe(Disposable.empty());
+            Schedulers.io().scheduleDirect(() -> {
+                it.onNext(1);
+            }, 0, TimeUnit.MILLISECONDS);
+            Schedulers.io().scheduleDirect(() -> {
+                it.onNext(2);
+            }, 0, TimeUnit.MILLISECONDS);
+        })
+        .doFinally(() -> {
+            inner.incrementAndGet();
+        });
+    }
 }