Fix expand current concurrency on cancel()

Author: akarnokd

src/main/java/hu/akarnokd/rxjava2/operators/FlowableExpand.java

@@ -185,6 +185,8 @@ static final class ExpandDepthSubscription<T>
 
         final AtomicLong requested;
 
+        final AtomicReference<Object> current;
+
         ArrayDeque<ExpandDepthSubscriber> subscriptionStack;
 
         volatile boolean cancelled;
@@ -193,8 +195,6 @@ static final class ExpandDepthSubscription<T>
 
         long consumed;
 
-        ExpandDepthSubscriber current;
-
         ExpandDepthSubscription(Subscriber<? super T> actual,
                 Function<? super T, ? extends Publisher<? extends T>> expander,
                         int capacityHint) {
@@ -204,6 +204,7 @@ static final class ExpandDepthSubscription<T>
             this.error = new AtomicThrowable();
             this.active = new AtomicInteger();
             this.requested = new AtomicLong();
+            this.current = new AtomicReference<Object>();
         }
 
         @Override
@@ -214,6 +215,7 @@ public void request(long n) {
             }
         }
 
+        @SuppressWarnings("unchecked")
         @Override
         public void cancel() {
             if (!cancelled) {
@@ -229,6 +231,11 @@ public void cancel() {
                         q.poll().dispose();
                     }
                 }
+
+                Object o = current.getAndSet(this);
+                if (o != this && o != null) {
+                    ((ExpandDepthSubscriber)o).dispose();
+                }
             }
         }
 
@@ -250,6 +257,21 @@ boolean push(ExpandDepthSubscriber subscriber) {
             }
         }
 
