2.x fix group-by eviction so that source is cancelled and reduce volatile reads (#5933) (#5947)

Author: davidmoten

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

@@ -104,7 +104,8 @@ public static final class GroupBySubscriber<T, K, V>
         final AtomicInteger groupCount = new AtomicInteger(1);
 
         Throwable error;
-        volatile boolean done;
+        volatile boolean finished;
+        boolean done;
 
         boolean outputFused;
 
@@ -178,12 +179,7 @@ public void onNext(T t) {
 
             group.onNext(v);
 
-            if (evictedGroups != null) {
-                GroupedUnicast<K, V> evictedGroup;
-                while ((evictedGroup = evictedGroups.poll()) != null) {
-                    evictedGroup.onComplete();
-                }
-            }
+            completeEvictions();
 
             if (newGroup) {
                 q.offer(group);
@@ -197,6 +193,7 @@ public void onError(Throwable t) {
                 RxJavaPlugins.onError(t);
                 return;
             }
+            done = true;
             for (GroupedUnicast<K, V> g : groups.values()) {
                 g.onError(t);
             }
@@ -205,7 +202,7 @@ public void onError(Throwable t) {
                 evictedGroups.clear();
             }
             error = t;
-            done = true;
+            finished = true;
             drain();
         }
 
@@ -220,6 +217,7 @@ public void onComplete() {
                     evictedGroups.clear();
                 }
                 done = true;
+                finished = true;
                 drain();
             }
         }
@@ -237,12 +235,27 @@ 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(false, true)) {
+                completeEvictions();
                 if (groupCount.decrementAndGet() == 0) {
                     s.cancel();
                 }
             }
         }
 
