Essential Effects: How do you "loop" with IO?

I recently tweeted a question

would “how do you loop in IO?” be something you’re interested in?

and due to the overwhelmingly positive response (5 likes!) I couldn’t hold myself back. I’m not exactly sure where this will live in the book, but it can be a good sidebar or extra section somewhere. I got the code wrong a few times, so it’s not exactly trivial.

If you have any questions or comments about this, please let me know!

Note: For these examples, we’ll delegate when we should stop to a method def shouldStop: Boolean.

In an imperative setting you might use the humble while loop:

def happyForever: Unit =
  while (!shouldStop) println("I'm happy!")

But this won’t work in the world of effects, because execution is delayed. We’re constructing an IO value that describes what we want to happen, but nothing executes it. (Usually there’s a visible or invisible call to the unsafeRunSync method.)

def happyForever: Unit =
  while (!shouldStop) IO(println("I'm happy!")) // <1>

<1> Nothing will ever be printed.

Let’s fix this and put the loop within the effect:

val happyForever: IO[Unit] =
  IO { while (!shouldStop) println("I'm happy!") }

There’s a problem with this definition though: it doesn’t respect cancelation. For example:

(
  happyForever,
  IO.raise(new RuntimeException("oh noes!"))
).parTupled

When this code is executed will run both effects in parallel, but the second effect will raise an error which will then cancel the other effects in the group. However, happyForever will continue to print I'm happy! forever. Why is that?

If cancelation did stop our effect–which it didn’t–it would have to do so at the level of the underlying thread, because the while loop of the effect body will never stop. However, nothing in our code checks any IO “cancelation status”, and the IO itself knows nothing about the code “inside” it. How can we get the concept of “cancelation status” pushed into our effect, so we could stop the while loop?

To answer, let’s examine the only mechanism we’ve defined the could stop the effect, the shouldStop method. If happyForever is cancelled, could we get shouldStop to return true?

// TODO: return true if the effect is cancelled
def shouldStop: Boolean = ???

Maybe we can have happyForever know if it is cancelled?

val happyForever: IO[Unit] =
  IO { while (!shouldStop(this)) println("I'm happy!") }

// TODO: return true if the effect is cancelled
def shouldStop[A](effect: IO[A]): Boolean = ???

But sadly–and for important reasons–there is no method on IO that reports its cancelation status. We do have, however, the static method IO.cancelBoundary:

def cancelBoundary: IO[Unit]

It doesn’t tell you if some effect has been cancelled or not, it only produces a Unit. Instead it won’t allow any subsequent execution if the “current” effect has been cancelled.

Let’s integrate it: the shouldStop method will be for application-level control of the effect, whereas we’ll add a separate cancelation check with IO.cancelBoundary:

val happyForever: IO[Unit] =
  IO { 
    while (!shouldStop || !IO.cancelBoundary) // <1>
      println("I'm happy!")
  }

// TODO: application-level stopping logic
def shouldStop: Boolean = ???

<1> This won’t work. It doesn’t even compile.

Inserting IO.cancelBoundary like this won’t work. It returns an IO[Unit], not a Boolean. We shouldn’t execute IO.cancelBoundary ourselves either. The problem lies in the while loop: we need it to act as an effect so it can work with other effects.

Let’s avoid that problem for a moment and focus on only one “iteration” of the loop:

val happyOnce: IO[Unit] =
  for {
    _ <- IO.cancelBoundary         // <1>
    _ <- IO(println("I'm happy!")) // <2>
  } yield ()

<1> Check if the effect has been cancelled, and if so, no further effects will be executed. <2> Otherwise perform our printing.

We’re not done–this only runs one iteration, and doesn’t check shouldStop. How do we do it repeatedly? We just call “ourselves” again. That is, we recurse:

val happyForever: IO[Unit] =
  for {
    _ <- IO.cancelBoundary
    _ <- IO(println("I'm happy!"))
    _ <- if (shouldStop) IO.unit
         else happyForever // <1>
  } yield ()

<1> If we shouldn’t stop, “do it again” via recursion.

Voilà!

(If you’re worried that we’ll get a stack overflow, we won’t because IO uses trampolined execution that trades stack space for heap space. That, plus the fact that happyForever is tail-recursive, makes this safe at runtime.)

Technically speaking, with the recursive structure we don’t need to explicitly add the IO.cancelBoundary effect to have our “loop” canceled, because Cats Effect automatically checks the cancelation status every 512 flatMap calls during execution. It’s a kind of fail-safe to attempt to ensure cancelled effects do get cancelled. By being explicit, however, we can exit from our loop sooner rather than later.

To conclude, here are a couple of variations of the solution that take advantage of some useful combinators in the Cats library. First, there’s the iterateUntil combinator if you want to avoid the explicit recursion:

import cats.implicits._

val happyForever: IO[Unit] = {
  val happyOnce =
    for {
      _ <- IO.cancelBoundary
      _ <- IO(println("I'm happy!"))
    } yield ()

  happyOnce.iterateUntil(shouldStop)
}

Or you could use the whileM_ combinator to more explicitly separate the effect from the stopping condition:

import cats.implicits._

val happyForever: IO[Unit] =
  IO(println("I'm happy!"))
    .whileM_(IO.cancelBoundary.map(_ => !shouldStop)) // <1>

<1> Check the cancelation status. If it isn’t cancelled, return true if the loop should continue within an effect.

(Yes, the M- and _-suffixes are a bit distracting. The M is for “monad”, as these combinators require a flatMap method to work. The _ suffix denotes methods that “throw away” their outputs and only produce Unit; we avoid the suffix-less whileM because it collects the result of each iteration.)