+        boolean setCurrent(ExpandDepthSubscriber inner) {
+            for (;;) {
+                Object o = current.get();
+                if (o == this) {
+                    if (inner != null) {
+                        inner.dispose();
+                    }
+                    return false;
+                }
+                if (current.compareAndSet(o, inner)) {
+                    return true;
+                }
+            }
+        }
+
         void drainQueue() {
             if (getAndIncrement() != 0) {
                 return;
@@ -261,13 +283,14 @@ void drainQueue() {
             AtomicInteger n = active;
 
             for (;;) {
-                if (cancelled) {
+                Object o = current.get();
+                if (cancelled || o == this) {
                     source = null;
-                    current = null;
                     return;
                 }
 
-                ExpandDepthSubscriber curr = current;
+                @SuppressWarnings("unchecked")
+                ExpandDepthSubscriber curr = (ExpandDepthSubscriber)o;
                 Publisher<? extends T> p = source;
 
                 if (curr == null && p != null) {
@@ -276,9 +299,11 @@ void drainQueue() {
 
                     ExpandDepthSubscriber eds = new ExpandDepthSubscriber();
                     curr = eds;
-                    current = eds;
-
-                    p.subscribe(eds);
+                    if (setCurrent(eds)) {
+                        p.subscribe(eds);
+                    } else {
+                        return;
+                    }
                 } else {
 
                     boolean currentDone = curr.done;
@@ -306,10 +331,12 @@ void drainQueue() {
                             if (push(curr)) {
                                 n.getAndIncrement();
                                 curr = new ExpandDepthSubscriber();
-                                current = curr;
-
-                                p.subscribe(curr);
-                                newSource = true;
+                                if (setCurrent(curr)) {
+                                    p.subscribe(curr);
+                                    newSource = true;
+                                } else {
+                                    return;
+                                }
                             }
                         }
                     }
@@ -326,9 +353,12 @@ void drainQueue() {
                                 return;
                             }
                             curr = pop();
-                            current = curr;
-                            curr.requestOne();
-                            continue;
+                            if (setCurrent(curr)) {
+                                curr.requestOne();
+                                continue;
+                            } else {
+                                return;
+                            }
                         }
                     }
                 }

src/test/java/hu/akarnokd/rxjava2/operators/FlowableExpandTest.java

@@ -18,11 +18,13 @@
 
 import java.io.IOException;
 import java.util.*;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.junit.*;
 import org.reactivestreams.*;
 
+import hu.akarnokd.rxjava2.test.TestHelper;
 import io.reactivex.*;
 import io.reactivex.functions.Function;
 import io.reactivex.internal.functions.Functions;
@@ -390,4 +392,121 @@ public void onComplete() {
 
         ts.assertResult(1);
     }
+
+    @Test
+    public void depthCancelRace() {
+        for (int i = 0; i < 1000; i++) {
+            final TestSubscriber<Integer> ts = Flowable.just(0)
+            .compose(FlowableTransformers.<Integer>expand(countDown, ExpandStrategy.DEPTH_FIRST))
+            .test(0);
+
+            Runnable r1 = new Runnable() {
+                @Override
+                public void run() {
+                    ts.request(1);
+                }
+            };
+            Runnable r2 = new Runnable() {
+                @Override
+                public void run() {
+                    ts.cancel();
+                }
+            };
+
+            TestHelper.race(r1, r2, Schedulers.single());
+        }
+    }
+
+    @Test
+    public void depthEmitCancelRace() {
+        for (int i = 0; i < 1000; i++) {
+
+            final PublishProcessor<Integer> pp = PublishProcessor.create();
+
+            final TestSubscriber<Integer> ts = Flowable.just(0)
+            .compose(FlowableTransformers.<Integer>expand(Functions.justFunction(pp), ExpandStrategy.DEPTH_FIRST))
+            .test(1);
+
+            Runnable r1 = new Runnable() {
+                @Override
+                public void run() {
+                    pp.onNext(1);
+                }
+            };
+            Runnable r2 = new Runnable() {
+                @Override
+                public void run() {
+                    ts.cancel();
+                }
+            };
+
+            TestHelper.race(r1, r2, Schedulers.single());
+        }
+    }
+
+    @Test
+    public void depthCompleteCancelRace() {
+        for (int i = 0; i < 1000; i++) {
+
+            final PublishProcessor<Integer> pp = PublishProcessor.create();
+
+            final TestSubscriber<Integer> ts = Flowable.just(0)
+            .compose(FlowableTransformers.<Integer>expand(Functions.justFunction(pp), ExpandStrategy.DEPTH_FIRST))
+            .test(1);
+
+            Runnable r1 = new Runnable() {
+                @Override
+                public void run() {
+                    pp.onComplete();
+                }
+            };
+            Runnable r2 = new Runnable() {
+                @Override
+                public void run() {
+                    ts.cancel();
+                }
+            };
+
+            TestHelper.race(r1, r2, Schedulers.single());
+        }
+    }
+
+    @Test
+    public void depthCancelRace2() throws Exception {
+        for (int i = 0; i < 1000; i++) {
+
+            final PublishProcessor<Integer> pp = PublishProcessor.create();
+
+            Flowable<Integer> source = Flowable.just(0)
+            .compose(FlowableTransformers.<Integer>expand(Functions.justFunction(pp), ExpandStrategy.DEPTH_FIRST));
+
+            final CountDownLatch cdl = new CountDownLatch(1);
+
+            TestSubscriber<Integer> ts = new TestSubscriber<Integer>() {
+                final AtomicInteger sync = new AtomicInteger(2);
+
+                @Override
+                public void onNext(Integer t) {
+                    super.onNext(t);
+                    Schedulers.single().scheduleDirect(new Runnable() {
+                        @Override
+                        public void run() {
+                            if (sync.decrementAndGet() != 0) {
+                                while (sync.get() != 0) { }
+                            }
+                            cancel();
+                            cdl.countDown();
+                        }
+                    });
+                    if (sync.decrementAndGet() != 0) {
+                        while (sync.get() != 0) { }
+                    }
+                }
+            };
+
+            source.subscribe(ts);
+
+            Assert.assertTrue(cdl.await(5, TimeUnit.SECONDS));
+        }
+    }
 }