From dbb318c5fccf414b2df61db56edc184c26c3d07e Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Fri, 2 Oct 2026 15:44:42 -0400 Subject: [PATCH 01/11] Add a performant groupChunksWithin --- .../GroupChunksWithinBenchmark.scala | 64 ++++ .../fs2/benchmark/GroupWithinBenchmark.scala | 12 + core/shared/src/main/scala/fs2/Chunk.scala | 18 ++ core/shared/src/main/scala/fs2/Stream.scala | 78 +++++ .../scala/fs2/concurrent/ConditionedRef.scala | 110 +++++++ .../src/test/scala/fs2/ChunkSuite.scala | 24 ++ .../scala/fs2/StreamCombinatorsSuite.scala | 281 ++++++++++++++++++ .../fs2/concurrent/ConditionedRefSuite.scala | 88 ++++++ .../src/test/scala/fs2/MemoryLeakSpec.scala | 21 ++ 9 files changed, 696 insertions(+) create mode 100644 benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala create mode 100644 core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala create mode 100644 core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala diff --git a/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala b/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala new file mode 100644 index 0000000000..9a008f5f47 --- /dev/null +++ b/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2013 Functional Streams for Scala + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER + * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN + * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +package fs2 +package benchmark + +import cats.effect.IO +import org.openjdk.jmh.annotations.{Benchmark, Param, Scope, State} + +import scala.concurrent.duration.* + +@State(Scope.Thread) +class GroupChunksWithinBenchmark { + + import cats.effect.unsafe.implicits.global + + val bufferWindow = 100.micros + + @Param(Array("100", "10000", "100000")) + var rangeLength: Int = _ + + @Param(Array("16", "256", "4096")) + var bufferSize: Int = _ + + @Benchmark + def groupChunksWithin(): Unit = + Stream + .range(0, rangeLength) + .covary[IO] + .groupChunksWithin(bufferSize, bufferWindow) + .compile + .drain + .unsafeRunSync() + + @Benchmark + def groupChunksWithinChunkedUpstream(): Unit = + Stream + .range(0, rangeLength) + .chunkN(bufferSize / 4 + 1) + .unchunks + .covary[IO] + .groupChunksWithin(bufferSize, bufferWindow) + .compile + .drain + .unsafeRunSync() +} diff --git a/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala b/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala index 9ffb7eb912..5e989539a3 100644 --- a/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala +++ b/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala @@ -59,4 +59,16 @@ class GroupWithinBenchmark { .compile .drain .unsafeRunSync() + + @Benchmark + def groupWithinChunkedUpstream(): Unit = + Stream + .range(0, rangeLength) + .chunkN(bufferSize / 4 + 1) + .unchunks + .covary[IO] + .groupWithin(bufferSize, bufferWindow) + .compile + .drain + .unsafeRunSync() } diff --git a/core/shared/src/main/scala/fs2/Chunk.scala b/core/shared/src/main/scala/fs2/Chunk.scala index 03aac9d5aa..c99d56f47d 100644 --- a/core/shared/src/main/scala/fs2/Chunk.scala +++ b/core/shared/src/main/scala/fs2/Chunk.scala @@ -207,6 +207,24 @@ abstract class Chunk[+O] extends Serializable with ChunkPlatform[O] with ChunkRu } } + /** Splits this chunk into groups of `n` elements, the last of which may be smaller. + * + * Like `List#grouped`, an empty chunk has no groups, and `n` must be positive. + */ + def grouped(n: Int): Vector[Chunk[O]] = { + require(n > 0, s"n must be positive, but got ${n.toString}") + + @tailrec + def loop(rest: Chunk[O], acc: Vector[Chunk[O]]): Vector[Chunk[O]] = + if (rest.size <= n) acc :+ rest + else { + val (group, remaining) = rest.splitAt(n) + loop(rest = remaining, acc = acc :+ group) + } + + if (isEmpty) Vector.empty else loop(rest = this, acc = Vector.empty) + } + /** Gets the first element of this chunk. */ def head: Option[O] = if (isEmpty) None else Some(apply(0)) diff --git a/core/shared/src/main/scala/fs2/Stream.scala b/core/shared/src/main/scala/fs2/Stream.scala index a3642794a8..2c7a31a10f 100644 --- a/core/shared/src/main/scala/fs2/Stream.scala +++ b/core/shared/src/main/scala/fs2/Stream.scala @@ -1616,6 +1616,84 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, } } + /** Like [[groupWithin]], but works on chunks for better performance. + */ + def groupChunksWithin[F2[x] >: F[x]]( + chunkSize: Int, + timeout: FiniteDuration + )(implicit F: Temporal[F2]): Stream[F2, Chunk[O]] = + Stream.force { + require(chunkSize > 0, s"chunkSize must be > 0, but got ${chunkSize.toString}") + + val StopConsumer = none[Chunk[O]].pure[F2] + val Skip = Chunk.empty[O].some.pure[F2] + + final case class Buffer[A](chunk: Chunk[A], done: Option[ExitCase]) { + def size: Int = chunk.size + def isEmpty: Boolean = chunk.isEmpty + def nonEmpty: Boolean = chunk.nonEmpty + def isFull: Boolean = size >= chunkSize + def isDone: Boolean = done.isDefined + } + + ConditionedRef.of[F2, Buffer[O]](Buffer(chunk = Chunk.empty[O], done = none)).map { buffer => + val producer = chunks + .evalMap { chunk => + buffer + .updateAndGet(b => b.copy(chunk = b.chunk ++ chunk)) + .flatMap(b => F.whenA(b.isFull)(buffer.waitUntil(!_.isFull))) + } + .onFinalizeCase(exitCase => buffer.update(_.copy(done = exitCase.some))) + .compile + .drain + + def take(b: Buffer[O], n: Int): (Buffer[O], F2[Option[Chunk[O]]]) = { + val (taken, remaining) = b.chunk.splitAt(n) + b.copy(chunk = remaining) -> taken.some.pure[F2] + } + + // None means upstream is done, and the buffer is drained. + def takeOrExit(all: Boolean): F2[Option[Chunk[O]]] = + buffer.modify { b => + if (b.nonEmpty && (all || b.isDone)) take(b, n = b.size) + // Take all potential batches so that producer doesn't need to wait for timeout + else if (b.isFull) take(b, n = b.size - b.size % chunkSize) + else + b -> (b.done match { + case None /* not full & no-timeout */ => Skip + case Some(ExitCase.Errored(e)) => F.raiseError[Option[Chunk[O]]](e) + case Some(_) /* empty */ => StopConsumer + }) + }.flatten + + val onTimeout = takeOrExit(all = true).flatMap { + case Some(batch) if batch.isEmpty => + buffer.waitUntil(b => b.nonEmpty || b.isDone) >> takeOrExit(all = true) + case result => result.pure[F2] + } + + val nextBatch: F2[Option[Chunk[O]]] = + // Potentially skip starting timer fiber if buffer is full + takeOrExit(all = false).flatMap { + case Some(batch) if batch.isEmpty => + F.race(F.sleep(timeout), buffer.waitUntil(b => b.isFull || b.isDone)) + .flatMap { + case Left(_ /* timeout */ ) => onTimeout + case Right(_ /* full batch or done */ ) => takeOrExit(all = false) + } + case result => result.pure[F2] + } + + def emitBatches: Pull[F2, Chunk[O], Unit] = + Pull.eval(nextBatch).flatMap { + case Some(batch) => Pull.output(Chunk.from(batch.grouped(chunkSize))) >> emitBatches + case None => Pull.done + } + + Stream.bracket(producer.start)(_.cancel) >> emitBatches.stream + } + } + /** If `this` terminates with `Stream.raiseError(e)`, invoke `h(e)`. * * @example {{{ diff --git a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala new file mode 100644 index 0000000000..e558ef01db --- /dev/null +++ b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala @@ -0,0 +1,110 @@ +/* + * Copyright (c) 2013 Functional Streams for Scala + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER + * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN + * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +package fs2 +package concurrent + +import cats.effect._ +import cats.effect.implicits._ +import cats.syntax.all._ + +/** A `Ref` whose value can be waited on. + * + * Waiters register a predicate on the value. Every `modify` evaluates the registered predicates + * against the new value and wakes exactly the waiters whose predicate is satisfied. Unlike [[SignallingRef]], + * waiters are not woken by unrelated updates. + */ +private[fs2] sealed trait ConditionedRef[F[_], A] { + + def get: F[A] + + /** Atomically updates the value and wakes every waiter whose predicate condition holds for the new value. */ + def modify[B](f: A => (A, B)): F[B] + + def update(f: A => A): F[Unit] + + def updateAndGet(f: A => A): F[A] = + modify { a => + val newA = f(a) + (newA, newA) + } + + def set(a: A): F[Unit] = update(_ => a) + + /** Completes if the predicate, `p` holds for the current value, or a later `modify` sets a value that satisfies `p`. + * + * `p` may no longer hold by the time this completes: act on the value through `modify`. + */ + def waitUntil(p: A => Boolean): F[Unit] +} + +private[fs2] object ConditionedRef { + + def of[F[_], A](initial: A)(implicit F: Concurrent[F]): F[ConditionedRef[F, A]] = + F.ref(State[F, A](value = initial, waiters = Nil)).map(new Impl(_)) + + private final class Waiter[F[_], A](val accepts: A => Boolean, val wake: Deferred[F, Unit]) { + def wakeUp: F[Boolean] = wake.complete(()) + } + + private final case class State[F[_], A](value: A, waiters: List[Waiter[F, A]]) { + def register(waiter: Waiter[F, A]): State[F, A] = copy(waiters = waiter :: waiters) + def deregister(waiter: Waiter[F, A]): State[F, A] = + copy(waiters = waiters.filterNot(_ eq waiter)) + } + + private final class Impl[F[_], A](state: Ref[F, State[F, A]])(implicit F: Concurrent[F]) + extends ConditionedRef[F, A] { + + def get: F[A] = state.get.map(_.value) + + def modify[B](f: A => (A, B)): F[B] = + state.flatModify { s => // uncancellable to avoid losing wake-up signals + val (value, result) = f(s.value) + if (!s.waiters.exists(_.accepts(value))) + State(value = value, waiters = s.waiters) -> result.pure[F] + else { + val (toWake, waiting) = s.waiters.partition(_.accepts(value)) + State(value = value, waiters = waiting) -> toWake.traverse_(_.wakeUp).as(result) + } + } + + def update(f: A => A): F[Unit] = modify(a => (f(a), ())) + + def waitUntil(p: A => Boolean): F[Unit] = + get.flatMap { value => + if (p(value)) F.unit + else + F.deferred[Unit].flatMap { wake => + val waiter = new Waiter(accepts = p, wake = wake) + F.uncancelable { poll => + state.modify { s => + if (p(s.value)) s -> F.unit + else + s.register(waiter) -> poll(wake.get).onCancel { + state.update(_.deregister(waiter)) + } + }.flatten + } + } + } + } +} diff --git a/core/shared/src/test/scala/fs2/ChunkSuite.scala b/core/shared/src/test/scala/fs2/ChunkSuite.scala index 33b7e9967e..78ef3b918b 100644 --- a/core/shared/src/test/scala/fs2/ChunkSuite.scala +++ b/core/shared/src/test/scala/fs2/ChunkSuite.scala @@ -131,6 +131,25 @@ class ChunkSuite extends Fs2Suite { } } + test("Chunk.grouped of an empty chunk has no groups, like List#grouped") { + assertEquals(Chunk.empty[Int].grouped(3), Vector.empty) + assertEquals(List.empty[Int].grouped(3).toList, Nil) + } + + test("Chunk.grouped rejects a non-positive group size, like List#grouped") { + List(0, -1).foreach { n => + intercept[IllegalArgumentException](Chunk(1, 2, 3).grouped(n)) + intercept[IllegalArgumentException](Chunk.empty[Int].grouped(n)) + intercept[IllegalArgumentException](List(1, 2, 3).grouped(n)) + } + } + + test("Chunk.grouped returns a chunk that fits in one group as is") { + forAll { (c: Chunk[Int]) => + if (c.nonEmpty) assert(c.grouped(c.size).head eq c) + } + } + class OddStringExtractor { val callCounter: AtomicInteger = new AtomicInteger(0) @@ -167,6 +186,11 @@ class ChunkSuite extends Fs2Suite { property("isEmpty") { forAll((c: Chunk[A]) => assertEquals(c.isEmpty, c.toList.isEmpty)) } + property("grouped") { + forAll(genChunk, Gen.choose(1, 50)) { (c: Chunk[A], n: Int) => + assertEquals(c.grouped(n).map(_.toList).toList, c.toList.grouped(n).toList) + } + } property("toArray") { forAll { (c: Chunk[A]) => assertEquals(c.toArray.toVector, c.toVector) diff --git a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala index 48ed6e1011..b0ea3ade29 100644 --- a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala +++ b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala @@ -987,6 +987,287 @@ class StreamCombinatorsSuite extends Fs2Suite { } } + group("groupChunksWithin") { + implicit val groupSizeArb: Arbitrary[Int] = Arbitrary(Gen.choose(1, 20)) + // a zero timeout is covered separately. Under TestControl it makes the idle loop spin forever + implicit val timeoutArb: Arbitrary[FiniteDuration] = Arbitrary(Gen.choose(1, 50).map(_.millis)) + + def sleep(d: Int): IO[Unit] = IO.sleep((d % 500).abs.micros) + + test("should never lose any elements") { + forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration) => + TestControl.executeEmbed { + s.covary[IO] + .evalTap(sleep) + .groupChunksWithin(chunkSize = 3, timeout) + .unchunks + .assertEmitsSameAs(s) + } + } + } + + test("should never emit empty chunks") { + forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration, groupSize: Int) => + TestControl.executeEmbed { + s.covary[IO] + .evalTap(sleep) + .groupChunksWithin(groupSize, timeout) + .assertForall(_.nonEmpty) + } + } + } + + test("makes progress with a zero timeout") { + Stream + .range(0, 100) + .covary[IO] + .evalTap(_ => IO.sleep(1.milli)) + .groupChunksWithin(chunkSize = 10, timeout = Duration.Zero) + .unchunks + .compile + .toList + .timeout(30.seconds) + .assertEquals((0 until 100).toList) + } + + test("every chunk but the last has at least chunkSize elements when no timeout triggers") { + val source = Stream.range(0, 101) + val chunkSize = 5 + + TestControl.executeEmbed { + source + .covary[IO] + .groupChunksWithin(chunkSize, 1.day) + .compile + .toList + .map { chunks => + assertEquals(chunks.flatMap(_.toList), source.toList) + assert(chunks.init.forall(_.size >= chunkSize)) + assert(chunks.lastOption.forall(_.size < chunkSize)) + } + } + } + + test("splits upstream chunks larger than chunkSize") { + TestControl.executeEmbed { + Stream + .emits(1 to 12) + .covary[IO] + .groupChunksWithin(chunkSize = 5, timeout = 1.day) + .map(_.toList) + .assertEmits(List((1 to 5).toList, (6 to 10).toList, (11 to 12).toList)) + } + } + + test("emits every full group of a large upstream chunk without waiting for the timeout") { + TestControl.executeEmbed { + (Stream.emits(1 to 100) ++ Stream.never[IO]) + .groupChunksWithin(chunkSize = 10, timeout = 1.day) + .take(10) + .compile + .toList + .timed + .map { case (elapsed, chunks) => + assertEquals(elapsed, Duration.Zero) + assertEquals(chunks.map(_.toList), (1 to 100).toList.grouped(10).toList) + } + } + } + + test( + "should return a finite stream back in a single chunk given a chunk size equal to the stream size and an absurdly high duration" + ) { + forAllF { (streamAsList0: List[Int]) => + val streamAsList = 0 :: streamAsList0 + TestControl.executeEmbed { + Stream + .emits(streamAsList) + .covary[IO] + .groupChunksWithin(streamAsList.size, (Int.MaxValue - 1L).nanoseconds) + .compile + .toList + .map(_.head.toList) + .assertEquals(streamAsList) + } + } + } + + test("accumulation with splitting") { + val t = 200.millis + val chunkSize = 5 + + def chunk(from: Int, to: Int) = + Stream.range(from, to + 1).chunkAll.unchunks + + val source = + chunk(from = 1, to = 3) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 4, to = 15) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 16, to = 22) + + val expected = List( + (1 to 3).toList, + (4 to 8).toList, + (9 to 13).toList, + (14 to 15).toList, + (16 to 20).toList, + (21 to 22).toList + ) + + TestControl.executeEmbed { + source + .groupChunksWithin(chunkSize, t) + .map(_.toList) + .assertEmits(expected) + } + } + + test("emits a burst arriving after several idle timeouts") { + val timeout = 100.millis + val chunkSize = 5 + + val source = + Stream.sleep_[IO](timeout * 10) ++ + Stream.emits(1 to 3) ++ + Stream.sleep_[IO](timeout * 10) ++ + Stream.emits(4 to 6) + + TestControl.executeEmbed { + source + .groupChunksWithin(chunkSize, timeout) + .map(_.toList) + .assertEmits(List(List(1, 2, 3), List(4, 5, 6))) + } + } + + test("does not reset timeout if nothing is emitted") { + TestControl + .executeEmbed( + Ref[IO] + .of(0.millis) + .flatMap { ref => + val timeout = 5.seconds + + def measureEmission[A]: Pipe[IO, A, A] = + _.chunks + .evalTap(_ => IO.monotonic.flatMap(ref.set)) + .unchunks + + // emits elements after the timeout has expired + val source = + Stream.sleep_[IO](timeout + 200.millis) ++ + Stream(4, 5) ++ + Stream.never[IO] // avoids emission due to source termination + + source + .through(measureEmission) + .groupChunksWithin(5, timeout) + .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) + .interruptAfter(timeout * 3) + .compile + .lastOrError + } + ) + .assertEquals(0.millis) // The stream emits after the timeout has expired + // on an empty buffer, so groupChunksWithin should re-emit with zero delay + // rather than start a fresh timeout window + } + + test("Edge case: should not introduce unnecessary delays when chunkSize == chunk size") { + TestControl + .executeEmbed( + Ref[IO] + .of(0.millis) + .flatMap { ref => + val timeout = 5.seconds + + def measureEmission[A]: Pipe[IO, A, A] = + _.chunks + .evalTap(_ => IO.monotonic.flatMap(ref.set)) + .unchunks + + val source = + Stream(1, 2, 3) ++ + Stream.sleep_[IO](timeout + 200.millis) + + source + .through(measureEmission) + .groupChunksWithin(3, timeout) + .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) + .compile + .lastOrError + } + ) + .assertEquals(0.millis) + } + + test("upstream failures are propagated downstream") { + TestControl.executeEmbed { + case object SevenNotAllowed extends NoStackTrace + + val source = Stream + .iterate(0)(_ + 1) + .covary[IO] + .evalTap(n => IO.raiseError(SevenNotAllowed).whenA(n == 7)) + + val downstream = source.groupChunksWithin(100, 2.seconds).map(_.toList) + + val expected = List((0 to 6).toList) + + downstream.assertEmits(expected).intercept[SevenNotAllowed.type] + } + } + + test( + "upstream interruption causes immediate downstream termination with all elements being emitted" + ) { + val sourceTimeout = 5.5.seconds + val downstreamTimeout = sourceTimeout + 2.seconds + + TestControl + .executeEmbed { + val source: Stream[IO, Int] = + Stream + .iterate(0)(_ + 1) + .covary[IO] + .meteredStartImmediately(1.second) + .interruptAfter(sourceTimeout) + + // large chunkSize and timeout: no emissions expected in the window + // specified, unless the source ends, due to interruption or natural + // termination (i.e. runs out of elements) + val downstream: Stream[IO, Chunk[Int]] = + source.groupChunksWithin(Int.MaxValue, 1.day) + + downstream.compile.lastOrError + .timeout(downstreamTimeout) + .map(_.toList) + .timed + } + .assertEquals((sourceTimeout, List(0, 1, 2, 3, 4, 5))) + } + + test("stress test: all elements are processed") { + val rangeLength = 10000 + + TestControl.executeEmbed { + Stream + .eval(Ref.of[IO, Int](0)) + .flatMap { counter => + Stream + .range(0, rangeLength) + .covary[IO] + .groupChunksWithin(4096, 100.micros) + .evalMap(ch => counter.updateAndGet(_ + ch.size)) + } + .compile + .lastOrError + .assertEquals(rangeLength) + } + } + } + property("head")(forAll((s: Stream[Pure, Int]) => assertEquals(s.head.toList, s.toList.take(1)))) group("ifEmpty") { diff --git a/core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala b/core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala new file mode 100644 index 0000000000..87238b36e9 --- /dev/null +++ b/core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala @@ -0,0 +1,88 @@ +/* + * Copyright (c) 2013 Functional Streams for Scala + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER + * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN + * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +package fs2 +package concurrent + +import cats.effect.IO +import cats.effect.testkit.TestControl +import cats.syntax.all._ + +import scala.concurrent.duration._ + +class ConditionedRefSuite extends Fs2Suite { + + test("waitUntil completes immediately if the condition already holds") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 1).flatMap { ref => + ref.waitUntil(_ > 0).timed.map { case (elapsed, _) => assertEquals(elapsed, Duration.Zero) } + } + } + } + + test("waitUntil completes on the update that makes the condition true") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + val updates = + IO.sleep(1.second) >> ref.update(_ + 1) >> IO.sleep(1.second) >> ref.update(_ + 1) + updates.background.surround { + ref.waitUntil(_ >= 2).timed.map { case (elapsed, _) => assertEquals(elapsed, 2.seconds) } + } + } + } + } + + test("waitUntil does not miss a condition that becomes true and then false again") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + ref.waitUntil(_ == 1).start.flatMap { waiter => + IO.sleep(1.second) >> ref.set(1) >> ref.set(2) >> + waiter.joinWithNever.timeout(1.second) >> ref.get.assertEquals(2) + } + } + } + } + + test("only waiters whose condition holds are woken") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + (ref.waitUntil(_ >= 1).start, ref.waitUntil(_ >= 2).start).flatMapN { (first, second) => + IO.sleep(1.second) >> ref.set(1) >> + first.joinWithNever.timeout(1.second) >> + IO.race(second.joinWithNever, IO.sleep(1.second)).map(r => assert(r.isRight)) >> + ref.set(2) >> + second.joinWithNever.timeout(1.second) + } + } + } + } + + test("a cancelled waiter does not affect later updates and waiters") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + ref.waitUntil(_ > 0).timeoutTo(duration = 1.second, fallback = IO.unit) >> + ref.update(_ + 1) >> + ref.waitUntil(_ > 0).timeout(1.second) >> + ref.get.assertEquals(1) + } + } + } +} diff --git a/integration/src/test/scala/fs2/MemoryLeakSpec.scala b/integration/src/test/scala/fs2/MemoryLeakSpec.scala index 4e2fa022f3..71afabe646 100644 --- a/integration/src/test/scala/fs2/MemoryLeakSpec.scala +++ b/integration/src/test/scala/fs2/MemoryLeakSpec.scala @@ -153,6 +153,27 @@ class MemoryLeakSpec extends FunSuite { .groupWithin(256, 1.second) } + leakTest("groupChunksWithin") { + Stream + .eval(IO.never) + .covary[IO] + .groupChunksWithin(Int.MaxValue, 1.millis) + } + + leakTest("groupChunksWithin --- Issue 2328") { + Stream + .range(0, 1000000) + .covary[IO] + .groupChunksWithin(256, 1.second) + } + + leakTest("groupChunksWithin --- repeated idle timeouts") { + Stream + .awakeEvery[IO](1.second) + .as(1) + .groupChunksWithin(256, 1.millis) + } + leakTest("topic continuous publish") { Stream .eval(Topic[IO, Int]) From 5e2d95d0fc8c8c285450f7d2f4b6b19b0559a832 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 12:32:28 -0400 Subject: [PATCH 02/11] Chunk.grouped returns Chunk[Chunk[O]] --- core/shared/src/main/scala/fs2/Chunk.scala | 22 ++++++++++++------- core/shared/src/main/scala/fs2/Stream.scala | 2 +- .../src/test/scala/fs2/ChunkSuite.scala | 9 ++++---- 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/core/shared/src/main/scala/fs2/Chunk.scala b/core/shared/src/main/scala/fs2/Chunk.scala index c99d56f47d..4c781d711f 100644 --- a/core/shared/src/main/scala/fs2/Chunk.scala +++ b/core/shared/src/main/scala/fs2/Chunk.scala @@ -211,18 +211,24 @@ abstract class Chunk[+O] extends Serializable with ChunkPlatform[O] with ChunkRu * * Like `List#grouped`, an empty chunk has no groups, and `n` must be positive. */ - def grouped(n: Int): Vector[Chunk[O]] = { + def grouped(n: Int): Chunk[Chunk[O]] = { require(n > 0, s"n must be positive, but got ${n.toString}") - @tailrec - def loop(rest: Chunk[O], acc: Vector[Chunk[O]]): Vector[Chunk[O]] = - if (rest.size <= n) acc :+ rest - else { - val (group, remaining) = rest.splitAt(n) - loop(rest = remaining, acc = acc :+ group) + if (isEmpty) Chunk.empty + else { + val numGroups = (size - 1) / n + 1 // ceil divide + val groups = new Array[Chunk[O]](numGroups) + var i = 0 + var chunk = this + while (i < numGroups) { + val (group, remaining) = chunk.splitAt(n) + groups(i) = group + chunk = remaining + i += 1 } - if (isEmpty) Vector.empty else loop(rest = this, acc = Vector.empty) + Chunk.array(groups) + } } /** Gets the first element of this chunk. */ diff --git a/core/shared/src/main/scala/fs2/Stream.scala b/core/shared/src/main/scala/fs2/Stream.scala index 2c7a31a10f..efdbe468f7 100644 --- a/core/shared/src/main/scala/fs2/Stream.scala +++ b/core/shared/src/main/scala/fs2/Stream.scala @@ -1686,7 +1686,7 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, def emitBatches: Pull[F2, Chunk[O], Unit] = Pull.eval(nextBatch).flatMap { - case Some(batch) => Pull.output(Chunk.from(batch.grouped(chunkSize))) >> emitBatches + case Some(batch) => Pull.output(batch.grouped(chunkSize)) >> emitBatches case None => Pull.done } diff --git a/core/shared/src/test/scala/fs2/ChunkSuite.scala b/core/shared/src/test/scala/fs2/ChunkSuite.scala index 78ef3b918b..3f59a45fb2 100644 --- a/core/shared/src/test/scala/fs2/ChunkSuite.scala +++ b/core/shared/src/test/scala/fs2/ChunkSuite.scala @@ -131,12 +131,11 @@ class ChunkSuite extends Fs2Suite { } } - test("Chunk.grouped of an empty chunk has no groups, like List#grouped") { - assertEquals(Chunk.empty[Int].grouped(3), Vector.empty) - assertEquals(List.empty[Int].grouped(3).toList, Nil) + test("Chunk.grouped of an empty chunk has no groups") { + assertEquals(Chunk.empty[Int].grouped(3), Chunk.empty) } - test("Chunk.grouped rejects a non-positive group size, like List#grouped") { + test("Chunk.grouped rejects a non-positive group size") { List(0, -1).foreach { n => intercept[IllegalArgumentException](Chunk(1, 2, 3).grouped(n)) intercept[IllegalArgumentException](Chunk.empty[Int].grouped(n)) @@ -146,7 +145,7 @@ class ChunkSuite extends Fs2Suite { test("Chunk.grouped returns a chunk that fits in one group as is") { forAll { (c: Chunk[Int]) => - if (c.nonEmpty) assert(c.grouped(c.size).head eq c) + if (c.nonEmpty) assert(c.grouped(c.size).head.get eq c) } } From d728d2e7da9c42a96483d5ec0b5cb4653fd05212 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 12:34:20 -0400 Subject: [PATCH 03/11] Use existing ref method --- .../main/scala/fs2/concurrent/ConditionedRef.scala | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala index e558ef01db..092d5aac94 100644 --- a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala +++ b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala @@ -95,14 +95,12 @@ private[fs2] object ConditionedRef { else F.deferred[Unit].flatMap { wake => val waiter = new Waiter(accepts = p, wake = wake) - F.uncancelable { poll => - state.modify { s => - if (p(s.value)) s -> F.unit - else - s.register(waiter) -> poll(wake.get).onCancel { - state.update(_.deregister(waiter)) - } - }.flatten + state.flatModifyFull { (poll, s) => + if (p(s.value)) s -> F.unit + else + s.register(waiter) -> poll(wake.get).onCancel { + state.update(_.deregister(waiter)) + } } } } From b38d8474680f58324fa83c66c733c7ae420d42e9 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 12:55:20 -0400 Subject: [PATCH 04/11] Try to justfy impl choice --- core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala | 2 ++ 1 file changed, 2 insertions(+) diff --git a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala index 092d5aac94..ffdad7bef6 100644 --- a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala +++ b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala @@ -79,6 +79,8 @@ private[fs2] object ConditionedRef { def modify[B](f: A => (A, B)): F[B] = state.flatModify { s => // uncancellable to avoid losing wake-up signals val (value, result) = f(s.value) + // Scan first instead of calling partition to avoid extra allocations which have a small impact on perf. + // This optimizes for the case where writes infrequently trigger wake-ups. if (!s.waiters.exists(_.accepts(value))) State(value = value, waiters = s.waiters) -> result.pure[F] else { From 39678fdec6e29edd5bc804e3041e7038ac46fb94 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 14:28:58 -0400 Subject: [PATCH 05/11] Use chunk-based impl for groupWithin; align chunking behavior around timeouts to existing api --- .../GroupChunksWithinBenchmark.scala | 64 ----- core/shared/src/main/scala/fs2/Stream.scala | 124 +------- .../scala/fs2/StreamCombinatorsSuite.scala | 265 +++--------------- .../src/test/scala/fs2/MemoryLeakSpec.scala | 16 +- 4 files changed, 44 insertions(+), 425 deletions(-) delete mode 100644 benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala diff --git a/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala b/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala deleted file mode 100644 index 9a008f5f47..0000000000 --- a/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala +++ /dev/null @@ -1,64 +0,0 @@ -/* - * Copyright (c) 2013 Functional Streams for Scala - * - * Permission is hereby granted, free of charge, to any person obtaining a copy of - * this software and associated documentation files (the "Software"), to deal in - * the Software without restriction, including without limitation the rights to - * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of - * the Software, and to permit persons to whom the Software is furnished to do so, - * subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS - * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR - * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER - * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN - * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - */ - -package fs2 -package benchmark - -import cats.effect.IO -import org.openjdk.jmh.annotations.{Benchmark, Param, Scope, State} - -import scala.concurrent.duration.* - -@State(Scope.Thread) -class GroupChunksWithinBenchmark { - - import cats.effect.unsafe.implicits.global - - val bufferWindow = 100.micros - - @Param(Array("100", "10000", "100000")) - var rangeLength: Int = _ - - @Param(Array("16", "256", "4096")) - var bufferSize: Int = _ - - @Benchmark - def groupChunksWithin(): Unit = - Stream - .range(0, rangeLength) - .covary[IO] - .groupChunksWithin(bufferSize, bufferWindow) - .compile - .drain - .unsafeRunSync() - - @Benchmark - def groupChunksWithinChunkedUpstream(): Unit = - Stream - .range(0, rangeLength) - .chunkN(bufferSize / 4 + 1) - .unchunks - .covary[IO] - .groupChunksWithin(bufferSize, bufferWindow) - .compile - .drain - .unsafeRunSync() -} diff --git a/core/shared/src/main/scala/fs2/Stream.scala b/core/shared/src/main/scala/fs2/Stream.scala index efdbe468f7..701dca9e10 100644 --- a/core/shared/src/main/scala/fs2/Stream.scala +++ b/core/shared/src/main/scala/fs2/Stream.scala @@ -1511,116 +1511,6 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, def groupWithin[F2[x] >: F[x]]( chunkSize: Int, timeout: FiniteDuration - )(implicit F: Temporal[F2]): Stream[F2, Chunk[O]] = { - - case class JunctionBuffer[T]( - data: Vector[T], - endOfSupply: Option[Either[Throwable, Unit]], - endOfDemand: Option[Either[Throwable, Unit]] - ) { - def splitAt(n: Int): (JunctionBuffer[T], JunctionBuffer[T]) = - if (this.data.size >= n) { - val (head, tail) = this.data.splitAt(n.toInt) - (this.copy(tail), this.copy(head)) - } else { - (this.copy(Vector.empty), this) - } - } - - val outputLong = chunkSize.toLong - fs2.Stream.force { - for { - demand <- Semaphore[F2](outputLong) - supply <- Semaphore[F2](0L) - buffer <- Ref[F2].of( - JunctionBuffer[O](Vector.empty[O], endOfSupply = None, endOfDemand = None) - ) - } yield { - /* - Buffer: stores items from input to be sent on next output chunk - * - Demand Semaphore: to avoid adding too many items to buffer - * - Supply: counts filled positions for next output chunk */ - def enqueue(t: O): F2[Boolean] = - for { - _ <- demand.acquire - buf <- buffer.modify(buf => (buf.copy(buf.data :+ t), buf)) - _ <- supply.release - } yield buf.endOfDemand.isEmpty - - val dequeueNextOutput: F2[Option[Vector[O]]] = { - // Trigger: waits until the supply buffer is full (with acquireN) - val waitSupply = supply.acquireN(outputLong).guaranteeCase { - case Outcome.Succeeded(_) => supply.releaseN(outputLong) - case _ => F.unit - } - - val onTimeout: F2[Long] = - for { - _ <- supply.acquire // waits until there is at least one element in buffer - m <- supply.available - k = m.min(outputLong - 1) - b <- supply.tryAcquireN(k) - } yield if (b) k + 1 else 1 - - // in JS cancellation doesn't always seem to run, so race conditions should restore state on their own - for { - acq <- F.race(F.sleep(timeout), waitSupply).flatMap { - case Left(_) => onTimeout - case Right(_) => supply.acquireN(outputLong).as(outputLong) - } - buf <- buffer.modify(_.splitAt(acq.toInt)) - _ <- demand.releaseN(buf.data.size.toLong) - res <- buf.endOfSupply match { - case Some(Left(error)) => F.raiseError(error) - case Some(Right(_)) if buf.data.isEmpty => F.pure(None) - case _ => F.pure(Some(buf.data)) - } - } yield res - } - - def endSupply(result: Either[Throwable, Unit]): F2[Unit] = - buffer.update(_.copy(endOfSupply = Some(result))) *> supply.releaseN( - // enough supply for 2 iterations of the race loop in case of upstream - // interruption: so that downstream can terminate immediately - outputLong * 2 - ) - - def endDemand(result: Either[Throwable, Unit]): F2[Unit] = - buffer.update(_.copy(endOfDemand = Some(result))) *> demand.releaseN(Int.MaxValue) - - def toEnding(ec: ExitCase): Either[Throwable, Unit] = ec match { - case ExitCase.Succeeded => Right(()) - case ExitCase.Errored(e) => Left(e) - case ExitCase.Canceled => Right(()) - } - - val enqueueAsync = F.start { - this - .evalMap(enqueue) - .forall(identity) - .onFinalizeCase(ec => endSupply(toEnding(ec))) - .compile - .drain - } - - val outputStream: Stream[F2, Chunk[O]] = - Stream - .eval(dequeueNextOutput) - .repeat - .collectWhile { case Some(data) => Chunk.from(data) } - - Stream - .bracketCase(enqueueAsync) { case (upstream, exitCase) => - endDemand(toEnding(exitCase)) *> upstream.cancel - } >> outputStream - } - } - } - - /** Like [[groupWithin]], but works on chunks for better performance. - */ - def groupChunksWithin[F2[x] >: F[x]]( - chunkSize: Int, - timeout: FiniteDuration )(implicit F: Temporal[F2]): Stream[F2, Chunk[O]] = Stream.force { require(chunkSize > 0, s"chunkSize must be > 0, but got ${chunkSize.toString}") @@ -1653,11 +1543,11 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, } // None means upstream is done, and the buffer is drained. - def takeOrExit(all: Boolean): F2[Option[Chunk[O]]] = + def takeOrExit(timedOut: Boolean): F2[Option[Chunk[O]]] = buffer.modify { b => - if (b.nonEmpty && (all || b.isDone)) take(b, n = b.size) - // Take all potential batches so that producer doesn't need to wait for timeout + if (b.nonEmpty && b.isDone) take(b, n = b.size) else if (b.isFull) take(b, n = b.size - b.size % chunkSize) + else if (b.nonEmpty && timedOut) take(b, n = b.size) else b -> (b.done match { case None /* not full & no-timeout */ => Skip @@ -1666,20 +1556,20 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, }) }.flatten - val onTimeout = takeOrExit(all = true).flatMap { + val onTimeout = takeOrExit(timedOut = true).flatMap { case Some(batch) if batch.isEmpty => - buffer.waitUntil(b => b.nonEmpty || b.isDone) >> takeOrExit(all = true) + buffer.waitUntil(b => b.nonEmpty || b.isDone) >> takeOrExit(timedOut = true) case result => result.pure[F2] } val nextBatch: F2[Option[Chunk[O]]] = // Potentially skip starting timer fiber if buffer is full - takeOrExit(all = false).flatMap { + takeOrExit(timedOut = false).flatMap { case Some(batch) if batch.isEmpty => F.race(F.sleep(timeout), buffer.waitUntil(b => b.isFull || b.isDone)) .flatMap { case Left(_ /* timeout */ ) => onTimeout - case Right(_ /* full batch or done */ ) => takeOrExit(all = false) + case Right(_ /* full batch or done */ ) => takeOrExit(timedOut = false) } case result => result.pure[F2] } diff --git a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala index b0ea3ade29..a7cae53751 100644 --- a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala +++ b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala @@ -763,7 +763,7 @@ class StreamCombinatorsSuite extends Fs2Suite { s.covary[IO] .evalTap(sleep) .groupWithin(groupSize, timeout) - .flatMap(Stream.chunk) + .unchunks .assertEmitsSameAs(s) } } @@ -856,6 +856,37 @@ class StreamCombinatorsSuite extends Fs2Suite { } } + test("accumulation and splitting with timeouts between bursts") { + val t = 200.millis + val chunkSize = 5 + + def chunk(from: Int, to: Int) = + Stream.range(from, to + 1).chunkAll.unchunks + + val source = + chunk(from = 1, to = 3) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 4, to = 15) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 16, to = 22) + + val expected = List( + (1 to 3).toList, + (4 to 8).toList, + (9 to 13).toList, + (14 to 15).toList, + (16 to 20).toList, + (21 to 22).toList + ) + + TestControl.executeEmbed { + source + .groupWithin(chunkSize, t) + .map(_.toList) + .assertEmits(expected) + } + } + test("does not reset timeout if nothing is emitted") { TestControl .executeEmbed( @@ -985,44 +1016,13 @@ class StreamCombinatorsSuite extends Fs2Suite { .assertEquals(rangeLength) } } - } - - group("groupChunksWithin") { - implicit val groupSizeArb: Arbitrary[Int] = Arbitrary(Gen.choose(1, 20)) - // a zero timeout is covered separately. Under TestControl it makes the idle loop spin forever - implicit val timeoutArb: Arbitrary[FiniteDuration] = Arbitrary(Gen.choose(1, 50).map(_.millis)) - - def sleep(d: Int): IO[Unit] = IO.sleep((d % 500).abs.micros) - - test("should never lose any elements") { - forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration) => - TestControl.executeEmbed { - s.covary[IO] - .evalTap(sleep) - .groupChunksWithin(chunkSize = 3, timeout) - .unchunks - .assertEmitsSameAs(s) - } - } - } - - test("should never emit empty chunks") { - forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration, groupSize: Int) => - TestControl.executeEmbed { - s.covary[IO] - .evalTap(sleep) - .groupChunksWithin(groupSize, timeout) - .assertForall(_.nonEmpty) - } - } - } test("makes progress with a zero timeout") { Stream .range(0, 100) .covary[IO] .evalTap(_ => IO.sleep(1.milli)) - .groupChunksWithin(chunkSize = 10, timeout = Duration.Zero) + .groupWithin(chunkSize = 10, timeout = Duration.Zero) .unchunks .compile .toList @@ -1030,30 +1030,12 @@ class StreamCombinatorsSuite extends Fs2Suite { .assertEquals((0 until 100).toList) } - test("every chunk but the last has at least chunkSize elements when no timeout triggers") { - val source = Stream.range(0, 101) - val chunkSize = 5 - - TestControl.executeEmbed { - source - .covary[IO] - .groupChunksWithin(chunkSize, 1.day) - .compile - .toList - .map { chunks => - assertEquals(chunks.flatMap(_.toList), source.toList) - assert(chunks.init.forall(_.size >= chunkSize)) - assert(chunks.lastOption.forall(_.size < chunkSize)) - } - } - } - test("splits upstream chunks larger than chunkSize") { TestControl.executeEmbed { Stream .emits(1 to 12) .covary[IO] - .groupChunksWithin(chunkSize = 5, timeout = 1.day) + .groupWithin(chunkSize = 5, timeout = 1.day) .map(_.toList) .assertEmits(List((1 to 5).toList, (6 to 10).toList, (11 to 12).toList)) } @@ -1062,7 +1044,7 @@ class StreamCombinatorsSuite extends Fs2Suite { test("emits every full group of a large upstream chunk without waiting for the timeout") { TestControl.executeEmbed { (Stream.emits(1 to 100) ++ Stream.never[IO]) - .groupChunksWithin(chunkSize = 10, timeout = 1.day) + .groupWithin(chunkSize = 10, timeout = 1.day) .take(10) .compile .toList @@ -1074,55 +1056,6 @@ class StreamCombinatorsSuite extends Fs2Suite { } } - test( - "should return a finite stream back in a single chunk given a chunk size equal to the stream size and an absurdly high duration" - ) { - forAllF { (streamAsList0: List[Int]) => - val streamAsList = 0 :: streamAsList0 - TestControl.executeEmbed { - Stream - .emits(streamAsList) - .covary[IO] - .groupChunksWithin(streamAsList.size, (Int.MaxValue - 1L).nanoseconds) - .compile - .toList - .map(_.head.toList) - .assertEquals(streamAsList) - } - } - } - - test("accumulation with splitting") { - val t = 200.millis - val chunkSize = 5 - - def chunk(from: Int, to: Int) = - Stream.range(from, to + 1).chunkAll.unchunks - - val source = - chunk(from = 1, to = 3) ++ - Stream.sleep_[IO](t + t / 2) ++ - chunk(from = 4, to = 15) ++ - Stream.sleep_[IO](t + t / 2) ++ - chunk(from = 16, to = 22) - - val expected = List( - (1 to 3).toList, - (4 to 8).toList, - (9 to 13).toList, - (14 to 15).toList, - (16 to 20).toList, - (21 to 22).toList - ) - - TestControl.executeEmbed { - source - .groupChunksWithin(chunkSize, t) - .map(_.toList) - .assertEmits(expected) - } - } - test("emits a burst arriving after several idle timeouts") { val timeout = 100.millis val chunkSize = 5 @@ -1135,137 +1068,11 @@ class StreamCombinatorsSuite extends Fs2Suite { TestControl.executeEmbed { source - .groupChunksWithin(chunkSize, timeout) + .groupWithin(chunkSize, timeout) .map(_.toList) .assertEmits(List(List(1, 2, 3), List(4, 5, 6))) } } - - test("does not reset timeout if nothing is emitted") { - TestControl - .executeEmbed( - Ref[IO] - .of(0.millis) - .flatMap { ref => - val timeout = 5.seconds - - def measureEmission[A]: Pipe[IO, A, A] = - _.chunks - .evalTap(_ => IO.monotonic.flatMap(ref.set)) - .unchunks - - // emits elements after the timeout has expired - val source = - Stream.sleep_[IO](timeout + 200.millis) ++ - Stream(4, 5) ++ - Stream.never[IO] // avoids emission due to source termination - - source - .through(measureEmission) - .groupChunksWithin(5, timeout) - .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) - .interruptAfter(timeout * 3) - .compile - .lastOrError - } - ) - .assertEquals(0.millis) // The stream emits after the timeout has expired - // on an empty buffer, so groupChunksWithin should re-emit with zero delay - // rather than start a fresh timeout window - } - - test("Edge case: should not introduce unnecessary delays when chunkSize == chunk size") { - TestControl - .executeEmbed( - Ref[IO] - .of(0.millis) - .flatMap { ref => - val timeout = 5.seconds - - def measureEmission[A]: Pipe[IO, A, A] = - _.chunks - .evalTap(_ => IO.monotonic.flatMap(ref.set)) - .unchunks - - val source = - Stream(1, 2, 3) ++ - Stream.sleep_[IO](timeout + 200.millis) - - source - .through(measureEmission) - .groupChunksWithin(3, timeout) - .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) - .compile - .lastOrError - } - ) - .assertEquals(0.millis) - } - - test("upstream failures are propagated downstream") { - TestControl.executeEmbed { - case object SevenNotAllowed extends NoStackTrace - - val source = Stream - .iterate(0)(_ + 1) - .covary[IO] - .evalTap(n => IO.raiseError(SevenNotAllowed).whenA(n == 7)) - - val downstream = source.groupChunksWithin(100, 2.seconds).map(_.toList) - - val expected = List((0 to 6).toList) - - downstream.assertEmits(expected).intercept[SevenNotAllowed.type] - } - } - - test( - "upstream interruption causes immediate downstream termination with all elements being emitted" - ) { - val sourceTimeout = 5.5.seconds - val downstreamTimeout = sourceTimeout + 2.seconds - - TestControl - .executeEmbed { - val source: Stream[IO, Int] = - Stream - .iterate(0)(_ + 1) - .covary[IO] - .meteredStartImmediately(1.second) - .interruptAfter(sourceTimeout) - - // large chunkSize and timeout: no emissions expected in the window - // specified, unless the source ends, due to interruption or natural - // termination (i.e. runs out of elements) - val downstream: Stream[IO, Chunk[Int]] = - source.groupChunksWithin(Int.MaxValue, 1.day) - - downstream.compile.lastOrError - .timeout(downstreamTimeout) - .map(_.toList) - .timed - } - .assertEquals((sourceTimeout, List(0, 1, 2, 3, 4, 5))) - } - - test("stress test: all elements are processed") { - val rangeLength = 10000 - - TestControl.executeEmbed { - Stream - .eval(Ref.of[IO, Int](0)) - .flatMap { counter => - Stream - .range(0, rangeLength) - .covary[IO] - .groupChunksWithin(4096, 100.micros) - .evalMap(ch => counter.updateAndGet(_ + ch.size)) - } - .compile - .lastOrError - .assertEquals(rangeLength) - } - } } property("head")(forAll((s: Stream[Pure, Int]) => assertEquals(s.head.toList, s.toList.take(1)))) diff --git a/integration/src/test/scala/fs2/MemoryLeakSpec.scala b/integration/src/test/scala/fs2/MemoryLeakSpec.scala index 71afabe646..4e00614322 100644 --- a/integration/src/test/scala/fs2/MemoryLeakSpec.scala +++ b/integration/src/test/scala/fs2/MemoryLeakSpec.scala @@ -153,25 +153,11 @@ class MemoryLeakSpec extends FunSuite { .groupWithin(256, 1.second) } - leakTest("groupChunksWithin") { - Stream - .eval(IO.never) - .covary[IO] - .groupChunksWithin(Int.MaxValue, 1.millis) - } - - leakTest("groupChunksWithin --- Issue 2328") { - Stream - .range(0, 1000000) - .covary[IO] - .groupChunksWithin(256, 1.second) - } - leakTest("groupChunksWithin --- repeated idle timeouts") { Stream .awakeEvery[IO](1.second) .as(1) - .groupChunksWithin(256, 1.millis) + .groupWithin(256, 1.millis) } leakTest("topic continuous publish") { From 78b87db6eacc4c0c914ad47208f042f915241577 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Fri, 2 Oct 2026 15:44:42 -0400 Subject: [PATCH 06/11] Add a performant groupChunksWithin --- .../GroupChunksWithinBenchmark.scala | 64 ++++ .../fs2/benchmark/GroupWithinBenchmark.scala | 12 + core/shared/src/main/scala/fs2/Chunk.scala | 18 ++ core/shared/src/main/scala/fs2/Stream.scala | 78 +++++ .../scala/fs2/concurrent/ConditionedRef.scala | 110 +++++++ .../src/test/scala/fs2/ChunkSuite.scala | 24 ++ .../scala/fs2/StreamCombinatorsSuite.scala | 281 ++++++++++++++++++ .../fs2/concurrent/ConditionedRefSuite.scala | 88 ++++++ .../src/test/scala/fs2/MemoryLeakSpec.scala | 21 ++ 9 files changed, 696 insertions(+) create mode 100644 benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala create mode 100644 core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala create mode 100644 core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala diff --git a/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala b/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala new file mode 100644 index 0000000000..9a008f5f47 --- /dev/null +++ b/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2013 Functional Streams for Scala + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER + * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN + * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +package fs2 +package benchmark + +import cats.effect.IO +import org.openjdk.jmh.annotations.{Benchmark, Param, Scope, State} + +import scala.concurrent.duration.* + +@State(Scope.Thread) +class GroupChunksWithinBenchmark { + + import cats.effect.unsafe.implicits.global + + val bufferWindow = 100.micros + + @Param(Array("100", "10000", "100000")) + var rangeLength: Int = _ + + @Param(Array("16", "256", "4096")) + var bufferSize: Int = _ + + @Benchmark + def groupChunksWithin(): Unit = + Stream + .range(0, rangeLength) + .covary[IO] + .groupChunksWithin(bufferSize, bufferWindow) + .compile + .drain + .unsafeRunSync() + + @Benchmark + def groupChunksWithinChunkedUpstream(): Unit = + Stream + .range(0, rangeLength) + .chunkN(bufferSize / 4 + 1) + .unchunks + .covary[IO] + .groupChunksWithin(bufferSize, bufferWindow) + .compile + .drain + .unsafeRunSync() +} diff --git a/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala b/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala index 9ffb7eb912..5e989539a3 100644 --- a/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala +++ b/benchmark/src/main/scala/fs2/benchmark/GroupWithinBenchmark.scala @@ -59,4 +59,16 @@ class GroupWithinBenchmark { .compile .drain .unsafeRunSync() + + @Benchmark + def groupWithinChunkedUpstream(): Unit = + Stream + .range(0, rangeLength) + .chunkN(bufferSize / 4 + 1) + .unchunks + .covary[IO] + .groupWithin(bufferSize, bufferWindow) + .compile + .drain + .unsafeRunSync() } diff --git a/core/shared/src/main/scala/fs2/Chunk.scala b/core/shared/src/main/scala/fs2/Chunk.scala index 03aac9d5aa..c99d56f47d 100644 --- a/core/shared/src/main/scala/fs2/Chunk.scala +++ b/core/shared/src/main/scala/fs2/Chunk.scala @@ -207,6 +207,24 @@ abstract class Chunk[+O] extends Serializable with ChunkPlatform[O] with ChunkRu } } + /** Splits this chunk into groups of `n` elements, the last of which may be smaller. + * + * Like `List#grouped`, an empty chunk has no groups, and `n` must be positive. + */ + def grouped(n: Int): Vector[Chunk[O]] = { + require(n > 0, s"n must be positive, but got ${n.toString}") + + @tailrec + def loop(rest: Chunk[O], acc: Vector[Chunk[O]]): Vector[Chunk[O]] = + if (rest.size <= n) acc :+ rest + else { + val (group, remaining) = rest.splitAt(n) + loop(rest = remaining, acc = acc :+ group) + } + + if (isEmpty) Vector.empty else loop(rest = this, acc = Vector.empty) + } + /** Gets the first element of this chunk. */ def head: Option[O] = if (isEmpty) None else Some(apply(0)) diff --git a/core/shared/src/main/scala/fs2/Stream.scala b/core/shared/src/main/scala/fs2/Stream.scala index c95539b649..468cece318 100644 --- a/core/shared/src/main/scala/fs2/Stream.scala +++ b/core/shared/src/main/scala/fs2/Stream.scala @@ -1625,6 +1625,84 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, } } + /** Like [[groupWithin]], but works on chunks for better performance. + */ + def groupChunksWithin[F2[x] >: F[x]]( + chunkSize: Int, + timeout: FiniteDuration + )(implicit F: Temporal[F2]): Stream[F2, Chunk[O]] = + Stream.force { + require(chunkSize > 0, s"chunkSize must be > 0, but got ${chunkSize.toString}") + + val StopConsumer = none[Chunk[O]].pure[F2] + val Skip = Chunk.empty[O].some.pure[F2] + + final case class Buffer[A](chunk: Chunk[A], done: Option[ExitCase]) { + def size: Int = chunk.size + def isEmpty: Boolean = chunk.isEmpty + def nonEmpty: Boolean = chunk.nonEmpty + def isFull: Boolean = size >= chunkSize + def isDone: Boolean = done.isDefined + } + + ConditionedRef.of[F2, Buffer[O]](Buffer(chunk = Chunk.empty[O], done = none)).map { buffer => + val producer = chunks + .evalMap { chunk => + buffer + .updateAndGet(b => b.copy(chunk = b.chunk ++ chunk)) + .flatMap(b => F.whenA(b.isFull)(buffer.waitUntil(!_.isFull))) + } + .onFinalizeCase(exitCase => buffer.update(_.copy(done = exitCase.some))) + .compile + .drain + + def take(b: Buffer[O], n: Int): (Buffer[O], F2[Option[Chunk[O]]]) = { + val (taken, remaining) = b.chunk.splitAt(n) + b.copy(chunk = remaining) -> taken.some.pure[F2] + } + + // None means upstream is done, and the buffer is drained. + def takeOrExit(all: Boolean): F2[Option[Chunk[O]]] = + buffer.modify { b => + if (b.nonEmpty && (all || b.isDone)) take(b, n = b.size) + // Take all potential batches so that producer doesn't need to wait for timeout + else if (b.isFull) take(b, n = b.size - b.size % chunkSize) + else + b -> (b.done match { + case None /* not full & no-timeout */ => Skip + case Some(ExitCase.Errored(e)) => F.raiseError[Option[Chunk[O]]](e) + case Some(_) /* empty */ => StopConsumer + }) + }.flatten + + val onTimeout = takeOrExit(all = true).flatMap { + case Some(batch) if batch.isEmpty => + buffer.waitUntil(b => b.nonEmpty || b.isDone) >> takeOrExit(all = true) + case result => result.pure[F2] + } + + val nextBatch: F2[Option[Chunk[O]]] = + // Potentially skip starting timer fiber if buffer is full + takeOrExit(all = false).flatMap { + case Some(batch) if batch.isEmpty => + F.race(F.sleep(timeout), buffer.waitUntil(b => b.isFull || b.isDone)) + .flatMap { + case Left(_ /* timeout */ ) => onTimeout + case Right(_ /* full batch or done */ ) => takeOrExit(all = false) + } + case result => result.pure[F2] + } + + def emitBatches: Pull[F2, Chunk[O], Unit] = + Pull.eval(nextBatch).flatMap { + case Some(batch) => Pull.output(Chunk.from(batch.grouped(chunkSize))) >> emitBatches + case None => Pull.done + } + + Stream.bracket(producer.start)(_.cancel) >> emitBatches.stream + } + } + /** If `this` terminates with `Stream.raiseError(e)`, invoke `h(e)`. * * @example {{{ diff --git a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala new file mode 100644 index 0000000000..e558ef01db --- /dev/null +++ b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala @@ -0,0 +1,110 @@ +/* + * Copyright (c) 2013 Functional Streams for Scala + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER + * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN + * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +package fs2 +package concurrent + +import cats.effect._ +import cats.effect.implicits._ +import cats.syntax.all._ + +/** A `Ref` whose value can be waited on. + * + * Waiters register a predicate on the value. Every `modify` evaluates the registered predicates + * against the new value and wakes exactly the waiters whose predicate is satisfied. Unlike [[SignallingRef]], + * waiters are not woken by unrelated updates. + */ +private[fs2] sealed trait ConditionedRef[F[_], A] { + + def get: F[A] + + /** Atomically updates the value and wakes every waiter whose predicate condition holds for the new value. */ + def modify[B](f: A => (A, B)): F[B] + + def update(f: A => A): F[Unit] + + def updateAndGet(f: A => A): F[A] = + modify { a => + val newA = f(a) + (newA, newA) + } + + def set(a: A): F[Unit] = update(_ => a) + + /** Completes if the predicate, `p` holds for the current value, or a later `modify` sets a value that satisfies `p`. + * + * `p` may no longer hold by the time this completes: act on the value through `modify`. + */ + def waitUntil(p: A => Boolean): F[Unit] +} + +private[fs2] object ConditionedRef { + + def of[F[_], A](initial: A)(implicit F: Concurrent[F]): F[ConditionedRef[F, A]] = + F.ref(State[F, A](value = initial, waiters = Nil)).map(new Impl(_)) + + private final class Waiter[F[_], A](val accepts: A => Boolean, val wake: Deferred[F, Unit]) { + def wakeUp: F[Boolean] = wake.complete(()) + } + + private final case class State[F[_], A](value: A, waiters: List[Waiter[F, A]]) { + def register(waiter: Waiter[F, A]): State[F, A] = copy(waiters = waiter :: waiters) + def deregister(waiter: Waiter[F, A]): State[F, A] = + copy(waiters = waiters.filterNot(_ eq waiter)) + } + + private final class Impl[F[_], A](state: Ref[F, State[F, A]])(implicit F: Concurrent[F]) + extends ConditionedRef[F, A] { + + def get: F[A] = state.get.map(_.value) + + def modify[B](f: A => (A, B)): F[B] = + state.flatModify { s => // uncancellable to avoid losing wake-up signals + val (value, result) = f(s.value) + if (!s.waiters.exists(_.accepts(value))) + State(value = value, waiters = s.waiters) -> result.pure[F] + else { + val (toWake, waiting) = s.waiters.partition(_.accepts(value)) + State(value = value, waiters = waiting) -> toWake.traverse_(_.wakeUp).as(result) + } + } + + def update(f: A => A): F[Unit] = modify(a => (f(a), ())) + + def waitUntil(p: A => Boolean): F[Unit] = + get.flatMap { value => + if (p(value)) F.unit + else + F.deferred[Unit].flatMap { wake => + val waiter = new Waiter(accepts = p, wake = wake) + F.uncancelable { poll => + state.modify { s => + if (p(s.value)) s -> F.unit + else + s.register(waiter) -> poll(wake.get).onCancel { + state.update(_.deregister(waiter)) + } + }.flatten + } + } + } + } +} diff --git a/core/shared/src/test/scala/fs2/ChunkSuite.scala b/core/shared/src/test/scala/fs2/ChunkSuite.scala index 33b7e9967e..78ef3b918b 100644 --- a/core/shared/src/test/scala/fs2/ChunkSuite.scala +++ b/core/shared/src/test/scala/fs2/ChunkSuite.scala @@ -131,6 +131,25 @@ class ChunkSuite extends Fs2Suite { } } + test("Chunk.grouped of an empty chunk has no groups, like List#grouped") { + assertEquals(Chunk.empty[Int].grouped(3), Vector.empty) + assertEquals(List.empty[Int].grouped(3).toList, Nil) + } + + test("Chunk.grouped rejects a non-positive group size, like List#grouped") { + List(0, -1).foreach { n => + intercept[IllegalArgumentException](Chunk(1, 2, 3).grouped(n)) + intercept[IllegalArgumentException](Chunk.empty[Int].grouped(n)) + intercept[IllegalArgumentException](List(1, 2, 3).grouped(n)) + } + } + + test("Chunk.grouped returns a chunk that fits in one group as is") { + forAll { (c: Chunk[Int]) => + if (c.nonEmpty) assert(c.grouped(c.size).head eq c) + } + } + class OddStringExtractor { val callCounter: AtomicInteger = new AtomicInteger(0) @@ -167,6 +186,11 @@ class ChunkSuite extends Fs2Suite { property("isEmpty") { forAll((c: Chunk[A]) => assertEquals(c.isEmpty, c.toList.isEmpty)) } + property("grouped") { + forAll(genChunk, Gen.choose(1, 50)) { (c: Chunk[A], n: Int) => + assertEquals(c.grouped(n).map(_.toList).toList, c.toList.grouped(n).toList) + } + } property("toArray") { forAll { (c: Chunk[A]) => assertEquals(c.toArray.toVector, c.toVector) diff --git a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala index 48ed6e1011..b0ea3ade29 100644 --- a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala +++ b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala @@ -987,6 +987,287 @@ class StreamCombinatorsSuite extends Fs2Suite { } } + group("groupChunksWithin") { + implicit val groupSizeArb: Arbitrary[Int] = Arbitrary(Gen.choose(1, 20)) + // a zero timeout is covered separately. Under TestControl it makes the idle loop spin forever + implicit val timeoutArb: Arbitrary[FiniteDuration] = Arbitrary(Gen.choose(1, 50).map(_.millis)) + + def sleep(d: Int): IO[Unit] = IO.sleep((d % 500).abs.micros) + + test("should never lose any elements") { + forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration) => + TestControl.executeEmbed { + s.covary[IO] + .evalTap(sleep) + .groupChunksWithin(chunkSize = 3, timeout) + .unchunks + .assertEmitsSameAs(s) + } + } + } + + test("should never emit empty chunks") { + forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration, groupSize: Int) => + TestControl.executeEmbed { + s.covary[IO] + .evalTap(sleep) + .groupChunksWithin(groupSize, timeout) + .assertForall(_.nonEmpty) + } + } + } + + test("makes progress with a zero timeout") { + Stream + .range(0, 100) + .covary[IO] + .evalTap(_ => IO.sleep(1.milli)) + .groupChunksWithin(chunkSize = 10, timeout = Duration.Zero) + .unchunks + .compile + .toList + .timeout(30.seconds) + .assertEquals((0 until 100).toList) + } + + test("every chunk but the last has at least chunkSize elements when no timeout triggers") { + val source = Stream.range(0, 101) + val chunkSize = 5 + + TestControl.executeEmbed { + source + .covary[IO] + .groupChunksWithin(chunkSize, 1.day) + .compile + .toList + .map { chunks => + assertEquals(chunks.flatMap(_.toList), source.toList) + assert(chunks.init.forall(_.size >= chunkSize)) + assert(chunks.lastOption.forall(_.size < chunkSize)) + } + } + } + + test("splits upstream chunks larger than chunkSize") { + TestControl.executeEmbed { + Stream + .emits(1 to 12) + .covary[IO] + .groupChunksWithin(chunkSize = 5, timeout = 1.day) + .map(_.toList) + .assertEmits(List((1 to 5).toList, (6 to 10).toList, (11 to 12).toList)) + } + } + + test("emits every full group of a large upstream chunk without waiting for the timeout") { + TestControl.executeEmbed { + (Stream.emits(1 to 100) ++ Stream.never[IO]) + .groupChunksWithin(chunkSize = 10, timeout = 1.day) + .take(10) + .compile + .toList + .timed + .map { case (elapsed, chunks) => + assertEquals(elapsed, Duration.Zero) + assertEquals(chunks.map(_.toList), (1 to 100).toList.grouped(10).toList) + } + } + } + + test( + "should return a finite stream back in a single chunk given a chunk size equal to the stream size and an absurdly high duration" + ) { + forAllF { (streamAsList0: List[Int]) => + val streamAsList = 0 :: streamAsList0 + TestControl.executeEmbed { + Stream + .emits(streamAsList) + .covary[IO] + .groupChunksWithin(streamAsList.size, (Int.MaxValue - 1L).nanoseconds) + .compile + .toList + .map(_.head.toList) + .assertEquals(streamAsList) + } + } + } + + test("accumulation with splitting") { + val t = 200.millis + val chunkSize = 5 + + def chunk(from: Int, to: Int) = + Stream.range(from, to + 1).chunkAll.unchunks + + val source = + chunk(from = 1, to = 3) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 4, to = 15) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 16, to = 22) + + val expected = List( + (1 to 3).toList, + (4 to 8).toList, + (9 to 13).toList, + (14 to 15).toList, + (16 to 20).toList, + (21 to 22).toList + ) + + TestControl.executeEmbed { + source + .groupChunksWithin(chunkSize, t) + .map(_.toList) + .assertEmits(expected) + } + } + + test("emits a burst arriving after several idle timeouts") { + val timeout = 100.millis + val chunkSize = 5 + + val source = + Stream.sleep_[IO](timeout * 10) ++ + Stream.emits(1 to 3) ++ + Stream.sleep_[IO](timeout * 10) ++ + Stream.emits(4 to 6) + + TestControl.executeEmbed { + source + .groupChunksWithin(chunkSize, timeout) + .map(_.toList) + .assertEmits(List(List(1, 2, 3), List(4, 5, 6))) + } + } + + test("does not reset timeout if nothing is emitted") { + TestControl + .executeEmbed( + Ref[IO] + .of(0.millis) + .flatMap { ref => + val timeout = 5.seconds + + def measureEmission[A]: Pipe[IO, A, A] = + _.chunks + .evalTap(_ => IO.monotonic.flatMap(ref.set)) + .unchunks + + // emits elements after the timeout has expired + val source = + Stream.sleep_[IO](timeout + 200.millis) ++ + Stream(4, 5) ++ + Stream.never[IO] // avoids emission due to source termination + + source + .through(measureEmission) + .groupChunksWithin(5, timeout) + .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) + .interruptAfter(timeout * 3) + .compile + .lastOrError + } + ) + .assertEquals(0.millis) // The stream emits after the timeout has expired + // on an empty buffer, so groupChunksWithin should re-emit with zero delay + // rather than start a fresh timeout window + } + + test("Edge case: should not introduce unnecessary delays when chunkSize == chunk size") { + TestControl + .executeEmbed( + Ref[IO] + .of(0.millis) + .flatMap { ref => + val timeout = 5.seconds + + def measureEmission[A]: Pipe[IO, A, A] = + _.chunks + .evalTap(_ => IO.monotonic.flatMap(ref.set)) + .unchunks + + val source = + Stream(1, 2, 3) ++ + Stream.sleep_[IO](timeout + 200.millis) + + source + .through(measureEmission) + .groupChunksWithin(3, timeout) + .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) + .compile + .lastOrError + } + ) + .assertEquals(0.millis) + } + + test("upstream failures are propagated downstream") { + TestControl.executeEmbed { + case object SevenNotAllowed extends NoStackTrace + + val source = Stream + .iterate(0)(_ + 1) + .covary[IO] + .evalTap(n => IO.raiseError(SevenNotAllowed).whenA(n == 7)) + + val downstream = source.groupChunksWithin(100, 2.seconds).map(_.toList) + + val expected = List((0 to 6).toList) + + downstream.assertEmits(expected).intercept[SevenNotAllowed.type] + } + } + + test( + "upstream interruption causes immediate downstream termination with all elements being emitted" + ) { + val sourceTimeout = 5.5.seconds + val downstreamTimeout = sourceTimeout + 2.seconds + + TestControl + .executeEmbed { + val source: Stream[IO, Int] = + Stream + .iterate(0)(_ + 1) + .covary[IO] + .meteredStartImmediately(1.second) + .interruptAfter(sourceTimeout) + + // large chunkSize and timeout: no emissions expected in the window + // specified, unless the source ends, due to interruption or natural + // termination (i.e. runs out of elements) + val downstream: Stream[IO, Chunk[Int]] = + source.groupChunksWithin(Int.MaxValue, 1.day) + + downstream.compile.lastOrError + .timeout(downstreamTimeout) + .map(_.toList) + .timed + } + .assertEquals((sourceTimeout, List(0, 1, 2, 3, 4, 5))) + } + + test("stress test: all elements are processed") { + val rangeLength = 10000 + + TestControl.executeEmbed { + Stream + .eval(Ref.of[IO, Int](0)) + .flatMap { counter => + Stream + .range(0, rangeLength) + .covary[IO] + .groupChunksWithin(4096, 100.micros) + .evalMap(ch => counter.updateAndGet(_ + ch.size)) + } + .compile + .lastOrError + .assertEquals(rangeLength) + } + } + } + property("head")(forAll((s: Stream[Pure, Int]) => assertEquals(s.head.toList, s.toList.take(1)))) group("ifEmpty") { diff --git a/core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala b/core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala new file mode 100644 index 0000000000..87238b36e9 --- /dev/null +++ b/core/shared/src/test/scala/fs2/concurrent/ConditionedRefSuite.scala @@ -0,0 +1,88 @@ +/* + * Copyright (c) 2013 Functional Streams for Scala + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER + * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN + * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ + +package fs2 +package concurrent + +import cats.effect.IO +import cats.effect.testkit.TestControl +import cats.syntax.all._ + +import scala.concurrent.duration._ + +class ConditionedRefSuite extends Fs2Suite { + + test("waitUntil completes immediately if the condition already holds") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 1).flatMap { ref => + ref.waitUntil(_ > 0).timed.map { case (elapsed, _) => assertEquals(elapsed, Duration.Zero) } + } + } + } + + test("waitUntil completes on the update that makes the condition true") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + val updates = + IO.sleep(1.second) >> ref.update(_ + 1) >> IO.sleep(1.second) >> ref.update(_ + 1) + updates.background.surround { + ref.waitUntil(_ >= 2).timed.map { case (elapsed, _) => assertEquals(elapsed, 2.seconds) } + } + } + } + } + + test("waitUntil does not miss a condition that becomes true and then false again") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + ref.waitUntil(_ == 1).start.flatMap { waiter => + IO.sleep(1.second) >> ref.set(1) >> ref.set(2) >> + waiter.joinWithNever.timeout(1.second) >> ref.get.assertEquals(2) + } + } + } + } + + test("only waiters whose condition holds are woken") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + (ref.waitUntil(_ >= 1).start, ref.waitUntil(_ >= 2).start).flatMapN { (first, second) => + IO.sleep(1.second) >> ref.set(1) >> + first.joinWithNever.timeout(1.second) >> + IO.race(second.joinWithNever, IO.sleep(1.second)).map(r => assert(r.isRight)) >> + ref.set(2) >> + second.joinWithNever.timeout(1.second) + } + } + } + } + + test("a cancelled waiter does not affect later updates and waiters") { + TestControl.executeEmbed { + ConditionedRef.of[IO, Int](initial = 0).flatMap { ref => + ref.waitUntil(_ > 0).timeoutTo(duration = 1.second, fallback = IO.unit) >> + ref.update(_ + 1) >> + ref.waitUntil(_ > 0).timeout(1.second) >> + ref.get.assertEquals(1) + } + } + } +} diff --git a/integration/src/test/scala/fs2/MemoryLeakSpec.scala b/integration/src/test/scala/fs2/MemoryLeakSpec.scala index 4e2fa022f3..71afabe646 100644 --- a/integration/src/test/scala/fs2/MemoryLeakSpec.scala +++ b/integration/src/test/scala/fs2/MemoryLeakSpec.scala @@ -153,6 +153,27 @@ class MemoryLeakSpec extends FunSuite { .groupWithin(256, 1.second) } + leakTest("groupChunksWithin") { + Stream + .eval(IO.never) + .covary[IO] + .groupChunksWithin(Int.MaxValue, 1.millis) + } + + leakTest("groupChunksWithin --- Issue 2328") { + Stream + .range(0, 1000000) + .covary[IO] + .groupChunksWithin(256, 1.second) + } + + leakTest("groupChunksWithin --- repeated idle timeouts") { + Stream + .awakeEvery[IO](1.second) + .as(1) + .groupChunksWithin(256, 1.millis) + } + leakTest("topic continuous publish") { Stream .eval(Topic[IO, Int]) From 30aef1a7d90abf1496731f7c3085a5cb69779e49 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 12:32:28 -0400 Subject: [PATCH 07/11] Chunk.grouped returns Chunk[Chunk[O]] --- core/shared/src/main/scala/fs2/Chunk.scala | 22 ++++++++++++------- core/shared/src/main/scala/fs2/Stream.scala | 2 +- .../src/test/scala/fs2/ChunkSuite.scala | 9 ++++---- 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/core/shared/src/main/scala/fs2/Chunk.scala b/core/shared/src/main/scala/fs2/Chunk.scala index c99d56f47d..4c781d711f 100644 --- a/core/shared/src/main/scala/fs2/Chunk.scala +++ b/core/shared/src/main/scala/fs2/Chunk.scala @@ -211,18 +211,24 @@ abstract class Chunk[+O] extends Serializable with ChunkPlatform[O] with ChunkRu * * Like `List#grouped`, an empty chunk has no groups, and `n` must be positive. */ - def grouped(n: Int): Vector[Chunk[O]] = { + def grouped(n: Int): Chunk[Chunk[O]] = { require(n > 0, s"n must be positive, but got ${n.toString}") - @tailrec - def loop(rest: Chunk[O], acc: Vector[Chunk[O]]): Vector[Chunk[O]] = - if (rest.size <= n) acc :+ rest - else { - val (group, remaining) = rest.splitAt(n) - loop(rest = remaining, acc = acc :+ group) + if (isEmpty) Chunk.empty + else { + val numGroups = (size - 1) / n + 1 // ceil divide + val groups = new Array[Chunk[O]](numGroups) + var i = 0 + var chunk = this + while (i < numGroups) { + val (group, remaining) = chunk.splitAt(n) + groups(i) = group + chunk = remaining + i += 1 } - if (isEmpty) Vector.empty else loop(rest = this, acc = Vector.empty) + Chunk.array(groups) + } } /** Gets the first element of this chunk. */ diff --git a/core/shared/src/main/scala/fs2/Stream.scala b/core/shared/src/main/scala/fs2/Stream.scala index 468cece318..3fd6af3fb7 100644 --- a/core/shared/src/main/scala/fs2/Stream.scala +++ b/core/shared/src/main/scala/fs2/Stream.scala @@ -1695,7 +1695,7 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, def emitBatches: Pull[F2, Chunk[O], Unit] = Pull.eval(nextBatch).flatMap { - case Some(batch) => Pull.output(Chunk.from(batch.grouped(chunkSize))) >> emitBatches + case Some(batch) => Pull.output(batch.grouped(chunkSize)) >> emitBatches case None => Pull.done } diff --git a/core/shared/src/test/scala/fs2/ChunkSuite.scala b/core/shared/src/test/scala/fs2/ChunkSuite.scala index 78ef3b918b..3f59a45fb2 100644 --- a/core/shared/src/test/scala/fs2/ChunkSuite.scala +++ b/core/shared/src/test/scala/fs2/ChunkSuite.scala @@ -131,12 +131,11 @@ class ChunkSuite extends Fs2Suite { } } - test("Chunk.grouped of an empty chunk has no groups, like List#grouped") { - assertEquals(Chunk.empty[Int].grouped(3), Vector.empty) - assertEquals(List.empty[Int].grouped(3).toList, Nil) + test("Chunk.grouped of an empty chunk has no groups") { + assertEquals(Chunk.empty[Int].grouped(3), Chunk.empty) } - test("Chunk.grouped rejects a non-positive group size, like List#grouped") { + test("Chunk.grouped rejects a non-positive group size") { List(0, -1).foreach { n => intercept[IllegalArgumentException](Chunk(1, 2, 3).grouped(n)) intercept[IllegalArgumentException](Chunk.empty[Int].grouped(n)) @@ -146,7 +145,7 @@ class ChunkSuite extends Fs2Suite { test("Chunk.grouped returns a chunk that fits in one group as is") { forAll { (c: Chunk[Int]) => - if (c.nonEmpty) assert(c.grouped(c.size).head eq c) + if (c.nonEmpty) assert(c.grouped(c.size).head.get eq c) } } From 96508a39fa6cb16a2d45245657225cb3b9d3fd40 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 12:34:20 -0400 Subject: [PATCH 08/11] Use existing ref method --- .../main/scala/fs2/concurrent/ConditionedRef.scala | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala index e558ef01db..092d5aac94 100644 --- a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala +++ b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala @@ -95,14 +95,12 @@ private[fs2] object ConditionedRef { else F.deferred[Unit].flatMap { wake => val waiter = new Waiter(accepts = p, wake = wake) - F.uncancelable { poll => - state.modify { s => - if (p(s.value)) s -> F.unit - else - s.register(waiter) -> poll(wake.get).onCancel { - state.update(_.deregister(waiter)) - } - }.flatten + state.flatModifyFull { (poll, s) => + if (p(s.value)) s -> F.unit + else + s.register(waiter) -> poll(wake.get).onCancel { + state.update(_.deregister(waiter)) + } } } } From 6fd83218f91060f8cf5d24f2b1de15ecefe7b653 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 12:55:20 -0400 Subject: [PATCH 09/11] Try to justfy impl choice --- core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala | 2 ++ 1 file changed, 2 insertions(+) diff --git a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala index 092d5aac94..ffdad7bef6 100644 --- a/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala +++ b/core/shared/src/main/scala/fs2/concurrent/ConditionedRef.scala @@ -79,6 +79,8 @@ private[fs2] object ConditionedRef { def modify[B](f: A => (A, B)): F[B] = state.flatModify { s => // uncancellable to avoid losing wake-up signals val (value, result) = f(s.value) + // Scan first instead of calling partition to avoid extra allocations which have a small impact on perf. + // This optimizes for the case where writes infrequently trigger wake-ups. if (!s.waiters.exists(_.accepts(value))) State(value = value, waiters = s.waiters) -> result.pure[F] else { From e807c31d125ce7798c35c0bfcebf9ff182435544 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Mon, 5 Oct 2026 14:28:58 -0400 Subject: [PATCH 10/11] Use chunk-based impl for groupWithin; align chunking behavior around timeouts to existing api --- .../GroupChunksWithinBenchmark.scala | 64 ----- core/shared/src/main/scala/fs2/Stream.scala | 124 +------- .../scala/fs2/StreamCombinatorsSuite.scala | 265 +++--------------- .../src/test/scala/fs2/MemoryLeakSpec.scala | 16 +- 4 files changed, 44 insertions(+), 425 deletions(-) delete mode 100644 benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala diff --git a/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala b/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala deleted file mode 100644 index 9a008f5f47..0000000000 --- a/benchmark/src/main/scala/fs2/benchmark/GroupChunksWithinBenchmark.scala +++ /dev/null @@ -1,64 +0,0 @@ -/* - * Copyright (c) 2013 Functional Streams for Scala - * - * Permission is hereby granted, free of charge, to any person obtaining a copy of - * this software and associated documentation files (the "Software"), to deal in - * the Software without restriction, including without limitation the rights to - * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of - * the Software, and to permit persons to whom the Software is furnished to do so, - * subject to the following conditions: - * - * The above copyright notice and this permission notice shall be included in all - * copies or substantial portions of the Software. - * - * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS - * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR - * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER - * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN - * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - */ - -package fs2 -package benchmark - -import cats.effect.IO -import org.openjdk.jmh.annotations.{Benchmark, Param, Scope, State} - -import scala.concurrent.duration.* - -@State(Scope.Thread) -class GroupChunksWithinBenchmark { - - import cats.effect.unsafe.implicits.global - - val bufferWindow = 100.micros - - @Param(Array("100", "10000", "100000")) - var rangeLength: Int = _ - - @Param(Array("16", "256", "4096")) - var bufferSize: Int = _ - - @Benchmark - def groupChunksWithin(): Unit = - Stream - .range(0, rangeLength) - .covary[IO] - .groupChunksWithin(bufferSize, bufferWindow) - .compile - .drain - .unsafeRunSync() - - @Benchmark - def groupChunksWithinChunkedUpstream(): Unit = - Stream - .range(0, rangeLength) - .chunkN(bufferSize / 4 + 1) - .unchunks - .covary[IO] - .groupChunksWithin(bufferSize, bufferWindow) - .compile - .drain - .unsafeRunSync() -} diff --git a/core/shared/src/main/scala/fs2/Stream.scala b/core/shared/src/main/scala/fs2/Stream.scala index 3fd6af3fb7..aefd629408 100644 --- a/core/shared/src/main/scala/fs2/Stream.scala +++ b/core/shared/src/main/scala/fs2/Stream.scala @@ -1520,116 +1520,6 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, def groupWithin[F2[x] >: F[x]]( chunkSize: Int, timeout: FiniteDuration - )(implicit F: Temporal[F2]): Stream[F2, Chunk[O]] = { - - case class JunctionBuffer[T]( - data: Vector[T], - endOfSupply: Option[Either[Throwable, Unit]], - endOfDemand: Option[Either[Throwable, Unit]] - ) { - def splitAt(n: Int): (JunctionBuffer[T], JunctionBuffer[T]) = - if (this.data.size >= n) { - val (head, tail) = this.data.splitAt(n.toInt) - (this.copy(tail), this.copy(head)) - } else { - (this.copy(Vector.empty), this) - } - } - - val outputLong = chunkSize.toLong - fs2.Stream.force { - for { - demand <- Semaphore[F2](outputLong) - supply <- Semaphore[F2](0L) - buffer <- Ref[F2].of( - JunctionBuffer[O](Vector.empty[O], endOfSupply = None, endOfDemand = None) - ) - } yield { - /* - Buffer: stores items from input to be sent on next output chunk - * - Demand Semaphore: to avoid adding too many items to buffer - * - Supply: counts filled positions for next output chunk */ - def enqueue(t: O): F2[Boolean] = - for { - _ <- demand.acquire - buf <- buffer.modify(buf => (buf.copy(buf.data :+ t), buf)) - _ <- supply.release - } yield buf.endOfDemand.isEmpty - - val dequeueNextOutput: F2[Option[Vector[O]]] = { - // Trigger: waits until the supply buffer is full (with acquireN) - val waitSupply = supply.acquireN(outputLong).guaranteeCase { - case Outcome.Succeeded(_) => supply.releaseN(outputLong) - case _ => F.unit - } - - val onTimeout: F2[Long] = - for { - _ <- supply.acquire // waits until there is at least one element in buffer - m <- supply.available - k = m.min(outputLong - 1) - b <- supply.tryAcquireN(k) - } yield if (b) k + 1 else 1 - - // in JS cancellation doesn't always seem to run, so race conditions should restore state on their own - for { - acq <- F.race(F.sleep(timeout), waitSupply).flatMap { - case Left(_) => onTimeout - case Right(_) => supply.acquireN(outputLong).as(outputLong) - } - buf <- buffer.modify(_.splitAt(acq.toInt)) - _ <- demand.releaseN(buf.data.size.toLong) - res <- buf.endOfSupply match { - case Some(Left(error)) => F.raiseError(error) - case Some(Right(_)) if buf.data.isEmpty => F.pure(None) - case _ => F.pure(Some(buf.data)) - } - } yield res - } - - def endSupply(result: Either[Throwable, Unit]): F2[Unit] = - buffer.update(_.copy(endOfSupply = Some(result))) *> supply.releaseN( - // enough supply for 2 iterations of the race loop in case of upstream - // interruption: so that downstream can terminate immediately - outputLong * 2 - ) - - def endDemand(result: Either[Throwable, Unit]): F2[Unit] = - buffer.update(_.copy(endOfDemand = Some(result))) *> demand.releaseN(Int.MaxValue) - - def toEnding(ec: ExitCase): Either[Throwable, Unit] = ec match { - case ExitCase.Succeeded => Right(()) - case ExitCase.Errored(e) => Left(e) - case ExitCase.Canceled => Right(()) - } - - val enqueueAsync = F.start { - this - .evalMap(enqueue) - .forall(identity) - .onFinalizeCase(ec => endSupply(toEnding(ec))) - .compile - .drain - } - - val outputStream: Stream[F2, Chunk[O]] = - Stream - .eval(dequeueNextOutput) - .repeat - .collectWhile { case Some(data) => Chunk.from(data) } - - Stream - .bracketCase(enqueueAsync) { case (upstream, exitCase) => - endDemand(toEnding(exitCase)) *> upstream.cancel - } >> outputStream - } - } - } - - /** Like [[groupWithin]], but works on chunks for better performance. - */ - def groupChunksWithin[F2[x] >: F[x]]( - chunkSize: Int, - timeout: FiniteDuration )(implicit F: Temporal[F2]): Stream[F2, Chunk[O]] = Stream.force { require(chunkSize > 0, s"chunkSize must be > 0, but got ${chunkSize.toString}") @@ -1662,11 +1552,11 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, } // None means upstream is done, and the buffer is drained. - def takeOrExit(all: Boolean): F2[Option[Chunk[O]]] = + def takeOrExit(timedOut: Boolean): F2[Option[Chunk[O]]] = buffer.modify { b => - if (b.nonEmpty && (all || b.isDone)) take(b, n = b.size) - // Take all potential batches so that producer doesn't need to wait for timeout + if (b.nonEmpty && b.isDone) take(b, n = b.size) else if (b.isFull) take(b, n = b.size - b.size % chunkSize) + else if (b.nonEmpty && timedOut) take(b, n = b.size) else b -> (b.done match { case None /* not full & no-timeout */ => Skip @@ -1675,20 +1565,20 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, }) }.flatten - val onTimeout = takeOrExit(all = true).flatMap { + val onTimeout = takeOrExit(timedOut = true).flatMap { case Some(batch) if batch.isEmpty => - buffer.waitUntil(b => b.nonEmpty || b.isDone) >> takeOrExit(all = true) + buffer.waitUntil(b => b.nonEmpty || b.isDone) >> takeOrExit(timedOut = true) case result => result.pure[F2] } val nextBatch: F2[Option[Chunk[O]]] = // Potentially skip starting timer fiber if buffer is full - takeOrExit(all = false).flatMap { + takeOrExit(timedOut = false).flatMap { case Some(batch) if batch.isEmpty => F.race(F.sleep(timeout), buffer.waitUntil(b => b.isFull || b.isDone)) .flatMap { case Left(_ /* timeout */ ) => onTimeout - case Right(_ /* full batch or done */ ) => takeOrExit(all = false) + case Right(_ /* full batch or done */ ) => takeOrExit(timedOut = false) } case result => result.pure[F2] } diff --git a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala index b0ea3ade29..a7cae53751 100644 --- a/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala +++ b/core/shared/src/test/scala/fs2/StreamCombinatorsSuite.scala @@ -763,7 +763,7 @@ class StreamCombinatorsSuite extends Fs2Suite { s.covary[IO] .evalTap(sleep) .groupWithin(groupSize, timeout) - .flatMap(Stream.chunk) + .unchunks .assertEmitsSameAs(s) } } @@ -856,6 +856,37 @@ class StreamCombinatorsSuite extends Fs2Suite { } } + test("accumulation and splitting with timeouts between bursts") { + val t = 200.millis + val chunkSize = 5 + + def chunk(from: Int, to: Int) = + Stream.range(from, to + 1).chunkAll.unchunks + + val source = + chunk(from = 1, to = 3) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 4, to = 15) ++ + Stream.sleep_[IO](t + t / 2) ++ + chunk(from = 16, to = 22) + + val expected = List( + (1 to 3).toList, + (4 to 8).toList, + (9 to 13).toList, + (14 to 15).toList, + (16 to 20).toList, + (21 to 22).toList + ) + + TestControl.executeEmbed { + source + .groupWithin(chunkSize, t) + .map(_.toList) + .assertEmits(expected) + } + } + test("does not reset timeout if nothing is emitted") { TestControl .executeEmbed( @@ -985,44 +1016,13 @@ class StreamCombinatorsSuite extends Fs2Suite { .assertEquals(rangeLength) } } - } - - group("groupChunksWithin") { - implicit val groupSizeArb: Arbitrary[Int] = Arbitrary(Gen.choose(1, 20)) - // a zero timeout is covered separately. Under TestControl it makes the idle loop spin forever - implicit val timeoutArb: Arbitrary[FiniteDuration] = Arbitrary(Gen.choose(1, 50).map(_.millis)) - - def sleep(d: Int): IO[Unit] = IO.sleep((d % 500).abs.micros) - - test("should never lose any elements") { - forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration) => - TestControl.executeEmbed { - s.covary[IO] - .evalTap(sleep) - .groupChunksWithin(chunkSize = 3, timeout) - .unchunks - .assertEmitsSameAs(s) - } - } - } - - test("should never emit empty chunks") { - forAllF { (s: Stream[Pure, Int], timeout: FiniteDuration, groupSize: Int) => - TestControl.executeEmbed { - s.covary[IO] - .evalTap(sleep) - .groupChunksWithin(groupSize, timeout) - .assertForall(_.nonEmpty) - } - } - } test("makes progress with a zero timeout") { Stream .range(0, 100) .covary[IO] .evalTap(_ => IO.sleep(1.milli)) - .groupChunksWithin(chunkSize = 10, timeout = Duration.Zero) + .groupWithin(chunkSize = 10, timeout = Duration.Zero) .unchunks .compile .toList @@ -1030,30 +1030,12 @@ class StreamCombinatorsSuite extends Fs2Suite { .assertEquals((0 until 100).toList) } - test("every chunk but the last has at least chunkSize elements when no timeout triggers") { - val source = Stream.range(0, 101) - val chunkSize = 5 - - TestControl.executeEmbed { - source - .covary[IO] - .groupChunksWithin(chunkSize, 1.day) - .compile - .toList - .map { chunks => - assertEquals(chunks.flatMap(_.toList), source.toList) - assert(chunks.init.forall(_.size >= chunkSize)) - assert(chunks.lastOption.forall(_.size < chunkSize)) - } - } - } - test("splits upstream chunks larger than chunkSize") { TestControl.executeEmbed { Stream .emits(1 to 12) .covary[IO] - .groupChunksWithin(chunkSize = 5, timeout = 1.day) + .groupWithin(chunkSize = 5, timeout = 1.day) .map(_.toList) .assertEmits(List((1 to 5).toList, (6 to 10).toList, (11 to 12).toList)) } @@ -1062,7 +1044,7 @@ class StreamCombinatorsSuite extends Fs2Suite { test("emits every full group of a large upstream chunk without waiting for the timeout") { TestControl.executeEmbed { (Stream.emits(1 to 100) ++ Stream.never[IO]) - .groupChunksWithin(chunkSize = 10, timeout = 1.day) + .groupWithin(chunkSize = 10, timeout = 1.day) .take(10) .compile .toList @@ -1074,55 +1056,6 @@ class StreamCombinatorsSuite extends Fs2Suite { } } - test( - "should return a finite stream back in a single chunk given a chunk size equal to the stream size and an absurdly high duration" - ) { - forAllF { (streamAsList0: List[Int]) => - val streamAsList = 0 :: streamAsList0 - TestControl.executeEmbed { - Stream - .emits(streamAsList) - .covary[IO] - .groupChunksWithin(streamAsList.size, (Int.MaxValue - 1L).nanoseconds) - .compile - .toList - .map(_.head.toList) - .assertEquals(streamAsList) - } - } - } - - test("accumulation with splitting") { - val t = 200.millis - val chunkSize = 5 - - def chunk(from: Int, to: Int) = - Stream.range(from, to + 1).chunkAll.unchunks - - val source = - chunk(from = 1, to = 3) ++ - Stream.sleep_[IO](t + t / 2) ++ - chunk(from = 4, to = 15) ++ - Stream.sleep_[IO](t + t / 2) ++ - chunk(from = 16, to = 22) - - val expected = List( - (1 to 3).toList, - (4 to 8).toList, - (9 to 13).toList, - (14 to 15).toList, - (16 to 20).toList, - (21 to 22).toList - ) - - TestControl.executeEmbed { - source - .groupChunksWithin(chunkSize, t) - .map(_.toList) - .assertEmits(expected) - } - } - test("emits a burst arriving after several idle timeouts") { val timeout = 100.millis val chunkSize = 5 @@ -1135,137 +1068,11 @@ class StreamCombinatorsSuite extends Fs2Suite { TestControl.executeEmbed { source - .groupChunksWithin(chunkSize, timeout) + .groupWithin(chunkSize, timeout) .map(_.toList) .assertEmits(List(List(1, 2, 3), List(4, 5, 6))) } } - - test("does not reset timeout if nothing is emitted") { - TestControl - .executeEmbed( - Ref[IO] - .of(0.millis) - .flatMap { ref => - val timeout = 5.seconds - - def measureEmission[A]: Pipe[IO, A, A] = - _.chunks - .evalTap(_ => IO.monotonic.flatMap(ref.set)) - .unchunks - - // emits elements after the timeout has expired - val source = - Stream.sleep_[IO](timeout + 200.millis) ++ - Stream(4, 5) ++ - Stream.never[IO] // avoids emission due to source termination - - source - .through(measureEmission) - .groupChunksWithin(5, timeout) - .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) - .interruptAfter(timeout * 3) - .compile - .lastOrError - } - ) - .assertEquals(0.millis) // The stream emits after the timeout has expired - // on an empty buffer, so groupChunksWithin should re-emit with zero delay - // rather than start a fresh timeout window - } - - test("Edge case: should not introduce unnecessary delays when chunkSize == chunk size") { - TestControl - .executeEmbed( - Ref[IO] - .of(0.millis) - .flatMap { ref => - val timeout = 5.seconds - - def measureEmission[A]: Pipe[IO, A, A] = - _.chunks - .evalTap(_ => IO.monotonic.flatMap(ref.set)) - .unchunks - - val source = - Stream(1, 2, 3) ++ - Stream.sleep_[IO](timeout + 200.millis) - - source - .through(measureEmission) - .groupChunksWithin(3, timeout) - .evalMap(_ => (IO.monotonic, ref.get).mapN(_ - _)) - .compile - .lastOrError - } - ) - .assertEquals(0.millis) - } - - test("upstream failures are propagated downstream") { - TestControl.executeEmbed { - case object SevenNotAllowed extends NoStackTrace - - val source = Stream - .iterate(0)(_ + 1) - .covary[IO] - .evalTap(n => IO.raiseError(SevenNotAllowed).whenA(n == 7)) - - val downstream = source.groupChunksWithin(100, 2.seconds).map(_.toList) - - val expected = List((0 to 6).toList) - - downstream.assertEmits(expected).intercept[SevenNotAllowed.type] - } - } - - test( - "upstream interruption causes immediate downstream termination with all elements being emitted" - ) { - val sourceTimeout = 5.5.seconds - val downstreamTimeout = sourceTimeout + 2.seconds - - TestControl - .executeEmbed { - val source: Stream[IO, Int] = - Stream - .iterate(0)(_ + 1) - .covary[IO] - .meteredStartImmediately(1.second) - .interruptAfter(sourceTimeout) - - // large chunkSize and timeout: no emissions expected in the window - // specified, unless the source ends, due to interruption or natural - // termination (i.e. runs out of elements) - val downstream: Stream[IO, Chunk[Int]] = - source.groupChunksWithin(Int.MaxValue, 1.day) - - downstream.compile.lastOrError - .timeout(downstreamTimeout) - .map(_.toList) - .timed - } - .assertEquals((sourceTimeout, List(0, 1, 2, 3, 4, 5))) - } - - test("stress test: all elements are processed") { - val rangeLength = 10000 - - TestControl.executeEmbed { - Stream - .eval(Ref.of[IO, Int](0)) - .flatMap { counter => - Stream - .range(0, rangeLength) - .covary[IO] - .groupChunksWithin(4096, 100.micros) - .evalMap(ch => counter.updateAndGet(_ + ch.size)) - } - .compile - .lastOrError - .assertEquals(rangeLength) - } - } } property("head")(forAll((s: Stream[Pure, Int]) => assertEquals(s.head.toList, s.toList.take(1)))) diff --git a/integration/src/test/scala/fs2/MemoryLeakSpec.scala b/integration/src/test/scala/fs2/MemoryLeakSpec.scala index 71afabe646..4e00614322 100644 --- a/integration/src/test/scala/fs2/MemoryLeakSpec.scala +++ b/integration/src/test/scala/fs2/MemoryLeakSpec.scala @@ -153,25 +153,11 @@ class MemoryLeakSpec extends FunSuite { .groupWithin(256, 1.second) } - leakTest("groupChunksWithin") { - Stream - .eval(IO.never) - .covary[IO] - .groupChunksWithin(Int.MaxValue, 1.millis) - } - - leakTest("groupChunksWithin --- Issue 2328") { - Stream - .range(0, 1000000) - .covary[IO] - .groupChunksWithin(256, 1.second) - } - leakTest("groupChunksWithin --- repeated idle timeouts") { Stream .awakeEvery[IO](1.second) .as(1) - .groupChunksWithin(256, 1.millis) + .groupWithin(256, 1.millis) } leakTest("topic continuous publish") { From 05f824e9412cc715ff4813097fd55c9d3f774378 Mon Sep 17 00:00:00 2001 From: Michael Mienko Date: Wed, 7 Oct 2026 05:12:18 -0400 Subject: [PATCH 11/11] Update method docs --- core/shared/src/main/scala/fs2/Stream.scala | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/core/shared/src/main/scala/fs2/Stream.scala b/core/shared/src/main/scala/fs2/Stream.scala index aefd629408..39b1e7fd8c 100644 --- a/core/shared/src/main/scala/fs2/Stream.scala +++ b/core/shared/src/main/scala/fs2/Stream.scala @@ -1487,13 +1487,16 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F, go(None, this).stream } - /** Splits this stream into a stream of chunks of elements, such that + /** Splits this stream into a stream of chunks of elements, such that: + * * 1. each chunk in the output has at most `outputSize` elements, and + * * 2. the concatenation of those chunks, which is obtained by calling * `unchunks`, yields the same element sequence as this stream. * - * As `this` stream emits input elements, the result stream them in a - * waiting buffer, until it has enough elements to emit next chunk. + * As `this` stream emits input elements, the result stream accumulates + * them in a waiting buffer, until it has enough elements to emit the + * next chunk. Accumulation acts on chunks of input elements for efficiency. * * To avoid holding input elements for too long, this method takes a * `timeout`. This timeout is reset after each output chunk is emitted.