2.x: fix GroupBy MissingBackpressureException due to main/group overflow

The problem with GroupBy is that a request(1) from the main source may
result in a new group, a value for another group or value for the
intended group. The latter two is handled by unicasting but the former
was not properly handled.

Author: akarnokd

src/main/java/io/reactivex/internal/operators/OperatorGroupBy.java

@@ -21,10 +21,12 @@
 import org.reactivestreams.*;
 
 import io.reactivex.Observable.Operator;
-import io.reactivex.internal.queue.SpscLinkedArrayQueue;
+import io.reactivex.exceptions.MissingBackpressureException;
+import io.reactivex.internal.queue.*;
 import io.reactivex.internal.subscriptions.*;
-import io.reactivex.internal.util.BackpressureHelper;
+import io.reactivex.internal.util.*;
 import io.reactivex.observables.GroupedObservable;
+import io.reactivex.plugins.RxJavaPlugins;
 
 public final class OperatorGroupBy<T, K, V> implements Operator<GroupedObservable<K, V>, T>{
     final Function<? super T, ? extends K> keySelector;
@@ -44,7 +46,9 @@ public Subscriber<? super T> apply(Subscriber<? super GroupedObservable<K, V>> t
         return new GroupBySubscriber<>(t, keySelector, valueSelector, bufferSize, delayError);
     }
 
-    public static final class GroupBySubscriber<T, K, V> extends AtomicInteger implements Subscriber<T>, Subscription {
+    public static final class GroupBySubscriber<T, K, V> 
+    extends AtomicInteger
+    implements Subscriber<T>, Subscription {
         /** */
         private static final long serialVersionUID = -3688291656102519502L;
 
@@ -54,6 +58,7 @@ public static final class GroupBySubscriber<T, K, V> extends AtomicInteger imple
         final int bufferSize;
         final boolean delayError;
         final Map<Object, GroupedUnicast<K, V>> groups;
+        final Queue<T> queue;
 
         static final Object NULL_KEY = new Object();
 
@@ -64,14 +69,28 @@ public static final class GroupBySubscriber<T, K, V> extends AtomicInteger imple
         static final AtomicIntegerFieldUpdater<GroupBySubscriber> CANCELLED =
                 AtomicIntegerFieldUpdater.newUpdater(GroupBySubscriber.class, "cancelled");
 
+        volatile long requested;
+        @SuppressWarnings("rawtypes")
+        static final AtomicLongFieldUpdater<GroupBySubscriber> REQUESTED =
+                AtomicLongFieldUpdater.newUpdater(GroupBySubscriber.class, "requested");
+        
+        volatile int groupCount;
+        @SuppressWarnings("rawtypes")
+        static final AtomicIntegerFieldUpdater<GroupBySubscriber> GROUP_COUNT =
+                AtomicIntegerFieldUpdater.newUpdater(GroupBySubscriber.class, "groupCount");
+        
+        Throwable error;
+        volatile boolean done;
+        
         public GroupBySubscriber(Subscriber<? super GroupedObservable<K, V>> actual, Function<? super T, ? extends K> keySelector, Function<? super T, ? extends V> valueSelector, int bufferSize, boolean delayError) {
             this.actual = actual;
             this.keySelector = keySelector;
             this.valueSelector = valueSelector;
             this.bufferSize = bufferSize;
             this.delayError = delayError;
             this.groups = new ConcurrentHashMap<>();
-            this.lazySet(1);
+            this.queue = new SpscLinkedArrayQueue<>(bufferSize);
+            GROUP_COUNT.lazySet(this, 1);
         }
 
         @Override
@@ -82,90 +101,60 @@ public void onSubscribe(Subscription s) {
 
             this.s = s;
             actual.onSubscribe(this);
+            s.request(bufferSize);
         }
 
         @Override
         public void onNext(T t) {
-            K key;
-            try {
-                key = keySelector.apply(t);
-            } catch (Throwable e) {
-                s.cancel();
-                onError(e);
+            if (done) {
                 return;
             }
-            
-            boolean notNew = true;
-            Object mapKey = key != null ? key : NULL_KEY;
-            GroupedUnicast<K, V> group = groups.get(mapKey);
-            if (group == null) {
-                // if the main has been cancelled, stop creating groups
-                // and skip this value
-                if (cancelled != 0) {
-                    s.request(1);
-                    return;
-                }
-                notNew = true;
-                
-                group = GroupedUnicast.createWith(key, bufferSize, this, delayError);
-                groups.put(mapKey, group);
-                
-                getAndIncrement();
-                
-                actual.onNext(group);
-            }
-            
-            V v;
-            try {
-                v = valueSelector.apply(t);
-            } catch (Throwable e) {
+            if (!queue.offer(t)) {
                 s.cancel();
-                onError(e);
+                onError(new MissingBackpressureException("Queue full?!"));
                 return;
             }
-
-            group.onNext(v);
-            
-            if (notNew) {
-                s.request(1); // we spent this t on an existing group, request one more
-            }
+            drain();
         }
 
         @Override
         public void onError(Throwable t) {
-            List<GroupedUnicast<K, V>> list = new ArrayList<>(groups.values());
-            groups.clear();
-            
-            for (GroupedUnicast<K, V> e : list) {
-                e.onError(t);
+            if (done) {
+                RxJavaPlugins.onError(t);
+                return;
             }
-            
-            actual.onError(t);
+            error = t;
+            done = true;
+            GROUP_COUNT.decrementAndGet(this);
+            drain();
         }
 
         @Override
         public void onComplete() {
-            List<GroupedUnicast<K, V>> list = new ArrayList<>(groups.values());
-            groups.clear();
-            
-            for (GroupedUnicast<K, V> e : list) {
-                e.onComplete();
+            if (done) {
+                return;
             }
-            
-            actual.onComplete();
+            done = true;
+            GROUP_COUNT.decrementAndGet(this);
+            drain();
         }
 
         @Override
         public void request(long n) {
-            s.request(n);
+            if (SubscriptionHelper.validateRequest(n)) {
+                return;
+            }
+            
+            BackpressureHelper.add(REQUESTED, this, n);
+            drain();
         }
 
         @Override
         public void cancel() {
             // cancelling the main source means we don't want any more groups
             // but running groups still require new values
             if (CANCELLED.compareAndSet(this, 0, 1)) {
-                if (decrementAndGet() == 0) {
+                if (GROUP_COUNT.decrementAndGet(this) == 0) {
                     s.cancel();
                 }
             }
@@ -174,10 +163,142 @@ public void cancel() {
         public void cancel(K key) {
             Object mapKey = key != null ? key : NULL_KEY;
             groups.remove(mapKey);
-            if (decrementAndGet() == 0) {
+            if (GROUP_COUNT.decrementAndGet(this) == 0) {
                 s.cancel();
             }
         }
+        
+        void drain() {
+            if (getAndIncrement() != 0) {
+                return;
+            }
+            
+            int missed = 1;
+            
+            final Queue<T> q = this.queue;
+            final Subscriber<? super GroupedObservable<K, V>> a = this.actual;
+            
+            for (;;) {
+                
+                if (checkTerminated(done, q.isEmpty(), a, q)) {
+                    return;
+                }
+                
+                long r = requested;
+                boolean unbounded = r == Long.MAX_VALUE;
+                long e = 0L;
+                
+                while (r != 0) {
+                    boolean d = done;
+                    
+                    T t = q.poll();
+                    
+                    boolean empty = t == null;
+                    
+                    if (checkTerminated(d, empty, a, q)) {
+                        return;
+                    }
+                    
+                    if (empty) {
+                        break;
+                    }
+                    
+                    K key;
+                    try {
+                        key = keySelector.apply(t);
+                    } catch (Throwable ex) {
+                        s.cancel();
+                        errorAll(a, q, ex);
+                        return;
+                    }
+                    
+                    boolean notNew = true;
+                    Object mapKey = key != null ? key : NULL_KEY;
+                    GroupedUnicast<K, V> group = groups.get(mapKey);
+                    if (group == null) {
+                        // if the main has been cancelled, stop creating groups
+                        // and skip this value
+                        if (cancelled == 0) {
+                            group = GroupedUnicast.createWith(key, bufferSize, this, delayError);
+                            groups.put(mapKey, group);
+                            
+                            GROUP_COUNT.getAndIncrement(this);
+                            
+                            actual.onNext(group);
+
+                            notNew = false;
+                        } else {
+                            continue;
+                        }
+                    }
+                    
+                    V v;
+                    try {
+                        v = valueSelector.apply(t);
+                    } catch (Throwable ex) {
+                        s.cancel();
+                        errorAll(a, q, ex);
+                        return;
+                    }
+
+                    group.onNext(v);
+                    
+                    if (notNew) {
+                        s.request(1); // we spent this t on an existing group, request one more
+                    } else {
+                        r--;
+                        e--;
+                    }
+                }
+                
+                if (e != 0L) {
+                    if (!unbounded) {
+                        REQUESTED.addAndGet(this, e);
+                    }
+                    s.request(-e);
+                }
+                
+                missed = addAndGet(-missed);
+                if (missed == 0) {
+                    break;
+                }
+            }
+        }
+        
+        void errorAll(Subscriber<? super GroupedObservable<K, V>> a, Queue<T> q, Throwable ex) {
+            q.clear();
+            List<GroupedUnicast<K, V>> list = new ArrayList<>(groups.values());
+            groups.clear();
+            
+            for (GroupedUnicast<K, V> e : list) {
+                e.onError(ex);
+            }
+            
+            a.onError(ex);
+        }
+        
+        boolean checkTerminated(boolean d, boolean empty, 
+                Subscriber<? super GroupedObservable<K, V>> a, Queue<T> q) {
+            if (d) {
+                Throwable err = error;
+                if (err != null) {
+                    errorAll(a, q, err);
+                    return true;
+                } else
+                if (empty) {
+                    List<GroupedUnicast<K, V>> list = new ArrayList<>(groups.values());
+                    groups.clear();
+                    
+                    for (GroupedUnicast<K, V> e : list) {
+                        e.onComplete();
+                    }
+                    
+                    actual.onComplete();
+                    return true;
+                }
+            }
+            return false;
+        }
     }
 
     static final class GroupedUnicast<K, T> extends GroupedObservable<K, T> {
@@ -247,7 +368,6 @@ public void request(long n) {
                 return;
             }
             BackpressureHelper.add(REQUESTED, this, n);
-            parent.request(n);
             drain();
         }
 

src/test/java/io/reactivex/internal/operators/OperatorGroupByTest.java

@@ -22,8 +22,8 @@
 import java.util.concurrent.atomic.*;
 import java.util.function.*;
 
-import org.junit.*;
-import org.mockito.*;
+import org.junit.Test;
+import org.mockito.Matchers;
 import org.reactivestreams.*;
 
 import io.reactivex.Observable;
@@ -1288,7 +1288,9 @@ public void testgroupByBackpressure2() throws InterruptedException {
 
         TestSubscriber<String> ts = new TestSubscriber<>();
 
-        Observable.range(1, 4000).groupBy(IS_EVEN2).flatMap(new Function<GroupedObservable<Boolean, Integer>, Observable<String>>() {
+        Observable.range(1, 4000)
+            .doOnNext(v -> System.out.println("testgroupByBackpressure2 >> " + v))
+            .groupBy(IS_EVEN2).flatMap(new Function<GroupedObservable<Boolean, Integer>, Observable<String>>() {
 
             @Override
             public Observable<String> apply(final GroupedObservable<Boolean, Integer> g) {
@@ -1471,4 +1473,42 @@ public void onNext(Integer t) {
                     }});
         assertTrue(completed.get());
     }
+    
+    /**
+     * Issue #3425.
+     * 
+     * The problem is that a request of 1 may create a new group, emit to the desired group
+     * or emit to a completely different group. In this test, the merge requests N which
+     * must be produced by the range, however it will create a bunch of groups before the actual
+     * group receives a value.
+     */
+    @Test
+    public void testBackpressureObserveOnOuter() {
+        for (int j = 0; j < 1000; j++) {
+            Observable.merge(
+                    Observable.range(0, 500)
+                    .groupBy(i -> i % (Observable.bufferSize() + 2))
+                    .observeOn(Schedulers.computation())
+            ).toBlocking().last();
+        }
+    }
+    
+    /**
+     * Synchronous verification of issue #3425.
+     */
+    @Test
+    public void testBackpressureInnerDoesntOverflowOuter() {
+        TestSubscriber<GroupedObservable<Integer, Integer>> ts = new TestSubscriber<>((Long)null);
+        
+        Observable.fromArray(1, 2)
+                .groupBy(v -> v)
+                .doOnNext(g -> g.subscribe()) // this will request Long.MAX_VALUE
+                .subscribe(ts)
+                ;
+        ts.request(1);
+        
+        ts.assertNotComplete();
+        ts.assertNoErrors();
+        ts.assertValueCount(1);
+    }
 }