3.x: constrain upstream requests from take (#6569)

Author: davidmoten

src/main/java/io/reactivex/rxjava3/core/Flowable.java

@@ -11212,52 +11212,6 @@ public final <R> Flowable<R> lift(FlowableOperator<? extends R, ? super T> lifte
         return RxJavaPlugins.onAssembly(new FlowableLift<R, T>(this, lifter));
     }
 
-    /**
-     * Limits both the number of upstream items (after which the sequence completes)
-     * and the total downstream request amount requested from the upstream to
-     * possibly prevent the creation of excess items by the upstream.
-     * <p>
-     * The operator requests at most the given {@code count} of items from upstream even
-     * if the downstream requests more than that. For example, given a {@code limit(5)},
-     * if the downstream requests 1, a request of 1 is submitted to the upstream
-     * and the operator remembers that only 4 items can be requested now on. A request
-     * of 5 at this point will request 4 from the upstream and any subsequent requests will
-     * be ignored.
-     * <p>
-     * Note that requests are negotiated on an operator boundary and {@code limit}'s amount
-     * may not be preserved further upstream. For example,
-     * {@code source.observeOn(Schedulers.computation()).limit(5)} will still request the
-     * default (128) elements from the given {@code source}.
-     * <p>
-     * The main use of this operator is with sources that are async boundaries that
-     * don't interfere with request amounts, such as certain {@code Flowable}-based
-     * network endpoints that relay downstream request amounts unchanged and are, therefore,
-     * prone to trigger excessive item creation/transmission over the network.
-     * <dl>
-     *  <dt><b>Backpressure:</b></dt>
-     *  <dd>The operator requests a total of the given {@code count} items from the upstream.</dd>
-     *  <dt><b>Scheduler:</b></dt>
-     *  <dd>{@code limit} does not operate by default on a particular {@link Scheduler}.</dd>
-     * </dl>
-     * <p>History: 2.1.6 - experimental
-     * @param count the maximum number of items and the total request amount, non-negative.
-     *              Zero will immediately cancel the upstream on subscription and complete
-     *              the downstream.
-     * @return the new Flowable instance
-     * @see #take(long)
-     * @see #rebatchRequests(int)
-     * @since 2.2
-     */
-    @BackpressureSupport(BackpressureKind.SPECIAL)
-    @SchedulerSupport(SchedulerSupport.NONE)
-    @CheckReturnValue
-    public final Flowable<T> limit(long count) {
-        if (count < 0) {
-            throw new IllegalArgumentException("count >= 0 required but it was " + count);
-        }
-        return RxJavaPlugins.onAssembly(new FlowableLimit<T>(this, count));
-    }
-
     /**
      * Returns a Flowable that applies a specified function to each item emitted by the source Publisher and
      * emits the results of these function applications.
@@ -15383,9 +15337,7 @@ public final <R> Flowable<R> switchMapSingleDelayError(@NonNull Function<? super
      * {@link Subscriber#onComplete onComplete}.
      * <dl>
      *  <dt><b>Backpressure:</b></dt>
-     *  <dd>The operator doesn't interfere with backpressure which is determined by the source {@code Publisher}'s backpressure
-     *  behavior in case the first request is smaller than the {@code count}. Otherwise, the source {@code Publisher}
-     *  is consumed in an unbounded manner (i.e., without applying backpressure to it).</dd>
+     *  <dd>The source {@code Publisher} is consumed in a bounded manner.</dd>
      *  <dt><b>Scheduler:</b></dt>
      *  <dd>This version of {@code take} does not operate by default on a particular {@link Scheduler}.</dd>
      * </dl>
@@ -15397,7 +15349,7 @@ public final <R> Flowable<R> switchMapSingleDelayError(@NonNull Function<? super
      * @see <a href="http://reactivex.io/documentation/operators/take.html">ReactiveX operators documentation: Take</a>
      */
     @CheckReturnValue
-    @BackpressureSupport(BackpressureKind.SPECIAL) // may trigger UNBOUNDED_IN
+    @BackpressureSupport(BackpressureKind.FULL) 
     @SchedulerSupport(SchedulerSupport.NONE)
     public final Flowable<T> take(long count) {
         if (count < 0) {

src/main/java/io/reactivex/rxjava3/internal/operators/flowable/FlowableLimit.java

@@ -13,7 +13,7 @@
 
 package io.reactivex.rxjava3.internal.operators.flowable;
 
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.reactivestreams.*;
 
@@ -33,7 +33,7 @@ protected void subscribeActual(Subscriber<? super T> s) {
         source.subscribe(new TakeSubscriber<T>(s, limit));
     }
 
-    static final class TakeSubscriber<T> extends AtomicBoolean implements FlowableSubscriber<T>, Subscription {
+    static final class TakeSubscriber<T> extends AtomicLong implements FlowableSubscriber<T>, Subscription {
 
         private static final long serialVersionUID = -5636543848937116287L;
 
@@ -69,7 +69,7 @@ public void onSubscribe(Subscription s) {
 
         @Override
         public void onNext(T t) {
-            if (!done && remaining-- > 0) {
+            if (remaining-- > 0) {
                 boolean stop = remaining == 0;
                 downstream.onNext(t);
                 if (stop) {
@@ -103,13 +103,29 @@ public void request(long n) {
             if (!SubscriptionHelper.validate(n)) {
                 return;
             }
-            if (!get() && compareAndSet(false, true)) {
-                if (n >= limit) {
-                    upstream.request(Long.MAX_VALUE);
+            while (true) {
+                // get the total requests so far r
+                long r = get();
+                if (r == limit) {
                     return;
                 }
+                // calculate the new total of requests r2 after overflow 
+                // and limit have been accounted for
+                long r2 = r + n;
+                final long more;
+                if (r2 < 0) {
+                    // overflow
+                    r2 = limit;
+                    more = limit - r;
+                } else {
+                    r2 = Math.min(limit, r2);
+                    more = r2 - r;
+                }
+                if (compareAndSet(r, r2)) {
+                    upstream.request(more);
+                    break;
+                }
             }
-            upstream.request(n);
         }
 
         @Override

src/main/java/io/reactivex/rxjava3/internal/operators/flowable/FlowableTake.java

@@ -331,8 +331,7 @@ public void cancel() {
             }
         }).take(2).subscribe(ts);
 
-        // FIXME the take now requests Long.MAX_PATH if downstream requests at least the limit
-        assertEquals(Long.MAX_VALUE, requested.get());
+        assertEquals(2, requested.get());
     }
 
     @Test

src/test/java/io/reactivex/rxjava3/flowable/FlowableSubscriberTest.java

@@ -2047,30 +2047,6 @@ public void openCloseTake() {
         ts.assertResult(Collections.<Integer>emptyList());
     }
 
-    @Test
-    @SuppressWarnings("unchecked")
-    public void openCloseLimit() {
-        PublishProcessor<Integer> source = PublishProcessor.create();
-
-        PublishProcessor<Integer> openIndicator = PublishProcessor.create();
-
-        PublishProcessor<Integer> closeIndicator = PublishProcessor.create();
-
-        TestSubscriber<List<Integer>> ts = source
-        .buffer(openIndicator, Functions.justFunction(closeIndicator))
-        .limit(1)
-        .test(2);
-
-        openIndicator.onNext(1);
-        closeIndicator.onComplete();
-
-        assertFalse(source.hasSubscribers());
-        assertFalse(openIndicator.hasSubscribers());
-        assertFalse(closeIndicator.hasSubscribers());
-
-        ts.assertResult(Collections.<Integer>emptyList());
-    }
-
     @Test
     @SuppressWarnings("unchecked")
     public void openCloseEmptyBackpressure() {

src/test/java/io/reactivex/rxjava3/internal/operators/flowable/FlowableBufferTest.java

@@ -357,7 +357,7 @@ public void drainExactRequestCancel() {
         final MaybeSubject<Integer> cs = MaybeSubject.create();
 
         TestSubscriber<Integer> ts = pp.mergeWith(cs)
-                .limit(2)
+                .take(2)
                 .subscribeWith(new TestSubscriber<Integer>(2) {
             @Override
             public void onNext(Integer t) {

src/test/java/io/reactivex/rxjava3/internal/operators/flowable/FlowableMergeWithMaybeTest.java

@@ -353,7 +353,7 @@ public void drainExactRequestCancel() {
         final SingleSubject<Integer> cs = SingleSubject.create();
 
         TestSubscriber<Integer> ts = pp.mergeWith(cs)
-                .limit(2)
+                .take(2)
                 .subscribeWith(new TestSubscriber<Integer>(2) {
             @Override
             public void onNext(Integer t) {

src/test/java/io/reactivex/rxjava3/internal/operators/flowable/FlowableMergeWithSingleTest.java

@@ -142,6 +142,27 @@ public void subscribe(Subscriber<? super String> subscriber) {
             RxJavaPlugins.reset();
         }
     }
+    
+    @Test
+    public void takeEmitsErrors() {
+        Flowable.error(new TestException())
+            .take(1)
+            .test()
+            .assertNoValues()
+            .assertError(TestException.class);
+    }
+    
+    @Test
+    public void takeRequestOverflow() {
+        TestSubscriber<Integer> ts = Flowable.just(1, 2, 3)
+            .take(3)
+            .test(0);
+        ts.requestMore(1)
+          .assertValues(1)
+          .assertNotComplete()
+          .requestMore(Long.MAX_VALUE)
+          .assertValues(1, 2, 3);
+    }
 
     @Test
     public void unsubscribeAfterTake() {
@@ -305,7 +326,7 @@ public void cancel() {
             }
 
         }).take(3).subscribe(ts);
-        assertEquals(Long.MAX_VALUE, requested.get());
+        assertEquals(3, requested.get());
     }
 
     @Test
@@ -332,7 +353,7 @@ public void cancel() {
 
         }).take(1).subscribe(ts);
         //FIXME take triggers fast path if downstream requests more than the limit
-        assertEquals(Long.MAX_VALUE, requested.get());
+        assertEquals(1, requested.get());
     }
 
     @Test
@@ -383,7 +404,7 @@ public void accept(long n) {
         ts.awaitDone(5, TimeUnit.SECONDS);
         ts.assertComplete();
         ts.assertNoErrors();
-        assertEquals(3, requests.get());
+        assertEquals(2, requests.get());
     }
 
     @Test