4.x: Streamable ~ take native implementation

src/main/java/io/reactivex/rxjava4/core/Streamable.java

@@ -15,14 +15,12 @@
 
 import java.util.*;
 import java.util.concurrent.*;
-import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.*;
 
 import io.reactivex.rxjava4.annotations.*;
 import io.reactivex.rxjava4.core.config.*;
 import io.reactivex.rxjava4.disposables.*;
 import io.reactivex.rxjava4.functions.*;
-import io.reactivex.rxjava4.internal.functions.ObjectHelper;
 import io.reactivex.rxjava4.internal.operators.streamable.*;
 import io.reactivex.rxjava4.plugins.RxJavaPlugins;
 import io.reactivex.rxjava4.schedulers.Schedulers;
@@ -806,23 +804,22 @@ default Streamable<T> onErrorResumeNext(@NonNull Function<? super Throwable, ? e
     /**
      * Takes at most the given number of items from the upstream and relays it to the downstream,
      * then cancels the rest of the sequence.
-     * @param n the maximum number of items to relay
+     * <p>
+     * Note that cancellation of the upstream happens when the downstream
+     * calls {@link Streamer#next()} because unlike the push-based {@code take}
+     * implementations, the current upstream value has to remain accessible until
+     * the downstream calls {@code next} or {@link Streamer#finish()}.
+     * @param count the maximum number of items to relay
      * @return the new {@code Streamable} instance
-     * @throws IllegalArgumentException if {@code n} is non-positive
+     * @throws IllegalArgumentException if {@code count} is negative
      */
     @CheckReturnValue
     @NonNull
-    default Streamable<T> take(long n) {
-        ObjectHelper.verifyPositive(n, "n");
-        return defer(() -> {
-            var countdown = new AtomicLong(n);
-            return transform((item, emitter, stopper) -> {
-                emitter.emit(item);
-                if (countdown.decrementAndGet() <= 0) {
-                    stopper.dispose();
-                }
-            });
-        });
+    default Streamable<T> take(long count) {
+        if (count < 0) {
+            throw new IllegalArgumentException("count >= 0 required but it was " + count);
+        }
+        return RxJavaPlugins.onAssembly(new StreamableTake<>(this, count));
     }
 
     /**

src/main/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableTake.java

@@ -0,0 +1,64 @@
+/*
+ * 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.rxjava4.internal.operators.streamable;
+
+import java.util.concurrent.CompletionStage;
+
+import io.reactivex.rxjava4.annotations.NonNull;
+import io.reactivex.rxjava4.core.*;
+import io.reactivex.rxjava4.disposables.*;
+import io.reactivex.rxjava4.internal.fuseable.HasUpstreamStreamableSource;
+
+public record StreamableTake<T>(Streamable<T> source, long count)
+implements Streamable<T>, HasUpstreamStreamableSource<T> {
+
+    @Override
+    public @NonNull Streamer<@NonNull T> stream(@NonNull StreamerCancellation cancellation) {
+        var dsc = cancellation.derive();
+        return new TakeStreamer<>(source.stream(dsc), count, dsc);
+    }
+
+    static final class TakeStreamer<T> implements Streamer<T> {
+        final Streamer<T> upstream;
+
+        final Disposable upstreamDisposable;
+
+        long remaining;
+
+        TakeStreamer(Streamer<T> upstream, long count, Disposable upstreamDisposable) {
+            this.upstream = upstream;
+            this.upstreamDisposable = upstreamDisposable;
+            this.remaining = count;
+        }
+
+        @Override
+        public @NonNull CompletionStage<Boolean> next() {
+            if (remaining-- <= 0L) {
+                upstreamDisposable.dispose();
+                return NEXT_FALSE;
+            }
+            return upstream.next();
+        }
+
+        @Override
+        public @NonNull T current() {
+            return upstream.current();
+        }
+
+        @Override
+        public @NonNull CompletionStage<Void> finish() {
+            return upstream.finish();
+        }
+    }
+}

src/test/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableTakeTest.java

@@ -19,7 +19,10 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.jupiter.api.Test;
-import io.reactivex.rxjava4.core.Flowable;
+
+import io.reactivex.rxjava4.core.*;
+import io.reactivex.rxjava4.exceptions.TestException;
+import io.reactivex.rxjava4.processors.DispatchStreamProcessor;
 
 public class StreamableTakeTest extends StreamableBaseTest {
 
@@ -52,4 +55,45 @@ public void fewer() throws Throwable {
 
         assertFalse(isCancelled.get(), "Cancel was propagated!");
     }
+
+    @Test
+    public void error() {
+        Streamable.error(new TestException())
+        .take(5)
+        .test()
+        .awaitDone(5, TimeUnit.SECONDS)
+        .assertFailure(TestException.class);
+    }
+
+    @Test
+    public void doubleTake() {
+        Streamable.range(1, 5)
+        .take(3)
+        .take(1)
+        .test()
+        .awaitDone(5, TimeUnit.SECONDS)
+        .assertResult(1)
+        ;
+    }
+
+    @Test
+    public void cancelled() throws Throwable {
+        var dsp = new DispatchStreamProcessor<>();
+
+        var ts = dsp.take(3).test();
+
+        ts.awaitOnSubscribe(1, TimeUnit.SECONDS);
+        awaitStreamers(dsp, 1000);
+
+        dsp.next(1).toCompletableFuture().join();
+        dsp.next(2).toCompletableFuture().join();
+        dsp.next(3).toCompletableFuture().join();
+        dsp.next(4).toCompletableFuture().join();
+
+        awaitNoStreamers(dsp, 1000);
+
+        ts
+        .awaitDone(5, TimeUnit.SECONDS)
+        .assertResult(1, 2, 3);
+    }
 }

src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationNamingTest.java

@@ -318,6 +318,7 @@ static void processFile(Class<?> clazz) throws Exception {
                         if (linek.startsWith("public") || linek.startsWith("private")
                                 || linek.startsWith("protected")
                                 || linek.startsWith("static")
+                                || linek.startsWith("default")
                                 || linek.startsWith(baseClassName)) {
                             break;
                         }

src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationTest.java

@@ -634,6 +634,8 @@ public void checkStreamable() {
         addIgnore(new ParamIgnore(Streamable.class, "rangeLong", Long.TYPE, Long.TYPE));
         addIgnore(new ParamIgnore(Streamable.class, "intervalRange", Long.TYPE, Long.TYPE, Long.TYPE, Long.TYPE, TimeUnit.class, Scheduler.class));
         addIgnore(new ParamIgnore(Streamable.class, "intervalRange", Long.TYPE, Long.TYPE, Long.TYPE, Long.TYPE, TimeUnit.class, ExecutorService.class));
+        // zero take is allowed
+        addOverride(new ParamOverride(Streamable.class, 0, ParamMode.NON_NEGATIVE, "take", Long.TYPE));
 
         // -----------------------------------------------------------------------------------