merge: pull from upstream only when downstream has demand

core/shared/src/main/scala/fs2/Stream.scala

@@ -2001,16 +2001,41 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F,
     Stream.force(fstream)
   }
 
-  /** Implementation of [[merge]], however allows specifying how to combine the output stream.
-    * This can be used to control how chunks are emitted downstream. See [[mergeAndAwaitDownstream]] for example.
+  /** Interleaves the two inputs nondeterministically. The output stream
+    * halts after BOTH `s1` and `s2` terminate normally, or in the event
+    * of an uncaught failure on either `s1` or `s2`. Has the property that
+    * `merge(Stream.empty, s) == s` and `merge(raiseError(e), s)` will
+    * eventually terminate with `raiseError(e)`, possibly after emitting some
+    * elements of `s` first.
+    *
+    * Neither side pulls a further chunk from its upstream until the resulting
+    * stream is asked for the next element, so `merge` does not read ahead of
+    * downstream demand any more than it must in order to race the two sides.
+    * Concretely, the implementation always tries to pull one chunk from each
+    * side before waiting for one of them to be consumed, so there may be up to
+    * two chunks (one from each stream) waiting to be processed while the
+    * resulting stream is processing elements. Use [[prefetch]] on either side
+    * if you do want a chunk fetched ahead of demand.
     *
-    * @param f The function that combines the output stream and a finalizer for the chunk.
-    *          This way we can controll when to pull pull next chunk from upstream.
+    * Also note that if either side produces empty chunk,
+    * the processing on that side continues,
+    * w/o downstream requiring to consume result.
+    *
+    * If either side does not emit anything (i.e. as result of drain) that side
+    * will continue to run even when the resulting stream did not ask for more
+    * data. If a side needs to make progress independently of downstream demand,
+    * consider [[concurrently]] instead.
+    *
+    * @example {{{
+    * scala> import scala.concurrent.duration._, cats.effect.IO, cats.effect.unsafe.implicits.global
+    * scala> val s1 = Stream.awakeEvery[IO](500.millis).scan(0)((acc, _) => acc + 1)
+    * scala> val s = s1.merge(Stream.sleep_[IO](250.millis) ++ s1)
+    * scala> s.take(6).compile.toVector.unsafeRunSync()
+    * res0: Vector[Int] = Vector(0, 0, 1, 1, 2, 2)
+    * }}}
     */
-  private def merge_[F2[x] >: F[x], O2 >: O](
+  def merge[F2[x] >: F[x], O2 >: O](
       that: Stream[F2, O2]
-  )(
-      f: (Stream[F2, O2], F2[Unit]) => Stream[F2, O2]
   )(implicit F: Concurrent[F2]): Stream[F2, O2] =
     Stream.force {
       // `State` describes the state of an upstream stream (`this` and `that` are both upstream streams)
@@ -2041,10 +2066,11 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F,
           case (Some(r1), Some(r2)) => CompositeFailure.fromResults(r1, r2)
         }
         def run(s: Stream[F2, O2]): F2[Unit] =
-          // `guard` ensures we do not pull another chunk until the previous one has been produced for downstream.
+          // `guard` ensures we do not pull another chunk until the previous one
+          // has been consumed by the downstream.
           Semaphore[F2](1).flatMap { guard =>
             def sendChunk(chk: Chunk[O2]): F2[Unit] =
-              output.send(f(Stream.chunk(chk), guard.release)) >> guard.acquire
+              output.send(Stream.chunk(chk) ++ Stream.exec(guard.release)) >> guard.acquire
 
             (Stream.exec(guard.acquire) ++ s.chunks.foreach(sendChunk))
               // Stop when the other upstream has errored or the downstream has completed.
@@ -2079,64 +2105,16 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F,
       }
     }
 
-  /** Like [[merge]], but ensures that each chunk is fully consumed downstream before pulling the next chunk from the same side.
-    * This looses the equivalence with `Stream(this, that).parJoinUnbounded` but can be useful when we need to never read ahead from
-    * the merged streams.
-    *
-    * @note Pay attention to possible deadlocks of "this" or "that" when using this function, notably in parallel processing
-    *       as unless the chunk is fully processed / scope of the chunk is released, the next chunk will not be pulled.
+  /** Alias for [[merge]], retained for source compatibility.
     *
-    * @example {{{
-    * scala> import scala.concurrent.duration._, cats.effect.IO, cats.effect.unsafe.implicits.global
-    * scala> import cats.effect._
-    * scala> Ref.of[IO, Int](0).flatMap{ ref =>
-    *      |   fs2.Stream.never[IO].mergeAndAwaitDownstream(fs2.Stream.repeatEval(ref.get)).evalMap(value => {
-    *      |     IO.sleep(1.second) >> ref.set(value + 1) as value
-    *      |   }).take(6).compile.toVector
-    *      | }.unsafeRunSync()
-    * res0: Vector[Int] = Vector(0, 1, 2, 3, 4, 5)
-    * }}}
+    * This was introduced as a variant of `merge` that never reads ahead of
+    * downstream demand. `merge` now has that property itself, so the two are
+    * equivalent.
     */
   def mergeAndAwaitDownstream[F2[x] >: F[x], O2 >: O](
       that: Stream[F2, O2]
   )(implicit F: Concurrent[F2]): Stream[F2, O2] =
