diff --git a/src/main/java/io/reactivex/rxjava4/core/Streamable.java b/src/main/java/io/reactivex/rxjava4/core/Streamable.java index 924342b0fc..706b0b94d8 100644 --- a/src/main/java/io/reactivex/rxjava4/core/Streamable.java +++ b/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 onErrorResumeNext(@NonNull Function + * 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 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 take(long count) { + if (count < 0) { + throw new IllegalArgumentException("count >= 0 required but it was " + count); + } + return RxJavaPlugins.onAssembly(new StreamableTake<>(this, count)); } /** diff --git a/src/main/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableTake.java b/src/main/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableTake.java new file mode 100644 index 0000000000..7f0a20371b --- /dev/null +++ b/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(Streamable source, long count) +implements Streamable, HasUpstreamStreamableSource { + + @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 implements Streamer { + final Streamer upstream; + + final Disposable upstreamDisposable; + + long remaining; + + TakeStreamer(Streamer upstream, long count, Disposable upstreamDisposable) { + this.upstream = upstream; + this.upstreamDisposable = upstreamDisposable; + this.remaining = count; + } + + @Override + public @NonNull CompletionStage next() { + if (remaining-- <= 0L) { + upstreamDisposable.dispose(); + return NEXT_FALSE; + } + return upstream.next(); + } + + @Override + public @NonNull T current() { + return upstream.current(); + } + + @Override + public @NonNull CompletionStage finish() { + return upstream.finish(); + } + } +} diff --git a/src/test/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableTakeTest.java b/src/test/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableTakeTest.java index 1dc3a8a001..f3b6b311ad 100644 --- a/src/test/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableTakeTest.java +++ b/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); + } } diff --git a/src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationNamingTest.java b/src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationNamingTest.java index 4bcdabed08..8a67892d70 100644 --- a/src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationNamingTest.java +++ b/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; } diff --git a/src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationTest.java b/src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationTest.java index 89541434e0..9b72fd1501 100644 --- a/src/test/java/io/reactivex/rxjava4/validators/CheckParamValidationTest.java +++ b/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)); // -----------------------------------------------------------------------------------