interrupt behaviour change, test improvements

Author: Angel-O

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

@@ -1525,26 +1525,25 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F,
             buffer.tryTakeN(Some(groupSize.toInt)).map(Chunk.seq)
 
           // we need to check the buffer size, rather than the available supply since
-          // the supply is maxed out at the end so won't always indicate accurately the buffer size
+          // the supply is increased at the end so it won't always report the buffer size accurately
           val isBufferEmpty: F2[Boolean] =
             buffer.size.map(_ == 0)
 
           val streamExhausted: F2[Boolean] =
             (isBufferEmpty, supplyEnded.get).mapN(_ && _)
 
-          //  releasing a number of permits equal to {groupSize} should be enough,
-          //  but in order to ensure termination of the consumer when the producer
-          //  stream is cancelled we need to max out the supply
-          val maxOutSupply: F2[Unit] =
-            supply.available.flatMap(a => supply.releaseN(Long.MaxValue - a))
+          //  releasing a number of permits equal to {groupSize} is enough in most cases, but in
+          //  order to ensure prompt termination of the consumer on interruption when the timeout
+          //  has not kicked in yet nor we've seen enough elements we need to max out the supply
+          val maxOutSupply =
+            supply.available.flatMap(av => supply.releaseN(Long.MaxValue - av))
 
           // enabling termination of the consumer stream when the producer completes
           // naturally (i.e runs out of elements)
-          val markSupplyEnd: F2[Unit] = supplyEnded.set(true)
+          val endSupply: F2[Unit] = supplyEnded.set(true) *> maxOutSupply
 
           val enqueue: F2[Unit] =
-            foreach(buffer.offer(_) <* supply.release).compile.drain
-              .guarantee(markSupplyEnd *> maxOutSupply)
+            foreach(buffer.offer(_) <* supply.release).compile.drain.guarantee(endSupply)
 
           val awaitAndEmitNext: F2[Chunk[O]] = for {
             isEmpty <- isBufferEmpty

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

@@ -833,7 +833,7 @@ class StreamCombinatorsSuite extends Fs2Suite {
         .assertEquals(0.millis)
     }
 
-    test("Propagation: upstream failures are propagated downstream") {
+    test("Upstream failures are propagated downstream") {
 
       case object SevenNotAllowed extends NoStackTrace
 
@@ -847,25 +847,41 @@ class StreamCombinatorsSuite extends Fs2Suite {
       downstream.compile.lastOrError.intercept[SevenNotAllowed.type]
     }
 
-    test("Propagation: upstream cancellation is propagated downstream") {
+    test("Upstream interruption causes immediate downstream termination with all elements being released") {
 
-      def source(counter: Ref[IO, Int]): Stream[IO, Int] = {
-        Stream
-          .unfold(0)(s => Some(s, s + 1))
-          .covary[IO]
-          .meteredStartImmediately(1.second)
-          .evalTap(counter.set)
-          .interruptAfter(5.5.seconds)
-      }
+      val sourceTimeout = 5.5.seconds
+      val downstreamTimeout = sourceTimeout + 2.seconds
 
-      def downstream(counter: Ref[IO, Int]): Stream[IO, Chunk[Int]] =
-        source(counter).groupWithin(Int.MaxValue, 1.day)
-
-      (for {
-        counter <- Ref.of[IO, Int](0)
-        _ <- downstream(counter).compile.drain
-        c <- counter.get
-      } yield c).assertEquals(5)
+      TestControl
+        .executeEmbed(
+          Ref[IO]
+            .of(0.millis)
+            .flatMap { ref =>
+              val source: Stream[IO, Int] =
+                Stream
+                  .unfold(0)(s => Some(s, s + 1))
+                  .covary[IO]
+                  .meteredStartImmediately(1.second)
+                  .interruptAfter(sourceTimeout)
+
+              // large chunkSize and timeout (no emissions expected in the window
+              // specified, unless source ends, due to interruption or
+              // natural termination (i.e runs out of elements)
+              val downstream: Stream[IO, Chunk[Int]] =
+                source.groupWithin(Int.MaxValue / 2 + 4, 1.day)
+
+              downstream.compile.lastOrError
+                .map(_.toList)
+                .timeout(downstreamTimeout)
+                .flatTap(_ => IO.monotonic.flatMap(ref.set))
+                .flatMap(emit => ref.get.map(timeLapsed => (timeLapsed, emit)))
+            }
+        )
+        .assertEquals(
+          // downstream ended immediately (i.e timeLapsed = sourceTimeout)
+          // emitting whatever was accumulated at the time of interruption
+          (sourceTimeout, List(0, 1, 2, 3, 4, 5))
+        )
     }
   }