-    merge_(that) { case (s, fin) => s.onFinalize(fin) }
-
-  /** Interleaves the two inputs nondeterministically. The output stream
-    * halts after BOTH `s1` and `s2` terminate normally, or in the event
-    * of an uncaught failure on either `s1` or `s2`. Has the property that
-    * `merge(Stream.empty, s) == s` and `merge(raiseError(e), s)` will
-    * eventually terminate with `raiseError(e)`, possibly after emitting some
-    * elements of `s` first.
-    *
-    * The implementation always tries to pull one chunk from each side
-    * before waiting for it to be consumed by resulting stream.
-    * As such, there may be up to two chunks (one from each stream)
-    * waiting to be processed while the resulting stream
-    * is processing elements.
-    *
-    * Also note that if either side produces empty chunk,
-    * the processing on that side continues,
-    * w/o downstream requiring to consume result.
-    *
-    * If either side does not emit anything (i.e. as result of drain) that side
-    * will continue to run even when the resulting stream did not ask for more data.
-    *
-    * Note that even when this is equivalent to `Stream(this, that).parJoinUnbounded`,
-    * this implementation is little more efficient
-    *
-    * @example {{{
-    * scala> import scala.concurrent.duration._, cats.effect.IO, cats.effect.unsafe.implicits.global
-    * scala> val s1 = Stream.awakeEvery[IO](500.millis).scan(0)((acc, _) => acc + 1)
-    * scala> val s = s1.merge(Stream.sleep_[IO](250.millis) ++ s1)
-    * scala> s.take(6).compile.toVector.unsafeRunSync()
-    * res0: Vector[Int] = Vector(0, 0, 1, 1, 2, 2)
-    * }}}
-    */
-  def merge[F2[x] >: F[x], O2 >: O](
-      that: Stream[F2, O2]
-  )(implicit F: Concurrent[F2]): Stream[F2, O2] =
-    merge_(that) { case (s, fin) => Stream.exec(fin) ++ s }
+    merge(that)
 
   /** Like `merge`, but halts as soon as _either_ branch halts. */
   def mergeHaltBoth[F2[x] >: F[x]: Concurrent, O2 >: O](

core/shared/src/test/scala/fs2/StreamMergeSuite.scala

@@ -224,7 +224,7 @@ class StreamMergeSuite extends Fs2Suite {
     }
   }
 
-  test("merge not emit ahead more than 1 chunk") {
+  test("merge not emit ahead") {
     forAllF { (v: Int) =>
       Ref
         .of[IO, Int](v)
@@ -236,8 +236,8 @@ class StreamMergeSuite extends Fs2Suite {
             .repeatEval(ref.get)
             .merge(Stream.never[IO])
             .evalMap(sleepAndSet)
-            .take(6)
-            .assertEmits(List(v, v, v + 1, v + 1, v + 2, v + 2))
+            .take(3)
+            .assertEmits(List(v, v + 1, v + 2))
         }
     }
   }

core/shared/src/test/scala/fs2/TimedPullsSuite.scala

@@ -313,19 +313,8 @@ class TimedPullsSuite extends Fs2Suite {
   }
 
   test("After the first uncons, timeouts start immediately") {
-    // Time how often we generate data in the main stream.
-    // This is only started after the first uncons.
     val emissionTime = 100.millis
-
-    // Timeout which is registered before the first uncons, it is registered immediately
-    // But we do not expect it to trigger.
-    // This has to be longer than emissionTime, otherwise the first uncons would always timeout.
-    val initialTimeout = 200.millis
-
-    // Timeout registered after the first uncons, this one should be fired
-    val timeout = 50.millis
-
-    // Time we wait before doing uncons.
+    val timeout = 200.millis
     val timedPullPause = Pull.eval(IO.sleep(150.millis))
 
     val prog =
@@ -334,7 +323,7 @@ class TimedPullsSuite extends Fs2Suite {
         .repeatN(2)
         .pull
         .timed { tp =>
-          tp.timeout(initialTimeout) >>
+          tp.timeout(timeout) >>
             // If the first timeout started immediately, this pause
             // before uncons would cause a timeout to be emitted
             timedPullPause >>