prevent OperatorTake from requesting more than needed

Author: davidmoten

src/main/java/rx/internal/operators/OperatorTake.java

@@ -15,6 +15,8 @@
  */
 package rx.internal.operators;
 
+import java.util.concurrent.atomic.AtomicLong;
+
 import rx.Observable.Operator;
 import rx.Producer;
 import rx.Subscriber;
@@ -78,15 +80,24 @@ public void onNext(T i) {
             @Override
             public void setProducer(final Producer producer) {
                 child.setProducer(new Producer() {
-
+                    
+                    // keeps track of requests up to maximum of `limit`
+                    final AtomicLong requested = new AtomicLong(0);
+                    
                     @Override
                     public void request(long n) {
-                        if (!completed) {
-                            long c = limit - count;
-                            if (n < c) {
-                                producer.request(n);
-                            } else {
-                                producer.request(c);
+                        if (n >0 && !completed) {
+                            // because requests may happen concurrently use a CAS loop to 
+                            // ensure we only request as much as needed, no more no less
+                            while (true) {
+                                long r = requested.get();
+                                long c = Math.min(n, limit - r);
+                                if (c == 0)
+                                    break;
+                                else if (requested.compareAndSet(r, r + c)) {
+                                    producer.request(c);
+                                    break;
+                                }
                             }
                         }
                     }

src/test/java/rx/internal/operators/OperatorTakeTest.java

@@ -27,6 +27,7 @@
 
 import java.util.Arrays;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
@@ -388,4 +389,29 @@ public void call(Integer t1) {
         latch.await();
         assertNull(exception.get());
     }
+    
+    @Test
+    public void testDoesntRequestMoreThanNeededFromUpstream() throws InterruptedException {
+        final AtomicLong requests = new AtomicLong();
+        TestSubscriber<Long> ts = new TestSubscriber<Long>(0);
+        Observable.interval(100, TimeUnit.MILLISECONDS)
+            //
+            .doOnRequest(new Action1<Long>() {
+                @Override
+                public void call(Long n) {
+                    requests.addAndGet(n);
+            }})
+            //
+            .take(2)
+            //
+            .subscribe(ts);
+        Thread.sleep(50);
+        ts.requestMore(1);
+        ts.requestMore(1);
+        ts.requestMore(1);
+        ts.awaitTerminalEvent();
+        ts.assertCompleted();
+        ts.assertNoErrors();
+        assertEquals(2,requests.get());
+    }
 }