2.x: add limit() to limit both item count and request amount (#5655)

* 2.x: add limit() to limit both item count and request amount

* Address some feedback.

* Use the cancelled constant value.

Author: akarnokd

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

@@ -9676,6 +9676,53 @@ 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>
+
+     * @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.1.6 - experimental
+     */
+    @Experimental
+    @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.

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

@@ -0,0 +1,137 @@
+/**
+ * Copyright (c) 2016-present, RxJava Contributors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in
+ * compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is
+ * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
+ * the License for the specific language governing permissions and limitations under the License.
+ */
+
+package io.reactivex.internal.operators.flowable;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.reactivestreams.*;
+
+import io.reactivex.*;
+import io.reactivex.annotations.Experimental;
+import io.reactivex.internal.subscriptions.*;
+import io.reactivex.plugins.RxJavaPlugins;
+
+/**
+ * Limits both the total request amount and items received from the upstream.
+ *
+ * @param <T> the source and output value type
+ * @since 2.1.6 - experimental
+ */
+@Experimental
+public final class FlowableLimit<T> extends AbstractFlowableWithUpstream<T, T> {
+
+    final long n;
+
+    public FlowableLimit(Flowable<T> source, long n) {
+        super(source);
+        this.n = n;
+    }
+
+    @Override
+    protected void subscribeActual(Subscriber<? super T> s) {
+        source.subscribe(new LimitSubscriber<T>(s, n));
+    }
+
+    static final class LimitSubscriber<T>
+    extends AtomicLong
+    implements FlowableSubscriber<T>, Subscription {
+
+        private static final long serialVersionUID = 2288246011222124525L;
+
+        final Subscriber<? super T> actual;
+
+        long remaining;
+
+        Subscription upstream;
+
+        LimitSubscriber(Subscriber<? super T> actual, long remaining) {
+            this.actual = actual;
+            this.remaining = remaining;
+            lazySet(remaining);
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+            if (SubscriptionHelper.validate(this.upstream, s)) {
+                if (remaining == 0L) {
+                    s.cancel();
+                    EmptySubscription.complete(actual);
+                } else {
+                    this.upstream = s;
+                    actual.onSubscribe(this);
+                }
+            }
+        }
+
+        @Override
+        public void onNext(T t) {
+            long r = remaining;
+            if (r > 0L) {
+                remaining = --r;
+                actual.onNext(t);
+                if (r == 0L) {
+                    upstream.cancel();
+                    actual.onComplete();
+                }
+            }
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            if (remaining > 0L) {
+                remaining = 0L;
+                actual.onError(t);
+            } else {
+                RxJavaPlugins.onError(t);
+            }
+        }
+
+        @Override
+        public void onComplete() {
+            if (remaining > 0L) {
+                remaining = 0L;
+                actual.onComplete();
+            }
+        }
+
+        @Override
+        public void request(long n) {
+            if (SubscriptionHelper.validate(n)) {
+                for (;;) {
+                    long r = get();
+                    if (r == 0L) {
+                        break;
+                    }
+                    long toRequest;
+                    if (r <= n) {
+                        toRequest = r;
+                    } else {
+                        toRequest = n;
+                    }
+                    long u = r - toRequest;
+                    if (compareAndSet(r, u)) {
+                        upstream.request(toRequest);
+                        break;
+                    }
+                }
+            }
+        }
+
+        @Override
+        public void cancel() {
+            upstream.cancel();
+        }
+
+    }
+}

src/test/java/io/reactivex/ParamValidationCheckerTest.java

@@ -182,8 +182,9 @@ public void checkParallelFlowable() {
         addOverride(new ParamOverride(Flowable.class, 0, ParamMode.ANY, "take", Long.TYPE, TimeUnit.class));
         addOverride(new ParamOverride(Flowable.class, 0, ParamMode.ANY, "take", Long.TYPE, TimeUnit.class, Scheduler.class));
 
-        // zero retry is allowed
+        // zero take/limit is allowed
         addOverride(new ParamOverride(Flowable.class, 0, ParamMode.NON_NEGATIVE, "take", Long.TYPE));
+        addOverride(new ParamOverride(Flowable.class, 0, ParamMode.NON_NEGATIVE, "limit", Long.TYPE));
 
         // negative time is considered as zero time
         addOverride(new ParamOverride(Flowable.class, 0, ParamMode.ANY, "sample", Long.TYPE, TimeUnit.class));