Timed Pull

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

@@ -1472,97 +1472,56 @@ final class Stream[+F[_], +O] private[fs2] (private[fs2] val underlying: Pull[F,
     * Empty groups, which can occur if no elements can be pulled from upstream
     * in a given time window, will not be emitted.
     *
-    * Note: a time window starts each time downstream pulls.
+    * Note: a time window starts each time downstream pulls. TODO: correct this
     */
   def groupWithin[F2[x] >: F[x]](
       n: Int,
       d: FiniteDuration
   )(implicit F: Temporal[F2]): Stream[F2, Chunk[O]] =
-    Stream
-      .eval {
-        Queue
-          .synchronousNoneTerminated[F2, Either[Token, Chunk[O]]]
-          .product(F.ref(F.unit -> false))
-      }
-      .flatMap { case (q, currentTimeout) =>
-        def startTimeout: Stream[F2, Token] =
-          Stream.eval(Token[F2]).evalTap { token =>
-            val timeout = F.sleep(d) >> q.enqueue1(token.asLeft.some)
-
-            // We need to cancel outstanding timeouts to avoid leaks
-            // on interruption, but using `Stream.bracket` or
-            // derivatives causes a memory leak due to all the
-            // finalisers accumulating. Therefore we dispose of them
-            // manually, with a cooperative strategy between a single
-            // stream finaliser, and F finalisers on each timeout.
-            //
-            // Note that to avoid races, the correctness of the
-            // algorithm does not depend on timely cancellation of
-            // previous timeouts, but uses a versioning scheme to
-            // ensure stale timeouts are no-ops.
-            timeout.start
-              .bracket(_ => F.unit) { fiber =>
-                // note the this is in a `release` action, and therefore uninterruptible
-                currentTimeout.modify { case st @ (cancelInFlightTimeout, streamTerminated) =>
-                  if (streamTerminated)
-                    // the stream finaliser will cancel the in flight
-                    // timeout, we need to cancel the timeout we have
-                    // just started
-                    st -> fiber.cancel
-                  else
-                    // The stream finaliser hasn't run, so we cancel
-                    // the in flight timeout and store the finaliser for
-                    // the timeout we have just started
-                    (fiber.cancel, streamTerminated) -> cancelInFlightTimeout
-                }.flatten
-              }
-          }
-
-        def producer =
-          this.chunks.map(_.asRight.some).through(q.enqueue).onFinalize(q.enqueue1(None))
-
-        def emitNonEmpty(c: Chunk.Queue[O]): Stream[F2, Chunk[O]] =
-          if (c.size > 0) Stream.emit(c.toChunk)
-          else Stream.empty
-
-        def resize(c: Chunk[O], s: Stream[F2, Chunk[O]]): (Stream[F2, Chunk[O]], Chunk[O]) =
-          if (c.size < n) s -> c
-          else {
-            val (unit, rest) = c.splitAt(n)
-            resize(rest, s ++ Stream.emit(unit))
-          }
+    this.covary[F2].pull.timed { timedPull =>
+      def resize(c: Chunk[O], s: Pull[F2, Chunk[O], Unit]): (Pull[F2, Chunk[O], Unit], Chunk[O]) =
+        if (c.size < n) s -> c
+        else {
+          val (unit, rest) = c.splitAt(n)
+          resize(rest, s >> Pull.output1(unit))
+        }
 
-        def go(acc: Chunk.Queue[O], currentTimeout: Token): Stream[F2, Chunk[O]] =
-          Stream.eval(q.dequeue1).flatMap {
-            case None => emitNonEmpty(acc)
-            case Some(e) =>
-              e match {
-                case Left(t) if t == currentTimeout =>
-                  emitNonEmpty(acc) ++ startTimeout.flatMap { newTimeout =>
-                    go(Chunk.Queue.empty, newTimeout)
-                  }
-                case Left(_) => go(acc, currentTimeout)
-                case Right(c) =>
-                  val newAcc = acc :+ c
-                  if (newAcc.size < n)
-                    go(newAcc, currentTimeout)
-                  else {
-                    val (toEmit, rest) = resize(newAcc.toChunk, Stream.empty)
-                    toEmit ++ startTimeout.flatMap { newTimeout =>
-                      go(Chunk.Queue(rest), newTimeout)
-                    }
-                  }
-              }
-          }
+      // Invariants:
+      // acc.size < n, always
+      // hasTimedOut == true iff a timeout has been received, and acc.isEmpty
+      def go(acc: Chunk.Queue[O], timedPull: Stream.TimedPull[F2, O], hasTimedOut: Boolean = false): Pull[F2, Chunk[O], Unit] =
+        timedPull.uncons.flatMap {
+          case None =>
+            Pull.output1(acc.toChunk).whenA(acc.nonEmpty)
+          case Some((e, next)) =>
+            def resetTimerAndGo(q: Chunk.Queue[O]) =
+              timedPull.timeout(d) >> go(q, next)
+
+            e match {
+              case Left(_) =>
+                if (acc.nonEmpty)
+                  Pull.output1(acc.toChunk) >> resetTimerAndGo(Chunk.Queue.empty)
+                else
+                  go(Chunk.Queue.empty, next, hasTimedOut = true)
+              case Right(c) if hasTimedOut =>
+                // it has timed out without reset, so acc is empty
+                val (toEmit, rest) =
+                  if (c.size < n) Pull.output1(c) -> Chunk.empty
+                  else resize(c, Pull.done)
+                toEmit >> resetTimerAndGo(Chunk.Queue(rest))
+              case Right(c) if !hasTimedOut =>
+                val newAcc = acc :+ c
+                if (newAcc.size < n)
+                  go(newAcc, next)
+                else {
+                  val (toEmit, rest) = resize(newAcc.toChunk, Pull.done)
+                  toEmit >> resetTimerAndGo(Chunk.Queue(rest))
+                }
+            }
+        }
 
-        startTimeout
-          .flatMap(t => go(Chunk.Queue.empty, t).concurrently(producer))
-          .onFinalize {
-            currentTimeout
-              .getAndSet(F.unit -> true)
-              .flatMap { case (cancelInFlightTimeout, _) => cancelInFlightTimeout }
-          }
-      }
+      timedPull.timeout(d) >> go(Chunk.Queue.empty, timedPull)
+    }.stream
 
   /**
     * If `this` terminates with `Stream.raiseError(e)`, invoke `h(e)`.
@@ -4289,6 +4248,80 @@ object Stream extends StreamLowPriority {
               Pull.output(pfx) >> Pull.pure(Some(tl.cons(sfx)))
           }
       }
+
+    /**
+      * Allows expressing `Pull` computations whose `uncons` can receive
+      * a user-controlled, resettable `timeout`.
+      * See [[TimedPull]] for more info on timed `uncons` and `timeout`.
+      *
+      * As a quick example, let's write a timed pull which emits the
+      * string "late!" whenever a chunk of the stream is not emitted
+      * within 150 milliseconds:
+      *
+      * @example {{{
+      * scala> import cats.effect.IO
+      * scala> import cats.effect.unsafe.implicits.global
+      * scala> import scala.concurrent.duration._
+      * scala> val s = (Stream("elem") ++ Stream.sleep_[IO](200.millis)).repeat.take(3)
+      * scala> s.pull
+      *      |  .timed { timedPull =>
+      *      |     def go(timedPull: Stream.TimedPull[IO, String]): Pull[IO, String, Unit] =
+      *      |       timedPull.timeout(150.millis) >> // starts new timeout and stops the previous one
+      *      |       timedPull.uncons.flatMap {
+      *      |         case Some((Right(elems), next)) => Pull.output(elems) >> go(next)
+      *      |         case Some((Left(_), next)) => Pull.output1("late!") >> go(next)
+      *      |         case None => Pull.done
+      *      |       }
+      *      |     go(timedPull)
+      *      |  }.stream.compile.toVector.unsafeRunSync()
+      * res0: Vector[String] = Vector(elem, late!, elem, late!, elem)
+      * }}}
+      *
+      * For a more complex example, look at the implementation of [[Stream.groupWithin]].
+      */
+    def timed[O2, R](pull: Stream.TimedPull[F, O] => Pull[F, O2, R])(implicit F: Temporal[F]): Pull[F, O2, R] =
+      Pull
+        .eval { Token[F].mproduct(id => SignallingRef.of(id -> 0.millis)) }
+        .flatMap { case (initial, time) =>
+
+          def timeouts: Stream[F, Token] =
+            time
+              .discrete
+              .dropWhile { case (id, _) => id == initial }
+              .switchMap { case (id, duration) =>
+                // We cannot move this check into a `filter`:
+                // we want `switchMap` to execute and cancel the previous timeout
+                if (duration != 0.nanos)
+                  Stream.sleep(duration).as(id)
+                else
+                  Stream.empty
+              }
+
+          def output: Stream[F, Either[Token, Chunk[O]]] =
+            timeouts.map(_.asLeft)
+              .mergeHaltR(self.chunks.map(_.asRight))
+              .flatMap {
+                case chunk @ Right(_) => Stream.emit(chunk)
+                case timeout @ Left(id) =>
+                  Stream
+                    .eval(time.get)
+                    .collect { case (currentTimeout, _) if currentTimeout == id => timeout }
+              }
+
+          def toTimedPull(s: Stream[F, Either[Token, Chunk[O]]]): TimedPull[F, O] = new TimedPull[F, O] {
+            type Timeout = Token
+
+            def uncons: Pull[F, INothing, Option[(Either[Timeout, Chunk[O]], TimedPull[F, O])]] =
+              s.pull.uncons1
+                .map( _.map { case (r, next) => r -> toTimedPull(next) })
+
+            def timeout(t: FiniteDuration): Pull[F, INothing, Unit] = Pull.eval {
+              Token[F].tupleRight(t).flatMap(time.set)
+            }
+          }
+
+          pull(toTimedPull(output))
+        }
   }
 
   /** Projection of a `Stream` providing various ways to compile a `Stream[F,O]` to a `G[...]`. */
@@ -4560,6 +4593,39 @@ object Stream extends StreamLowPriority {
     def toVector: G[Vector[O]] = to(Vector)
   }
 
+/**
+  * An abstraction for writing `Pull` computations that can timeout
+  * while reading from a `Stream`.
+  *
+  * A `TimedPull` is not created or intepreted directly, but by calling [[timed]].
+  * {{{
+  * yourStream.pull.timed(tp => ...).stream
+  * }}}
+  *
+  * The argument to `timed` is a `TimedPull[F, O] => Pull[F, O2, R]`
+  * function which describes the pulling logic, and is often recursive, with shape:
+  *
+  * {{{
+  * def go(tp: TimedPull[F, A]): Pull[F, B, Unit] =
+  *   tp.uncons.flatMap {
+  *     case Some((Right(chunk), next)) => doSomething >> go(next)
+  *     case Some((Left(_), next)) => doSomethingElse >> go(next)
+  *     case None => Pull.done
+  *   }
+  * }}}
+  *
+  * Where `doSomething` and `doSomethingElse` are `Pull` computations
+  * such as `Pull.output`, in addition to `TimedPull.timeout`.
+  *
+  * See below for detailed descriptions of `timeout` and `uncons`, and
+  * look at the [[ToPull.timed]] scaladoc for an example of usage.
+  */
+  trait TimedPull[F[_], O] {
+    type Timeout
+    def uncons: Pull[F, INothing, Option[(Either[Timeout, Chunk[O]], TimedPull[F, O])]]
+    def timeout(t: FiniteDuration): Pull[F, INothing, Unit]
+  }
+
   /**
     * When merging multiple streams, this represents step of one leg.
     *

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

@@ -0,0 +1,47 @@
+/*
+ * 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
+
+
+object ex {
+
+import cats.effect.IO
+import cats.effect.unsafe.implicits.global
+import scala.concurrent.duration._
+
+val emitThenWait = Stream("elem") ++ Stream.sleep_[IO](200.millis)
+val s = emitThenWait.repeat.take(3)
+
+def a = s.pull
+ .timed { timedPull =>
+    def go(timedPull: Stream.TimedPull[IO, String]): Pull[IO, String, Unit] =
+      timedPull.timeout(150.millis) >> // starts a timeout timer, resets the previous one
+      timedPull.uncons.flatMap {
+        case Some((Right(elems), next)) => Pull.output(elems) >> go(next)
+        case Some((Left(_), next)) => Pull.output1("late!") >> go(next)
+        case None => Pull.done
+      }
+
+      go(timedPull)
+ }.stream.compile.toVector.unsafeRunSync()
+//val res0: Vector[String] = Vector(elem, late!, elem, late!, elem)
+}