+        private void completeEvictions() {
+            if (evictedGroups != null) {
+                int count = 0;
+                GroupedUnicast<K, V> evictedGroup;
+                while ((evictedGroup = evictedGroups.poll()) != null) {
+                    evictedGroup.onComplete();
+                    count++;
+                }
+                if (count != 0) {
+                    groupCount.addAndGet(-count);
+                }
+            }
+        }
+
         public void cancel(K key) {
             Object mapKey = key != null ? key : NULL_KEY;
             groups.remove(mapKey);
@@ -279,7 +292,7 @@ void drainFused() {
                     return;
                 }
 
-                boolean d = done;
+                boolean d = finished;
 
                 if (d && !delayError) {
                     Throwable ex = error;
@@ -321,7 +334,7 @@ void drainNormal() {
                 long e = 0L;
 
                 while (e != r) {
-                    boolean d = done;
+                    boolean d = finished;
 
                     GroupedFlowable<K, V> t = q.poll();
 
@@ -340,7 +353,7 @@ void drainNormal() {
                     e++;
                 }
 
-                if (e == r && checkTerminated(done, q.isEmpty(), a, q)) {
+                if (e == r && checkTerminated(finished, q.isEmpty(), a, q)) {
                     return;
                 }
 

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

@@ -26,6 +26,7 @@
 import org.mockito.Mockito;
 import org.reactivestreams.*;
 
+import com.google.common.base.Ticker;
 import com.google.common.cache.*;
 
 import io.reactivex.*;
@@ -1947,6 +1948,127 @@ public void run() throws Exception {
             }
         };
     }
+    
+    private static final class TestTicker extends Ticker {
+        long tick = 0;
+
+        @Override
+        public long read() {
+            return tick;
+        }
+    }
+    
+    @Test
+    public void testGroupByEvictionCancellationOfSource5933() {
+        PublishProcessor<Integer> source = PublishProcessor.create();
+        final TestTicker testTicker = new TestTicker();
+
+        Function<Consumer<Object>, Map<Integer, Object>> mapFactory = new Function<Consumer<Object>, Map<Integer, Object>>() {
+            @Override
+            public Map<Integer, Object> apply(final Consumer<Object> action) throws Exception {
+                return CacheBuilder.newBuilder() //
+                        .expireAfterAccess(5, TimeUnit.SECONDS).removalListener(new RemovalListener<Object, Object>() {
+                            @Override
+                            public void onRemoval(RemovalNotification<Object, Object> notification) {
+                                try {
+                                    action.accept(notification.getValue());
+                                } catch (Exception ex) {
+                                    throw new RuntimeException(ex);
+                                }
+                            }
+                        }).ticker(testTicker) //
+                        .<Integer, Object>build().asMap();
+            }
+        };
+
+        final List<String> list = new CopyOnWriteArrayList<String>();
+        Flowable<Integer> stream = source //
+                .doOnCancel(new Action() {
+                    @Override
+                    public void run() throws Exception {
+                        list.add("Source canceled");
+                    }
+                })
+                .<Integer, Integer>groupBy(Functions.<Integer>identity(), Functions.<Integer>identity(), false,
+                        Flowable.bufferSize(), mapFactory) //
+                .flatMap(new Function<GroupedFlowable<Integer, Integer>, Publisher<? extends Integer>>() {
+                    @Override
+                    public Publisher<? extends Integer> apply(GroupedFlowable<Integer, Integer> group)
+                            throws Exception {
+                        return group //
+                                .doOnComplete(new Action() {
+                                    @Override
+                                    public void run() throws Exception {
+                                        list.add("Group completed");
+                                    }
+                                }).doOnCancel(new Action() {
+                                    @Override
+                                    public void run() throws Exception {
+                                        list.add("Group canceled");
+                                    }
+                                });
+                    }
+                });
+        TestSubscriber<Integer> ts = stream //
+                .doOnCancel(new Action() {
+                    @Override
+                    public void run() throws Exception {
+                        list.add("Outer group by canceled");
+                    }
+                }).test();
+
+        // Send 3 in the same group and wait for them to be seen
+        source.onNext(1);
+        source.onNext(1);
+        source.onNext(1);
+        ts.awaitCount(3);
+        
+        // Advance time far enough to evict the group.
+        // NOTE -- Comment this line out to make the test "pass".
+        testTicker.tick = TimeUnit.SECONDS.toNanos(6);
+        
+        // Send more data in the group (triggering eviction and recreation)
+        source.onNext(1);
+
+        // Wait for the last 2 and then cancel the subscription
+        ts.awaitCount(4);
+        ts.cancel();
+
+        // Observe the result.  Note that right now the result differs depending on whether eviction occurred or
+        // not.  The observed sequence in that case is:  Group completed, Outer group by canceled., Group canceled.
+        // The addition of the "Group completed" is actually fine, but the fact that the cancel doesn't reach the
+        // source seems like a bug.  Commenting out the setting of "tick" above will produce the "expected" sequence.
+        System.out.println(list);
+        assertTrue(list.contains("Source canceled"));
+        assertEquals(Arrays.asList(
+                "Group completed", // this is here when eviction occurs
+                "Outer group by canceled", 
+                "Group canceled",
+                "Source canceled"  // This is *not* here when eviction occurs
+        ), list);
+    }
+    
+    @Test
+    public void testCancellationOfUpstreamWhenGroupedFlowableCompletes() {
+        final AtomicBoolean cancelled = new AtomicBoolean();
+        Flowable.just(1).repeat().doOnCancel(new Action() {
+            @Override
+            public void run() throws Exception {
+                cancelled.set(true);
+            }
+        })
+        .groupBy(Functions.<Integer>identity(), Functions.<Integer>identity()) //
+        .flatMap(new Function<GroupedFlowable<Integer, Integer>, Publisher<? extends Object>>() {
+            @Override
+            public Publisher<? extends Object> apply(GroupedFlowable<Integer, Integer> g) throws Exception {
+                return g.first(0).toFlowable();
+            }
+        })
+        .take(4) //
+        .test() //
+        .assertComplete();
+        assertTrue(cancelled.get());
+    }
 
     //not thread safe
     private static final class SingleThreadEvictingHashMap<K,V> implements Map<K,V> {