GroupBy & Time Gap Fixes

I ran head-on into the "time gap" (ReactiveX#844) issue while working on a stream processing use case (and new 'pivot' operator I'm writing).
This is a solution. It's still not ideal as the Javadocs of BufferUntilSubscriber mention, but this is working better than nothing and does not require blocking threads.
A better solution will come as part of the back pressure work where BufferUntilSubscriber will evolve to have a bounded buffer.

Author: benjchristensen

rxjava-core/src/main/java/rx/operators/BufferUntilSubscriber.java

@@ -15,152 +15,175 @@
  */
 package rx.operators;
 
-import java.util.LinkedList;
-import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicReference;
 
+import rx.Observable;
+import rx.Observer;
 import rx.Subscriber;
-import rx.subscriptions.CompositeSubscription;
 
 /**
- * Buffers the incoming events until notified, then replays the
- * buffered events and continues as a simple pass-through subscriber.
- * @param <T> the streamed value type
+ * A solution to the "time gap" problem that occurs with `groupBy` and `pivot` => https://github.com/Netflix/RxJava/issues/844
+ * 
+ * This currently has temporary unbounded buffers. It needs to become bounded and then do one of two things:
+ * 
+ * 1) blow up and make the user do something about it
+ * 2) work with the backpressure solution ... still to be implemented (such as co-routines)
+ * 
+ * Generally the buffer should be very short lived (milliseconds) and then stops being involved.
+ * It can become a memory leak though if a GroupedObservable backed by this class is emitted but never subscribed to (such as filtered out).
+ * In that case, either a time-bomb to throw away the buffer, or just blowing up and making the user do something about it is needed.
+ * 
+ * For example, to filter out GroupedObservables, perhaps they need a silent `subscribe()` on them to just blackhole the data.
+ * 
+ * This is an initial start at solving this problem and solves the immediate problem of `groupBy` and `pivot` and trades off the possibility of memory leak for deterministic functionality.
+ *
+ * @param <T>
  */
-public class BufferUntilSubscriber<T> extends Subscriber<T> {
-    /** The actual subscriber. */
-    private final Subscriber<? super T> actual;
-    /** Indicate the pass-through mode. */
-    private volatile boolean passthroughMode;
-    /** Protect mode transition. */
-    private final Object gate = new Object();
-    /** The buffered items. */
-    private final Queue<Object> queue = new LinkedList<Object>();
-    /** The queue capacity. */
-    private final int capacity;
-    private final NotificationLite<T> on = NotificationLite.instance();
+public class BufferUntilSubscriber<T> extends Observable<T> implements Observer<T> {
 
-    /**
-     * Constructor that wraps the actual subscriber and shares its subscription.
-     * @param capacity the queue capacity to accept before blocking, negative value indicates an unbounded queue
-     * @param actual
-     */
-    public BufferUntilSubscriber(int capacity, Subscriber<? super T> actual) {
-        super(actual);
-        this.actual = actual;
-        this.capacity = capacity;
-    }
-    /**
-     * Constructor that wraps the actual subscriber and uses the given composite
-     * subscription.
-     * @param capacity the queue capacity to accept before blocking, negative value indicates an unbounded queue
-     * @param actual
-     * @param cs 
-     */
-    public BufferUntilSubscriber(int capacity, Subscriber<? super T> actual, CompositeSubscription cs) {
-        super(cs);
-        this.actual = actual;
-        this.capacity = capacity;
+    public static <T> BufferUntilSubscriber<T> create() {
+        return new BufferUntilSubscriber<T>(new AtomicReference<Observer<? super T>>(new BufferedObserver<T>()));
     }
-    
-    /**
-     * Call this method to replay the buffered events and continue as a pass-through subscriber.
-     * If already in pass-through mode, this method is a no-op.
-     */
-    public void enterPassthroughMode() {
-        if (!passthroughMode) {
-            synchronized (gate) {
-                if (!passthroughMode) {
-                    while (!queue.isEmpty()) {
-                        Object o = queue.poll();
-                        if (!actual.isUnsubscribed()) {
-                            on.accept(actual, o);
-                        }
-                    }
-                    passthroughMode = true;
-                    gate.notifyAll();
+
+    private final AtomicReference<Observer<? super T>> observerRef;
+
+    private BufferUntilSubscriber(final AtomicReference<Observer<? super T>> observerRef) {
+        super(new OnSubscribe<T>() {
+
+            @Override
+            public void call(Subscriber<? super T> s) {
+                // drain queued notifications before subscription
+                // we do this here before PassThruObserver so the consuming thread can do this before putting itself in the line of the producer
+                BufferedObserver<T> buffered = (BufferedObserver<T>) observerRef.get();
+                Object o = null;
+                while ((o = buffered.buffer.poll()) != null) {
+                    emit(s, o);
                 }
+                // register real observer for pass-thru ... and drain any further events received on first notification
+                observerRef.set(new PassThruObserver<T>(s, buffered.buffer, observerRef));
             }
-        }
+
+        });
+        this.observerRef = observerRef;
+    }
+
+    @Override
+    public void onCompleted() {
+        observerRef.get().onCompleted();
     }
+
+    @Override
+    public void onError(Throwable e) {
+        observerRef.get().onError(e);
+    }
+
     @Override
     public void onNext(T t) {
-        if (!passthroughMode) {
-            synchronized (gate) {
-                if (!passthroughMode) {
-                    if (capacity < 0 || queue.size() < capacity) {
-                        queue.offer(on.next(t));
-                        return;
-                    }
-                    try {
-                        while (!passthroughMode) {
-                            gate.wait();
-                        }
-                        if (actual.isUnsubscribed()) {
-                            return;
-                        }
-                    } catch (InterruptedException ex) {
-                        Thread.currentThread().interrupt();
-                        actual.onError(ex);
-                        return;
-                    }
-                }
+        observerRef.get().onNext(t);
+    }
+
+    /**
+     * This is a temporary observer between buffering and the actual that gets into the line of notifications
+     * from the producer and will drain the queue of any items received during the race of the initial drain and
+     * switching this.
+     * 
+     * It will then immediately swap itself out for the actual (after a single notification), but since this is now
+     * being done on the same producer thread no further buffering will occur.
+     */
+    private static class PassThruObserver<T> implements Observer<T> {
+
+        private final Observer<? super T> actual;
+        // this assumes single threaded synchronous notifications (the Rx contract for a single Observer)
+        private final ConcurrentLinkedQueue<Object> buffer;
+        private final AtomicReference<Observer<? super T>> observerRef;
+
+        PassThruObserver(Observer<? super T> actual, ConcurrentLinkedQueue<Object> buffer, AtomicReference<Observer<? super T>> observerRef) {
+            this.actual = actual;
+            this.buffer = buffer;
+            this.observerRef = observerRef;
+        }
+
+        @Override
+        public void onCompleted() {
+            drainIfNeededAndSwitchToActual();
+            actual.onCompleted();
+        }
+
+        @Override
+        public void onError(Throwable e) {
+            drainIfNeededAndSwitchToActual();
+            actual.onError(e);
+        }
+
+        @Override
+        public void onNext(T t) {
+            drainIfNeededAndSwitchToActual();
+            actual.onNext(t);
+        }
+
+        private void drainIfNeededAndSwitchToActual() {
+            Object o = null;
+            while ((o = buffer.poll()) != null) {
+                emit(this, o);
             }
+            // now we can safely change over to the actual and get rid of the pass-thru
+            observerRef.set(actual);
         }
-        actual.onNext(t);
+
     }
 
-    @Override
-    public void onError(Throwable e) {
-        if (!passthroughMode) {
-            synchronized (gate) {
-                if (!passthroughMode) {
-                    if (capacity < 0 || queue.size() < capacity) {
-                        queue.offer(on.error(e));
-                        return;
-                    }
-                    try {
-                        while (!passthroughMode) {
-                            gate.wait();
-                        }
-                        if (actual.isUnsubscribed()) {
-                            return;
-                        }
-                    } catch (InterruptedException ex) {
-                        Thread.currentThread().interrupt();
-                        actual.onError(ex);
-                        return;
-                    }
-                }
+    private static class BufferedObserver<T> implements Observer<T> {
+        private final ConcurrentLinkedQueue<Object> buffer = new ConcurrentLinkedQueue<Object>();
+
+        @Override
+        public void onCompleted() {
+            buffer.add(COMPLETE_SENTINEL);
+        }
+
+        @Override
+        public void onError(Throwable e) {
+            buffer.add(new ErrorSentinel(e));
+        }
+
+        @Override
+        public void onNext(T t) {
+            if (t == null) {
+                buffer.add(NULL_SENTINEL);
+            } else {
+                buffer.add(t);
             }
         }
-        actual.onError(e);
+
     }
 
-    @Override
-    public void onCompleted() {
-        if (!passthroughMode) {
-            synchronized (gate) {
-                if (!passthroughMode) {
-                    if (capacity < 0 || queue.size() < capacity) {
-                        queue.offer(on.completed());
-                        return;
-                    }
-                    try {
-                        while (!passthroughMode) {
-                            gate.wait();
-                        }
-                        if (actual.isUnsubscribed()) {
-                            return;
-                        }
-                    } catch (InterruptedException ex) {
-                        Thread.currentThread().interrupt();
-                        actual.onError(ex);
-                        return;
-                    }
-                }
+    private final static <T> void emit(Observer<T> s, Object v) {
+        if (v instanceof Sentinel) {
+            if (v == NULL_SENTINEL) {
+                s.onNext(null);
+            } else if (v == COMPLETE_SENTINEL) {
+                s.onCompleted();
+            } else if (v instanceof ErrorSentinel) {
+                s.onError(((ErrorSentinel) v).e);
             }
+        } else {
+            s.onNext((T) v);
+        }
+    }
+
+    private static class Sentinel {
+
+    }
+
+    private static Sentinel NULL_SENTINEL = new Sentinel();
+    private static Sentinel COMPLETE_SENTINEL = new Sentinel();
+
+    private static class ErrorSentinel extends Sentinel {
+        final Throwable e;
+
+        ErrorSentinel(Throwable e) {
+            this.e = e;
         }
-        actual.onCompleted();
     }
 
 }

rxjava-core/src/main/java/rx/operators/OperatorGroupBy.java

@@ -27,8 +27,6 @@
 import rx.functions.Action0;
 import rx.functions.Func1;
 import rx.observables.GroupedObservable;
-import rx.subjects.PublishSubject;
-import rx.subjects.Subject;
 import rx.subscriptions.CompositeSubscription;
 import rx.subscriptions.Subscriptions;
 
@@ -51,43 +49,50 @@ public Subscriber<? super T> call(final Subscriber<? super GroupedObservable<K,
         // a new CompositeSubscription to decouple the subscription as the inner subscriptions need a separate lifecycle
         // and will unsubscribe on this parent if they are all unsubscribed
         return new Subscriber<T>(new CompositeSubscription()) {
-            private final Map<K, Subject<T, T>> groups = new HashMap<K, Subject<T, T>>();
+            private final Map<K, BufferUntilSubscriber<T>> groups = new HashMap<K, BufferUntilSubscriber<T>>();
             private final AtomicInteger completionCounter = new AtomicInteger(0);
-            private final AtomicBoolean completed = new AtomicBoolean(false);
+            private final AtomicBoolean completionEmitted = new AtomicBoolean(false);
+            private final AtomicBoolean terminated = new AtomicBoolean(false);
 
             @Override
             public void onCompleted() {
-                completed.set(true);
-                // if we receive onCompleted from our parent we onComplete children
-                for (Subject<T, T> ps : groups.values()) {
-                    ps.onCompleted();
-                }
+                if (terminated.compareAndSet(false, true)) {
+                    // if we receive onCompleted from our parent we onComplete children
+                    for (BufferUntilSubscriber<T> ps : groups.values()) {
+                        ps.onCompleted();
+                    }
 
-                // special case for empty (no groups emitted)
-                if (completionCounter.get() == 0) {
-                    childObserver.onCompleted();
+                    // special case for empty (no groups emitted)
+                    if (completionCounter.get() == 0) {
+                        // we must track 'completionEmitted' seperately from 'completed' since `completeInner` can result in childObserver.onCompleted() being emitted 
+                        if (completionEmitted.compareAndSet(false, true)) {
+                            childObserver.onCompleted();
+                        }
+                    }
                 }
             }
 
             @Override
             public void onError(Throwable e) {
-                // we immediately tear everything down if we receive an error
-                childObserver.onError(e);
+                if (terminated.compareAndSet(false, true)) {
+                    // we immediately tear everything down if we receive an error
+                    childObserver.onError(e);
+                }
             }
 
             @Override
             public void onNext(T t) {
                 try {
                     final K key = keySelector.call(t);
-                    Subject<T, T> gps = groups.get(key);
+                    BufferUntilSubscriber<T> gps = groups.get(key);
                     if (gps == null) {
                         // this group doesn't exist
                         if (childObserver.isUnsubscribed()) {
                             // we have been unsubscribed on the outer so won't send any  more groups 
                             return;
                         }
-                        gps = PublishSubject.create();
-                        final Subject<T, T> _gps = gps;
+                        gps = BufferUntilSubscriber.create();
+                        final BufferUntilSubscriber<T> _gps = gps;
 
                         GroupedObservable<K, T> go = new GroupedObservable<K, T>(key, new OnSubscribe<T>() {
 
@@ -136,15 +141,16 @@ public void onNext(T t) {
             }
 
             private void completeInner() {
-                if (completionCounter.decrementAndGet() == 0 && (completed.get() || childObserver.isUnsubscribed())) {
-                    if (childObserver.isUnsubscribed()) {
-                        // if the entire groupBy has been unsubscribed and children are completed we will propagate the unsubscribe up.
-                        unsubscribe();
-                    }
-                    for (Subject<T, T> ps : groups.values()) {
-                        ps.onCompleted();
+                // count can be < 0 because unsubscribe also calls this
+                if (completionCounter.decrementAndGet() <= 0 && (terminated.get() || childObserver.isUnsubscribed())) {
+                    // completionEmitted ensures we only emit onCompleted once
+                    if (completionEmitted.compareAndSet(false, true)) {
+                        if (childObserver.isUnsubscribed()) {
+                            // if the entire groupBy has been unsubscribed and children are completed we will propagate the unsubscribe up.
+                            unsubscribe();
+                        }
+                        childObserver.onCompleted();
                     }
-                    childObserver.onCompleted();
                 }
             }