Merge pull request ReactiveX#596 from akarnokd/BufferFix1

Fix for buffer not stopping when unsubscribed.
rickbw · Dec 10, 2013 · c8be88d · c8be88d
2 parents 65b3f6f + bcc779d
commit c8be88d
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 11 deletions.
diff --git a/rxjava-core/src/main/java/rx/operators/ChunkedOperation.java b/rxjava-core/src/main/java/rx/operators/ChunkedOperation.java
@@ -148,7 +148,7 @@ public OverlappingChunks(Observer<? super C> observer, Func0<? extends Chunk<T,
      *            The type of object all internal {@link rx.operators.ChunkedOperation.Chunk} objects record.
      *            <C> The type of object being tracked by the {@link Chunk}
      */
-    protected static class TimeAndSizeBasedChunks<T, C> extends Chunks<T, C> {
+    protected static class TimeAndSizeBasedChunks<T, C> extends Chunks<T, C> implements Subscription {
 
         private final ConcurrentMap<Chunk<T, C>, Subscription> subscriptions = new ConcurrentHashMap<Chunk<T, C>, Subscription>();
 
@@ -207,6 +207,12 @@ public void pushValue(T value) {
                 }
             }
         }
+        @Override
+        public void unsubscribe() {
+            for (Subscription s : subscriptions.values()) {
+                s.unsubscribe();
+            }
+        }
     }
 
     /**
@@ -218,7 +224,7 @@ public void pushValue(T value) {
      *            The type of object all internal {@link rx.operators.ChunkedOperation.Chunk} objects record.
      *            <C> The type of object being tracked by the {@link Chunk}
      */
