Propagate onError to all groups

Author: zsxwing

src/main/java/rx/internal/operators/OperatorGroupBy.java

@@ -15,6 +15,7 @@
  */
 package rx.internal.operators;
 
+import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
@@ -138,6 +139,8 @@ public Observer<T> getObserver() {
         @SuppressWarnings("rawtypes")
         static final AtomicLongFieldUpdater<GroupBySubscriber> BUFFERED_COUNT = AtomicLongFieldUpdater.newUpdater(GroupBySubscriber.class, "bufferedCount");
 
+        volatile boolean errorEmitted = false;
+
         @Override
         public void onStart() {
             REQUESTED.set(this, MAX_QUEUE_SIZE);
@@ -166,6 +169,13 @@ public void onCompleted() {
         @Override
         public void onError(Throwable e) {
             if (TERMINATED_UPDATER.compareAndSet(this, 0, 1)) {
+                errorEmitted = true;
+
+                // It's safe to access all groups and emit the error.
+                // onNext and onError are in sequence so no group will be created in the loop.
+                for (GroupState<K, T> group : groups.values()) {
+                    emitItem(group, nl.error(e));
+                }
                 try {
                     // we immediately tear everything down if we receive an error
                     child.onError(e);
@@ -259,6 +269,11 @@ public void onCompleted() {
                         @Override
                         public void onError(Throwable e) {
                             o.onError(e);
+                            // eagerly cleanup instead of waiting for unsubscribe
+                            if (once.compareAndSet(false, true)) {
+                                // done once per instance, either onComplete or onUnSubscribe
+                                cleanupGroup(key);
+                            }
                         }
 
                         @Override
@@ -386,7 +401,7 @@ private void completeInner() {
                     if (child.isUnsubscribed()) {
                         // if the entire groupBy has been unsubscribed and children are completed we will propagate the unsubscribe up.
                         unsubscribe();
-                    } else {
+                    } else if (!errorEmitted) {
                         child.onCompleted();
                     }
                 }

src/test/java/rx/internal/operators/OperatorGroupByTest.java

@@ -1407,4 +1407,51 @@ public Integer call(Integer integer) {
         }).subscribe().unsubscribe();
         verify(s).unsubscribe();
     }
+
+    @Test
+    public void testGroupByShouldPropagateError() {
+        final Throwable e = new RuntimeException("Oops");
+        final TestSubscriber<Integer> inner1 = new TestSubscriber<Integer>();
+        final TestSubscriber<Integer> inner2 = new TestSubscriber<Integer>();
+
+        final TestSubscriber<GroupedObservable<Integer, Integer>> outer
+                = new TestSubscriber<GroupedObservable<Integer, Integer>>(new Subscriber<GroupedObservable<Integer, Integer>>() {
+
+            @Override
+            public void onCompleted() {
+            }
+
+            @Override
+            public void onError(Throwable e) {
+            }
+
+            @Override
+            public void onNext(GroupedObservable<Integer, Integer> o) {
+                if (o.getKey() == 0) {
+                    o.subscribe(inner1);
+                } else {
+                    o.subscribe(inner2);
+                }
+            }
+        });
+        Observable.create(
+                new OnSubscribe<Integer>() {
+                    @Override
+                    public void call(Subscriber<? super Integer> subscriber) {
+                        subscriber.onNext(0);
+                        subscriber.onNext(1);
+                        subscriber.onError(e);
+                    }
+                }
+        ).groupBy(new Func1<Integer, Integer>() {
+
+            @Override
+            public Integer call(Integer i) {
+                return i % 2;
+            }
+        }).subscribe(outer);
+        assertEquals(Arrays.asList(e), outer.getOnErrorEvents());
+        assertEquals(Arrays.asList(e), inner1.getOnErrorEvents());
+        assertEquals(Arrays.asList(e), inner2.getOnErrorEvents());
+    }
 }