3.x: Fix Flowable.groupBy cancellation/cleanup/eviction race hangs (#6979)

Author: akarnokd

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

@@ -259,8 +259,9 @@ private void completeEvictions() {
                 int count = 0;
                 GroupedUnicast<K, V> evictedGroup;
                 while ((evictedGroup = evictedGroups.poll()) != null) {
-                    evictedGroup.onComplete();
-                    count++;
+                    if (evictedGroup.state.tryComplete()) {
+                        count++;
+                    }
                 }
                 if (count != 0) {
                     groupCount.addAndGet(-count);
@@ -383,6 +384,8 @@ static final class State<T, K> extends BasicIntQueueSubscription<T> implements P
         static final int ABANDONED = 2;
         static final int ABANDONED_HAS_SUBSCRIBER = ABANDONED | HAS_SUBSCRIBER;
 
+        final AtomicBoolean evictOnce = new AtomicBoolean();
+
         State(int bufferSize, GroupBySubscriber<?, K, T> parent, K key, boolean delayError) {
             this.queue = new SpscLinkedArrayQueue<>(bufferSize);
             this.parent = parent;
@@ -444,9 +447,18 @@ public void onComplete() {
             drain();
         }
 
+        boolean tryComplete() {
+            boolean canEvict = evictOnce.compareAndSet(false, true);
+            done = true;
+            drain();
+            return canEvict;
+        }
+
         void cancelParent() {
             if ((once.get() & ABANDONED) == 0) {
-                parent.cancel(key);
+                if (evictOnce.compareAndSet(false, true)) {
+                    parent.cancel(key);
+                }
             }
         }
 
@@ -518,37 +530,44 @@ void drainNormal() {
             final SpscLinkedArrayQueue<T> q = queue;
             final boolean delayError = this.delayError;
             Subscriber<? super T> a = actual.get();
+            final AtomicBoolean cancelled = this.cancelled;
+
+            outer:
             for (;;) {
-                if (a != null) {
-                    long r = requested.get();
-                    long e = 0;
+                if (cancelled.get()) {
+                    cleanupQueue(0, false);
+                } else {
+                    if (a != null) {
+                        long r = requested.get();
+                        long e = 0;
 
-                    while (e != r) {
-                        boolean d = done;
-                        T v = q.poll();
-                        boolean empty = v == null;
+                        while (e != r) {
+                            boolean d = done;
+                            T v = q.poll();
+                            boolean empty = v == null;
 
-                        if (checkTerminated(d, empty, a, delayError, e)) {
-                            return;
-                        }
+                            if (checkTerminated(d, empty, a, delayError, e, !empty)) {
+                                continue outer;
+                            }
 
-                        if (empty) {
-                            break;
-                        }
+                            if (empty) {
+                                break;
+                            }
 
-                        a.onNext(v);
+                            a.onNext(v);
 
-                        e++;
-                    }
+                            e++;
+                        }
 
-                    if (e == r && checkTerminated(done, q.isEmpty(), a, delayError, e)) {
-                        return;
-                    }
+                        if (e == r && checkTerminated(done, q.isEmpty(), a, delayError, e, false)) {
+                            continue outer;
+                        }
 
-                    if (e != 0L) {
-                        BackpressureHelper.produced(requested, e);
-                        // replenish based on this batch run
-                        requestParent(e);
+                        if (e != 0L) {
+                            BackpressureHelper.produced(requested, e);
+                            // replenish based on this batch run
+                            requestParent(e);
+                        }
                     }
                 }
 
@@ -568,23 +587,32 @@ void requestParent(long e) {
             }
         }
 
-        boolean checkTerminated(boolean d, boolean empty, Subscriber<? super T> a, boolean delayError, long emitted) {
+        void cleanupQueue(long emitted, boolean polled) {
+            // if this group is canceled, all accumulated emissions and
+            // remaining items in the queue should be requested
+            // so that other groups can proceed
+            while (queue.poll() != null) {
+                emitted++;
+            }
+            if (polled) {
+                emitted++;
+            }
+            if (emitted != 0L) {
+                requestParent(emitted);
+            }
+        }
+
+        boolean checkTerminated(boolean d, boolean empty, Subscriber<? super T> a,
+                boolean delayError, long emitted, boolean polled) {
             if (cancelled.get()) {
-                // if this group is canceled, all accumulated emissions and
-                // remaining items in the queue should be requested
-                // so that other groups can proceed
-                while (queue.poll() != null) {
-                    emitted++;
-                }
-                if (emitted != 0L) {
-                    requestParent(emitted);
-                }
+                cleanupQueue(emitted, polled);
                 return true;
             }
 
             if (d) {
                 if (delayError) {
                     if (empty) {
+                        cancelled.lazySet(true);
                         Throwable e = error;
                         if (e != null) {
                             a.onError(e);
@@ -597,10 +625,12 @@ boolean checkTerminated(boolean d, boolean empty, Subscriber<? super T> a, boole
                     Throwable e = error;
                     if (e != null) {
                         queue.clear();
+                        cancelled.lazySet(true);
                         a.onError(e);
                         return true;
                     } else
                     if (empty) {
+                        cancelled.lazySet(true);
                         a.onComplete();
                         return true;
                     }

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

@@ -2604,4 +2604,93 @@ static <T> Function<Consumer<Object>, Map<T, Object>> sizeCap(int maxCapacity, b
         })
         .<T, Object>build().asMap();
     }
+
+    static void issue6974RunPart2(int groupByBufferSize, int flatMapMaxConcurrency, int groups,
+            boolean notifyOnExplicitEviction) {
+        TestSubscriber<Integer> ts = Flowable
+        .range(1, 500_000)
+        .map(i -> i % groups)
+        .groupBy(i -> i, i -> i, false, groupByBufferSize,
+                // set cap too high
+                sizeCap(groups * 100, notifyOnExplicitEviction))
+        .flatMap(gf -> gf
+                .take(10, TimeUnit.MILLISECONDS)
+                , flatMapMaxConcurrency)
+        .test();
+
+        ts
+        .awaitDone(5, TimeUnit.SECONDS)
+        .assertNoErrors()
+        .assertComplete();
+    }
+
+    @Test
+    public void issue6974Part2Case1() {
+        final int groups = 20;
+
+        // Not completed (Timed out), buffer is too small
+        int groupByBufferSize = groups * 2;
+        int flatMapMaxConcurrency = 2 * groups;
+        boolean notifyOnExplicitEviction = false;
+        issue6974RunPart2(groupByBufferSize, flatMapMaxConcurrency, groups, notifyOnExplicitEviction);
+    }
+
+    @Test
+    public void issue6974Part2Case2() {
+        final int groups = 20;
+
+        // Timeout... explicit eviction notification makes difference
+        int groupByBufferSize = groups * 30;
+        int flatMapMaxConcurrency = 2 * groups;
+        boolean notifyOnExplicitEviction = true;
+        issue6974RunPart2(groupByBufferSize, flatMapMaxConcurrency, groups, notifyOnExplicitEviction);
+    }
+
+    /*
+     * Disabled: Takes very long. Run it locally only.
+    @Test
+    public void issue6974Part2Case2Loop() {
+        for (int i = 0; i < 1000; i++) {
+            issue6974Part2Case2();
+        }
+    }
+    */
+
+    static void issue6974RunPart2NoEvict(int groupByBufferSize, int flatMapMaxConcurrency, int groups,
+            boolean notifyOnExplicitEviction) {
+        TestSubscriber<Integer> ts = Flowable
+        .range(1, 500_000)
+        .map(i -> i % groups)
+        .groupBy(i -> i)
+        .flatMap(gf -> gf
+                .take(10, TimeUnit.MILLISECONDS)
+                , flatMapMaxConcurrency)
+        .test();
+
+        ts
+        .awaitDone(5, TimeUnit.SECONDS)
+        .assertNoErrors()
+        .assertComplete();
+    }
+
+    @Test
+    public void issue6974Part2Case1NoEvict() {
+        final int groups = 20;
+
+        // Not completed (Timed out), buffer is too small
+        int groupByBufferSize = groups * 2;
+        int flatMapMaxConcurrency = 2 * groups;
+        boolean notifyOnExplicitEviction = false;
+        issue6974RunPart2NoEvict(groupByBufferSize, flatMapMaxConcurrency, groups, notifyOnExplicitEviction);
+    }
+
+    /*
+     * Disabled: Takes very long. Run it locally only.
+    @Test
+    public void issue6974Part2Case1NoEvictLoop() {
+        for (int i = 0; i < 1000; i++) {
+            issue6974Part2Case1NoEvict();
+        }
+    }
+    */
 }