Merge pull request #2552 from akarnokd/OperatorPublishRequestFix

akarnokd · akarnokd · commit fe3cc750319b · 2015-01-31T17:53:15.000+01:00
Publish: fixed incorrect subscriber requested accounting
diff --git a/src/main/java/rx/internal/operators/OperatorPublish.java b/src/main/java/rx/internal/operators/OperatorPublish.java
@@ -206,7 +206,6 @@ public void onNext(T t) {
      */
     private static class State<T> {
         private long outstandingRequests = -1;
-        private long emittedSinceRequest = 0;
         private OriginSubscriber<T> origin;
         // using AtomicLong to simplify mutating it, not for thread-safety since we're synchronizing access to this class
         // using LinkedHashMap so the order of Subscribers having onNext invoked is deterministic (same each time the code is run)
@@ -225,15 +224,13 @@ public synchronized void setOrigin(OriginSubscriber<T> o) {
         public synchronized boolean canEmitWithDecrement() {
             if (outstandingRequests > 0) {
                 outstandingRequests--;
-                emittedSinceRequest++;
                 return true;
             }
             return false;
         }
 
         public synchronized void incrementOutstandingAfterFailedEmit() {
             outstandingRequests++;
-            emittedSinceRequest--;
         }
 
         public synchronized Subscriber<? super T>[] getSubscribers() {
@@ -243,50 +240,55 @@ public synchronized Subscriber<? super T>[] getSubscribers() {
         /**
          * @return long outstandingRequests
          */
-        public synchronized long requestFromSubscriber(Subscriber<? super T> subscriber, Long request) {
-            AtomicLong r = ss.get(subscriber);
+        public synchronized long requestFromSubscriber(Subscriber<? super T> subscriber, long request) {
+            Map<Subscriber<? super T>, AtomicLong> subs = ss;
+            AtomicLong r = subs.get(subscriber);
             if (r == null) {
-                ss.put(subscriber, new AtomicLong(request));
+                subs.put(subscriber, new AtomicLong(request));
             } else {
-                if (r.get() != Long.MAX_VALUE) {
-                    if (request == Long.MAX_VALUE) {
-                        r.set(Long.MAX_VALUE);
-                    } else {
-                        r.addAndGet(request.longValue());
+                do {
+                    long current = r.get();
+                    if (current == Long.MAX_VALUE) {
+                        break;
                     }
-                }
+                    long u = current + request;
+                    if (u < 0) {
+                        u = Long.MAX_VALUE;
+                    }
+                    if (r.compareAndSet(current, u)) {
+                        break;
+                    }
+                } while (true);
             }
 
-            return resetAfterSubscriberUpdate();
+            return resetAfterSubscriberUpdate(subs);
         }
 
         public synchronized void removeSubscriber(Subscriber<? super T> subscriber) {
-            ss.remove(subscriber);
-            resetAfterSubscriberUpdate();
+            Map<Subscriber<? super T>, AtomicLong> subs = ss;
+            subs.remove(subscriber);
+            resetAfterSubscriberUpdate(subs);
         }
 
         @SuppressWarnings("unchecked")
-        private long resetAfterSubscriberUpdate() {
-            subscribers = new Subscriber[ss.size()];
+        private long resetAfterSubscriberUpdate(Map<Subscriber<? super T>, AtomicLong> subs) {
+            Subscriber<? super T>[] subscriberArray = new Subscriber[subs.size()];
             int i = 0;
-            for (Subscriber<? super T> s : ss.keySet()) {
-                subscribers[i++] = s;
-            }
-
             long lowest = -1;
-            for (AtomicLong l : ss.values()) {
-                // decrement all we have emitted since last request
-                long c = l.addAndGet(-emittedSinceRequest);
+            for (Map.Entry<Subscriber<? super T>, AtomicLong> e : subs.entrySet()) {
+                subscriberArray[i++] = e.getKey();
+                AtomicLong l = e.getValue();
+                long c = l.get();
                 if (lowest == -1 || c < lowest) {
                     lowest = c;
                 }
             }
+            this.subscribers = subscriberArray;
             /*
              * when receiving a request from a subscriber we reset 'outstanding' to the lowest of all subscribers
              */
             outstandingRequests = lowest;
-            emittedSinceRequest = 0;
-            return outstandingRequests;
+            return lowest;
         }
     }
 
@@ -299,7 +301,7 @@ private static class RequestHandler<T> {
         @SuppressWarnings("rawtypes")
         static final AtomicLongFieldUpdater<RequestHandler> WIP = AtomicLongFieldUpdater.newUpdater(RequestHandler.class, "wip");
 
-        public void requestFromChildSubscriber(Subscriber<? super T> subscriber, Long request) {
+        public void requestFromChildSubscriber(Subscriber<? super T> subscriber, long request) {
             state.requestFromSubscriber(subscriber, request);
             OriginSubscriber<T> originSubscriber = state.getOrigin();
             if(originSubscriber != null) {
@@ -333,6 +335,11 @@ private void requestMoreAfterEmission(int emitted) {
 
         public void drainQueue(OriginSubscriber<T> originSubscriber) {
             if (WIP.getAndIncrement(this) == 0) {
+                State<T> localState = state;
+                Map<Subscriber<? super T>, AtomicLong> localMap = localState.ss;
+                RxRingBuffer localBuffer = originSubscriber.buffer;
+                NotificationLite<T> nl = notifier;
+                
                 int emitted = 0;
                 do {
                     /*
@@ -345,26 +352,24 @@ public void drainQueue(OriginSubscriber<T> originSubscriber) {
                      * If we want to batch this then we need to account for new subscribers arriving with a lower request count
                      * concurrently while iterating the batch ... or accept that they won't
                      */
+                    
                     while (true) {
-                        boolean shouldEmit = state.canEmitWithDecrement();
+                        boolean shouldEmit = localState.canEmitWithDecrement();
                         if (!shouldEmit) {
                             break;
                         }
-                        Object o = originSubscriber.buffer.poll();
+                        Object o = localBuffer.poll();
                         if (o == null) {
                             // nothing in buffer so increment outstanding back again
-                            state.incrementOutstandingAfterFailedEmit();
+                            localState.incrementOutstandingAfterFailedEmit();
                             break;
                         }
 
-                        if (notifier.isCompleted(o)) {
-                            for (Subscriber<? super T> s : state.getSubscribers()) {
-                                notifier.accept(s, o);
-                            }
-
-                        } else {
-                            for (Subscriber<? super T> s : state.getSubscribers()) {
-                                notifier.accept(s, o);
+                        for (Subscriber<? super T> s : localState.getSubscribers()) {
+                            AtomicLong req = localMap.get(s);
+                            if (req != null) { // null req indicates a concurrent unsubscription happened
+                                nl.accept(s, o);
+                                req.decrementAndGet();
                             }
                         }
                         emitted++;
diff --git a/src/test/java/rx/internal/operators/OperatorPublishTest.java b/src/test/java/rx/internal/operators/OperatorPublishTest.java
@@ -15,23 +15,17 @@
  */
 package rx.internal.operators;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
 
 import java.util.Arrays;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.*;
 
 import org.junit.Test;
 
-import rx.Observable;
+import rx.*;
 import rx.Observable.OnSubscribe;
-import rx.Subscriber;
-import rx.Subscription;
-import rx.functions.Action0;
-import rx.functions.Action1;
-import rx.functions.Func1;
+import rx.functions.*;
 import rx.internal.util.RxRingBuffer;
 import rx.observables.ConnectableObservable;
 import rx.observers.TestSubscriber;
@@ -187,4 +181,65 @@ public Boolean call(Integer i) {
         System.out.println(ts.getOnNextEvents());
     }
 
+    @Test(timeout = 10000)
+    public void testBackpressureTwoConsumers() {
+        final AtomicInteger sourceEmission = new AtomicInteger();
+        final AtomicBoolean sourceUnsubscribed = new AtomicBoolean();
+        final Observable<Integer> source = Observable.range(1, 100)
+                .doOnNext(new Action1<Integer>() {
+                    @Override
+                    public void call(Integer t1) {
+                        sourceEmission.incrementAndGet();
+                    }
+                })
+                .doOnUnsubscribe(new Action0() {
+                    @Override
+                    public void call() {
+                        sourceUnsubscribed.set(true);
+                    }
+                }).share();
+        ;
+        
+        final AtomicBoolean child1Unsubscribed = new AtomicBoolean();
+        final AtomicBoolean child2Unsubscribed = new AtomicBoolean();
+
+        final TestSubscriber<Integer> ts2 = new TestSubscriber<Integer>();
+
+        final TestSubscriber<Integer> ts1 = new TestSubscriber<Integer>() {
+            @Override
+            public void onNext(Integer t) {
+                if (getOnNextEvents().size() == 2) {
+                    source.doOnUnsubscribe(new Action0() {
+                        @Override
+                        public void call() {
+                            child2Unsubscribed.set(true);
+                        }
+                    }).take(5).subscribe(ts2);
+                }
+                super.onNext(t);
+            }
+        };
+        
+        source.doOnUnsubscribe(new Action0() {
+            @Override
+            public void call() {
+                child1Unsubscribed.set(true);
+            }
+        }).take(5).subscribe(ts1);
+        
+        ts1.awaitTerminalEvent();
+        ts2.awaitTerminalEvent();
+        
+        ts1.assertNoErrors();
+        ts2.assertNoErrors();
+        
+        assertTrue(sourceUnsubscribed.get());
+        assertTrue(child1Unsubscribed.get());
+        assertTrue(child2Unsubscribed.get());
+        
+        ts1.assertReceivedOnNext(Arrays.asList(1, 2, 3, 4, 5));
+        ts2.assertReceivedOnNext(Arrays.asList(4, 5, 6, 7, 8));
+        
+        assertEquals(8, sourceEmission.get());
+    }
 }
