2.x: Fix window(Observable|Callable) upstream handling (ReactiveX#5887)

Author: akarnokd

src/main/java/io/reactivex/internal/operators/observable/ObservableWindowBoundary.java

@@ -18,167 +18,199 @@
 import io.reactivex.*;
 import io.reactivex.disposables.Disposable;
 import io.reactivex.internal.disposables.DisposableHelper;
-import io.reactivex.internal.observers.QueueDrainObserver;
 import io.reactivex.internal.queue.MpscLinkedQueue;
-import io.reactivex.internal.util.NotificationLite;
-import io.reactivex.observers.*;
+import io.reactivex.internal.util.AtomicThrowable;
+import io.reactivex.observers.DisposableObserver;
 import io.reactivex.plugins.RxJavaPlugins;
 import io.reactivex.subjects.UnicastSubject;
 
 public final class ObservableWindowBoundary<T, B> extends AbstractObservableWithUpstream<T, Observable<T>> {
     final ObservableSource<B> other;
-    final int bufferSize;
+    final int capacityHint;
 
-    public ObservableWindowBoundary(ObservableSource<T> source, ObservableSource<B> other, int bufferSize) {
+    public ObservableWindowBoundary(ObservableSource<T> source, ObservableSource<B> other, int capacityHint) {
         super(source);
         this.other = other;
-        this.bufferSize = bufferSize;
+        this.capacityHint = capacityHint;
     }
 
     @Override
-    public void subscribeActual(Observer<? super Observable<T>> t) {
-        source.subscribe(new WindowBoundaryMainObserver<T, B>(new SerializedObserver<Observable<T>>(t), other, bufferSize));
+    public void subscribeActual(Observer<? super Observable<T>> observer) {
+        WindowBoundaryMainObserver<T, B> parent = new WindowBoundaryMainObserver<T, B>(observer, capacityHint);
+
+        observer.onSubscribe(parent);
+        other.subscribe(parent.boundaryObserver);
+
+        source.subscribe(parent);
     }
 
     static final class WindowBoundaryMainObserver<T, B>
-    extends QueueDrainObserver<T, Object, Observable<T>>
-    implements Disposable {
+    extends AtomicInteger
+    implements Observer<T>, Disposable, Runnable {
 
-        final ObservableSource<B> other;
-        final int bufferSize;
+        private static final long serialVersionUID = 2233020065421370272L;
 
-        Disposable s;
+        final Observer<? super Observable<T>> downstream;
 
-        final AtomicReference<Disposable> boundary = new AtomicReference<Disposable>();
+        final int capacityHint;
 
-        UnicastSubject<T> window;
+        final WindowBoundaryInnerObserver<T, B> boundaryObserver;
 
-        static final Object NEXT = new Object();
+        final AtomicReference<Disposable> upstream;
 
-        final AtomicLong windows = new AtomicLong();
+        final AtomicInteger windows;
 
-        WindowBoundaryMainObserver(Observer<? super Observable<T>> actual, ObservableSource<B> other,
-                int bufferSize) {
-            super(actual, new MpscLinkedQueue<Object>());
-            this.other = other;
-            this.bufferSize = bufferSize;
-            windows.lazySet(1);
-        }
+        final MpscLinkedQueue<Object> queue;
 
-        @Override
-        public void onSubscribe(Disposable s) {
-            if (DisposableHelper.validate(this.s, s)) {
-                this.s = s;
+        final AtomicThrowable errors;
 
-                Observer<? super Observable<T>> a = actual;
-                a.onSubscribe(this);
+        final AtomicBoolean stopWindows;
 
-                if (cancelled) {
-                    return;
-                }
+        static final Object NEXT_WINDOW = new Object();
 
-                UnicastSubject<T> w = UnicastSubject.create(bufferSize);
+        volatile boolean done;
 
-                window = w;
+        UnicastSubject<T> window;
 
-                a.onNext(w);
+        WindowBoundaryMainObserver(Observer<? super Observable<T>> downstream, int capacityHint) {
+            this.downstream = downstream;
+            this.capacityHint = capacityHint;
+            this.boundaryObserver = new WindowBoundaryInnerObserver<T, B>(this);
+            this.upstream = new AtomicReference<Disposable>();
+            this.windows = new AtomicInteger(1);
+            this.queue = new MpscLinkedQueue<Object>();
+            this.errors = new AtomicThrowable();
+            this.stopWindows = new AtomicBoolean();
+        }
 
-                WindowBoundaryInnerObserver<T, B> inner = new WindowBoundaryInnerObserver<T, B>(this);
+        @Override
+        public void onSubscribe(Disposable d) {
+            if (DisposableHelper.setOnce(upstream, d)) {
 
-                if (boundary.compareAndSet(null, inner)) {
-                    windows.getAndIncrement();
-                    other.subscribe(inner);
-                }
+                innerNext();
             }
         }
 
         @Override
         public void onNext(T t) {
-            if (fastEnter()) {
-                UnicastSubject<T> w = window;
-
-                w.onNext(t);
+            queue.offer(t);
+            drain();
+        }
 
-                if (leave(-1) == 0) {
-                    return;
-                }
+        @Override
+        public void onError(Throwable e) {
+            boundaryObserver.dispose();
+            if (errors.addThrowable(e)) {
+                done = true;
+                drain();
             } else {
-                queue.offer(NotificationLite.next(t));
-                if (!enter()) {
-                    return;
-                }
+                RxJavaPlugins.onError(e);
             }
-            drainLoop();
         }
 
         @Override
-        public void onError(Throwable t) {
-            if (done) {
-                RxJavaPlugins.onError(t);
-                return;
-            }
-            error = t;
+        public void onComplete() {
+            boundaryObserver.dispose();
             done = true;
-            if (enter()) {
-                drainLoop();
-            }
+            drain();
+        }
 
-            if (windows.decrementAndGet() == 0) {
-                DisposableHelper.dispose(boundary);
+        @Override
+        public void dispose() {
+            if (stopWindows.compareAndSet(false, true)) {
+                boundaryObserver.dispose();
+                if (windows.decrementAndGet() == 0) {
+                    DisposableHelper.dispose(upstream);
+                }
             }
-
-            actual.onError(t);
         }
 
         @Override
-        public void onComplete() {
-            if (done) {
-                return;
-            }
-            done = true;
-            if (enter()) {
-                drainLoop();
-            }
+        public boolean isDisposed() {
+            return stopWindows.get();
+        }
 
+        @Override
+        public void run() {
             if (windows.decrementAndGet() == 0) {
-                DisposableHelper.dispose(boundary);
+                DisposableHelper.dispose(upstream);
             }
+        }
 
-            actual.onComplete();
-
+        void innerNext() {
+            queue.offer(NEXT_WINDOW);
+            drain();
         }
 
-        @Override
-        public void dispose() {
-            cancelled = true;
+        void innerError(Throwable e) {
+            DisposableHelper.dispose(upstream);
+            if (errors.addThrowable(e)) {
+                done = true;
+                drain();
+            } else {
+                RxJavaPlugins.onError(e);
+            }
         }
 
-        @Override
-        public boolean isDisposed() {
-            return cancelled;
+        void innerComplete() {
+            DisposableHelper.dispose(upstream);
+            done = true;
+            drain();
         }
 
-        void drainLoop() {
-            final MpscLinkedQueue<Object> q = (MpscLinkedQueue<Object>)queue;
-            final Observer<? super Observable<T>> a = actual;
+        @SuppressWarnings("unchecked")
+        void drain() {
+            if (getAndIncrement() != 0) {
+                return;
+            }
+
             int missed = 1;
-            UnicastSubject<T> w = window;
+            Observer<? super Observable<T>> downstream = this.downstream;
+            MpscLinkedQueue<Object> queue = this.queue;
+            AtomicThrowable errors = this.errors;
+
             for (;;) {
 
                 for (;;) {
+                    if (windows.get() == 0) {
+                        queue.clear();
+                        window = null;
+                        return;
+                    }
+
+                    UnicastSubject<T> w = window;
+
                     boolean d = done;
 
-                    Object o = q.poll();
+                    if (d && errors.get() != null) {
+                        queue.clear();
+                        Throwable ex = errors.terminate();
+                        if (w != null) {
+                            window = null;
+                            w.onError(ex);
+                        }
+                        downstream.onError(ex);
+                        return;
+                    }
+
+                    Object v = queue.poll();
 
-                    boolean empty = o == null;
+                    boolean empty = v == null;
 
                     if (d && empty) {
-                        DisposableHelper.dispose(boundary);
-                        Throwable e = error;
-                        if (e != null) {
-                            w.onError(e);
+                        Throwable ex = errors.terminate();
+                        if (ex == null) {
+                            if (w != null) {
+                                window = null;
+                                w.onComplete();
+                            }
+                            downstream.onComplete();
                         } else {
-                            w.onComplete();
+                            if (w != null) {
+                                window = null;
+                                w.onError(ex);
+                            }
+                            downstream.onError(ex);
                         }
                         return;
                     }
@@ -187,48 +219,35 @@ void drainLoop() {
                         break;
                     }
 
-                    if (o == NEXT) {
-                        w.onComplete();
-
-                        if (windows.decrementAndGet() == 0) {
-                            DisposableHelper.dispose(boundary);
-                            return;
-                        }
-
-                        if (cancelled) {
-                            continue;
-                        }
-
-                        w = UnicastSubject.create(bufferSize);
+                    if (v != NEXT_WINDOW) {
+                        w.onNext((T)v);
+                        continue;
+                    }
 
-                        windows.getAndIncrement();
+                    if (w != null) {
+                        window = null;
+                        w.onComplete();
+                    }
 
+                    if (!stopWindows.get()) {
+                        w = UnicastSubject.create(capacityHint, this);
                         window = w;
+                        windows.getAndIncrement();
 
-                        a.onNext(w);
-
-                        continue;
+                        downstream.onNext(w);
                     }
-
-                    w.onNext(NotificationLite.<T>getValue(o));
                 }
 
-                missed = leave(-missed);
+                missed = addAndGet(-missed);
                 if (missed == 0) {
-                    return;
+                    break;
                 }
             }
         }
-
-        void next() {
-            queue.offer(NEXT);
-            if (enter()) {
-                drainLoop();
-            }
-        }
     }
 
     static final class WindowBoundaryInnerObserver<T, B> extends DisposableObserver<B> {
+
         final WindowBoundaryMainObserver<T, B> parent;
 
         boolean done;
@@ -242,7 +261,7 @@ public void onNext(B t) {
             if (done) {
                 return;
             }
-            parent.next();
+            parent.innerNext();
         }
 
         @Override
@@ -252,7 +271,7 @@ public void onError(Throwable t) {
                 return;
             }
             done = true;
-            parent.onError(t);
+            parent.innerError(t);
         }
 
         @Override
@@ -261,7 +280,7 @@ public void onComplete() {
                 return;
             }
             done = true;
-            parent.onComplete();
+            parent.innerComplete();
         }
     }
 }