-    protected static class TimeBasedChunks<T, C> extends OverlappingChunks<T, C> {
+    protected static class TimeBasedChunks<T, C> extends OverlappingChunks<T, C> implements Subscription {
 
         private final ConcurrentMap<Chunk<T, C>, Subscription> subscriptions = new ConcurrentHashMap<Chunk<T, C>, Subscription>();
 
@@ -250,6 +256,14 @@ public void emitChunk(Chunk<T, C> chunk) {
             subscriptions.remove(chunk);
             super.emitChunk(chunk);
         }
+
+        @Override
+        public void unsubscribe() {
+            for (Subscription s : subscriptions.values()) {
+                s.unsubscribe();
+            }
+        }
+
     }
 
     /**

diff --git a/rxjava-core/src/main/java/rx/operators/OperationBuffer.java b/rxjava-core/src/main/java/rx/operators/OperationBuffer.java
@@ -17,13 +17,15 @@
 
 import java.util.List;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import rx.Observable;
 import rx.Observable.OnSubscribeFunc;
 import rx.Observer;
 import rx.Scheduler;
 import rx.Subscription;
 import rx.concurrency.Schedulers;
+import rx.subscriptions.CompositeSubscription;
 import rx.util.functions.Func0;
 import rx.util.functions.Func1;
 
@@ -65,11 +67,14 @@ public static <T, TClosing> OnSubscribeFunc<List<T>> buffer(final Observable<T>
             public Subscription onSubscribe(Observer<? super List<T>> observer) {
                 NonOverlappingChunks<T, List<T>> buffers = new NonOverlappingChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker());
                 ChunkCreator creator = new ObservableBasedSingleChunkCreator<T, List<T>, TClosing>(buffers, bufferClosingSelector);
-                return source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator));
+                return new CompositeSubscription(
+                        new ChunkToSubscription(creator),
+                        source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator))
+                );
             }
         };
     }
-
+    
     /**
      * <p>This method creates a {@link Func1} object which represents the buffer operation. This operation takes
      * values from the specified {@link Observable} source and stores them in the currently active chunks. Initially
@@ -101,7 +106,10 @@ public static <T, TOpening, TClosing> OnSubscribeFunc<List<T>> buffer(final Obse
             public Subscription onSubscribe(final Observer<? super List<T>> observer) {
                 OverlappingChunks<T, List<T>> buffers = new OverlappingChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker());
                 ChunkCreator creator = new ObservableBasedMultiChunkCreator<T, List<T>, TOpening, TClosing>(buffers, bufferOpenings, bufferClosingSelector);
-                return source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator));
+                return new CompositeSubscription(
+                        new ChunkToSubscription(creator),
+                        source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator))
+                );
             }
         };
     }
@@ -156,7 +164,10 @@ public static <T> OnSubscribeFunc<List<T>> buffer(final Observable<T> source, fi
             public Subscription onSubscribe(final Observer<? super List<T>> observer) {
                 Chunks<T, List<T>> chunks = new SizeBasedChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker(), count);
                 ChunkCreator creator = new SkippingChunkCreator<T, List<T>>(chunks, skip);
-                return source.subscribe(new ChunkObserver<T, List<T>>(chunks, observer, creator));
+                return new CompositeSubscription(
+                        new ChunkToSubscription(creator),
+                        source.subscribe(new ChunkObserver<T, List<T>>(chunks, observer, creator))
+                );
             }
         };
     }
@@ -211,7 +222,10 @@ public static <T> OnSubscribeFunc<List<T>> buffer(final Observable<T> source, fi
             public Subscription onSubscribe(final Observer<? super List<T>> observer) {
                 NonOverlappingChunks<T, List<T>> buffers = new NonOverlappingChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker());
                 ChunkCreator creator = new TimeBasedChunkCreator<T, List<T>>(buffers, timespan, unit, scheduler);
-                return source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator));
+                return new CompositeSubscription(
+                        new ChunkToSubscription(creator),
+                        source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator))
+                );
             }
         };
     }
@@ -270,9 +284,13 @@ public static <T> OnSubscribeFunc<List<T>> buffer(final Observable<T> source, fi
         return new OnSubscribeFunc<List<T>>() {
             @Override
             public Subscription onSubscribe(final Observer<? super List<T>> observer) {
-                Chunks<T, List<T>> chunks = new TimeAndSizeBasedChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker(), count, timespan, unit, scheduler);
+                TimeAndSizeBasedChunks<T, List<T>> chunks = new TimeAndSizeBasedChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker(), count, timespan, unit, scheduler);
                 ChunkCreator creator = new SingleChunkCreator<T, List<T>>(chunks);
-                return source.subscribe(new ChunkObserver<T, List<T>>(chunks, observer, creator));
+                return new CompositeSubscription(
+                        chunks,
+                        new ChunkToSubscription(creator),
+                        source.subscribe(new ChunkObserver<T, List<T>>(chunks, observer, creator))
+                );
             }
         };
     }
@@ -331,9 +349,13 @@ public static <T> OnSubscribeFunc<List<T>> buffer(final Observable<T> source, fi
         return new OnSubscribeFunc<List<T>>() {
             @Override
             public Subscription onSubscribe(final Observer<? super List<T>> observer) {
-                OverlappingChunks<T, List<T>> buffers = new TimeBasedChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker(), timespan, unit, scheduler);
+                TimeBasedChunks<T, List<T>> buffers = new TimeBasedChunks<T, List<T>>(observer, OperationBuffer.<T> bufferMaker(), timespan, unit, scheduler);
                 ChunkCreator creator = new TimeBasedChunkCreator<T, List<T>>(buffers, timeshift, unit, scheduler);
-                return source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator));
+                return new CompositeSubscription(
+                        buffers,
+                        new ChunkToSubscription(creator),
+                        source.subscribe(new ChunkObserver<T, List<T>>(buffers, observer, creator))
+                );
             }
         };
     }
@@ -355,4 +377,24 @@ public List<T> getContents() {
             return contents;
         }
     }
+
+    /**
+     * Converts a chunk creator into a subscription which stops the chunk.
+     */
+    private static class ChunkToSubscription implements Subscription {
+        private ChunkCreator cc;
+        private final AtomicBoolean done;
+        public ChunkToSubscription(ChunkCreator cc) {
+            this.cc = cc;
+            this.done = new AtomicBoolean();
+        }
+        @Override
+        public void unsubscribe() {
+            if (done.compareAndSet(false, true)) {
+                ChunkCreator cc0 = cc;
+                cc = null;
+                cc0.stop();
+            }
+        }
+    }
 }
diff --git a/rxjava-core/src/test/java/rx/operators/OperationBufferTest.java b/rxjava-core/src/test/java/rx/operators/OperationBufferTest.java
@@ -19,6 +19,7 @@
 import static rx.operators.OperationBuffer.*;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -27,6 +28,8 @@
 import org.junit.Test;
 import org.mockito.InOrder;
 import org.mockito.Mockito;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 
 import rx.Observable;
 import rx.Observer;
@@ -359,4 +362,25 @@ public void call() {
             }
         }, delay, TimeUnit.MILLISECONDS);
     }
+
+    @Test
+    public void testBufferStopsWhenUnsubscribed1() {
+        Observable<Integer> source = Observable.never();
+
+        Observer<List<Integer>> o = mock(Observer.class);
+
+        Subscription s = source.buffer(100, 200, TimeUnit.MILLISECONDS, scheduler).subscribe(o);
+
+        InOrder inOrder = Mockito.inOrder(o);
+
+        scheduler.advanceTimeBy(1001, TimeUnit.MILLISECONDS);
+
+        inOrder.verify(o, times(5)).onNext(Arrays.<Integer>asList());
+
+        s.unsubscribe();
+
+        scheduler.advanceTimeBy(999, TimeUnit.MILLISECONDS);
+
+        inOrder.verifyNoMoreInteractions();
+    }
 }