Composing complex programs with ZIO

There is a fundamental mental shift a procedural programmer must make when embracing functional programming. In this lesson we'll learn to visualize complex ZIO programs to compose them correctly.

The art of type chasing

Types in Scala give us confidence that the program is doing what we want and that the quality of the data is high enough. With ZIO they also help us catching nasty bugs at compile time instead of runtime.

The canonical ZIO program has this type:

ZIO[ZEnv, Nothing, ExitCode]

This is our end game, every ZIO program we can possibly conceive, if we transform it into (chase) this type, then we can easily run it.

Usually what we start with will be something like

ZIO[ZEnv with MyDatabase, DomainSpecificErrors, Unit]

to make it coincide with our canonical ZIO, we need to act in two different directions.

• The environment type needs to become more generic (up to ZEnv)
• The error and return type need to become more specific (error down to Nothing and return to ExitCode)

By using ZIO provided methods we can creatively chase the canonical types.

Cat and mouse chasing

As fun example of type chasing we can create a program of a cat chasing a mouse.

We will have a mouse program and a cat program that will run concurrently. Both programs will share a variable to track the mouse distance from the cat.

Creating a variable

Our variable, tracking the lead the mouse has on the cat, needs to be shared by both the cat and the mouse program.

To create a variable that is safe to use concurrently we can write the following program

import zio.Ref

val makeMouseLead: UIO[Ref[Int]] = Ref.make(10)

This is a UIO program which expands to a ZIO[Any, Nothing, Ref[Int]]. It can run in any environment, never fails, and returns a Ref of an integer.

Ref is a reference to an immutable variable that we can check and swap. You can picture a Ref like a single place datum cradle. We will use it to read the distance between cat and mouse and update it whenever either of them advances.

Digital mouse that runs away

To represent the mouse getting away from the cat we can design a simple recursive program

import zio.clock.sleep
import zio.console.putStrLn
import zio.random.nextIntBounded
import zio.duration._
import zio._

def mouseProgram(mouseLead: Ref[Int]): URIO[ZEnv, Unit] =
  for {
    cm <- nextIntBounded(2)
    _  <- mouseLead.update(_ + cm)
    _  <- putStrLn(s"Mouse advances by $cm cm")
    _  <- sleep(100.milliseconds)
    _  <- mouseProgram(mouseLead)
  } yield ()

Let's dissect this automatic mouse:

• We are accepting the reference to an Int variable, the mouse lead on the cat.
• A random number, less than 2, is drawn, representing how many centimeters the mouse advances.
• We add the above number to the current lead and put the sum in the Ref. This will be the new mouse lead.
• Print on the screen the mouse progress
• We then wait 100 milliseconds
• Recursively call the mouse program

Notice that we are using functions coming from zio.clock, zio.console and zio.random, all contained in ZEnv so our final type is URIO[ZEnv, Unit].

I wrote ZEnv instead of Clock with Console with Random. This is allowed because I can always pick a more specific type as environment. The opposite is true for error type and return type.

This program is recursive, it will never return the Unit we promised. To make the return type more expressive it would have to return Nothing meaning it will never do. But, as said above we cannot put a more specific type for the return channel.

The problem is that the compiler didn't detect our infinite recursion and assumed at some point we are going to return Unit. There is a way we can explicitly say that the mouse will run forever.

def mouseProgram(mouseLead: Ref[Int]): URIO[ZEnv, Nothing] =
  for {
    cm    <- nextIntBounded(2)
    _     <- mouseLead.update(_ + cm)
    _     <- putStrLn(s"Mouse advances by $cm cm")
    _     <- sleep(100.milliseconds)
    recur <- mouseProgram(mouseLead)
  } yield recur

By making the return type the mouse program itself we can force write Nothing and the compiler will be happy.

Thanks to the ZIO library we can actually do this by skipping the recursive step entirely:

def mouseProgram(mouseLead: Ref[Int]): URIO[ZEnv, Nothing] =
  (for {
    cm    <- nextIntBounded(2)
    _     <- mouseLead.update(_ + cm)
    _     <- putStrLn(s"Mouse advances by $cm cm")
    _     <- sleep(100.milliseconds)
  } yield ()).forever

In this code, we deleted the recursive call and wrapped the for comprehension in parenthesis. With forever we created a new program that repeats the for comprehension forever.

Digital cat that chases

The third program we are going to write is the dual of the mouse. The cat is very similar, but it runs faster and, for every centimeter it advances, the mouse lead diminishes instead of increasing.

def catProgram(mouseLead: Ref[Int]): URIO[ZEnv, Unit] =
  for {
    cm   <- nextIntBounded(5)
    lead <- mouseLead.updateAndGet(_ - cm)
    _    <- putStrLn(s"Cat advances by $cm cm, mouse is $lead cm away")
    _    <- sleep(100.milliseconds)
    _    <- catchDetector(lead, mouseLead)
  } yield ()

The cat program will actually finish at some point: the last instruction uses a special program called catchDetector. It considers two possible scenarios: either the mouse has still some lead or the cat is at or past the mouse.

def catchDetector(currentLead: Int, mouseLead: Ref[Int]): ZIO[ZEnv, Nothing, Unit] =
  currentLead match {
    case lead if lead > 0  => catProgram(mouseLead)
    case lead if lead <= 0 => putStrLn("Cat catches the mouse!")
  }

Here, if the mouse is not reached yet, the cat program is called (recursively), otherwise the catchDetector returns with a program that prints "Cat catches the mouse!".

We can make the program even more fun by including a third possibility: if the cat goes past the mouse, instead of declaring victory, we assume the mouse escaped. We can express this as an exception. From the point of view of the cat program, this is certainly not a welcome outcome.

First, let's declare the exception as a case class.

case class MouseEscapedException(catLead: Int)

The variable catLead will tell us how much the cat went past the mouse.

Secondly, we need to add the third case in the catchDetector.

def catchDetector(currentLead: Int, mouseLead: Ref[Int]): ZIO[ZEnv, MouseEscapedException, Unit] =
  currentLead match {
    case 0                 => putStrLn("Cat catches the mouse!")
    case lead if lead > 0  => catProgram(mouseLead)
    case lead if lead < 0 => ZIO.fail(MouseEscapedException(lead))
  }

Now if the currentLead is strictly less than 0 we return a program that fails with MouseEscapedException. This makes our program fail in turn.

Since the cat program is calling this, we need to update the signature by adding the error modality:

def catProgram(mouseLead: Ref[Int]): ZIO[ZEnv, MouseEscapedException, Unit] = ...

Putting it all together

We have created a total of 4 ZIO programs. This is a good time to list them all.

• UIO[Ref[Int]] will give us an integer variable
• URIO[ZEnv, Nothing] programs the mouse to run
• ZIO[ZEnv, MouseEscapedException, Unit] programs the cat to chase
• ZIO[ZEnv, MouseEscapedException, Unit] subroutine of the cat that can either succeed with a Unit or trigger failure with a MouseEscapedException.

The diagram below is another representation of these programs.

Notice, at the center, our canonical ZIO program URIO[ZEnv, ExitCode]. In green I represented the environment channel. The canonical ZIO program will provide you a perfectly good ZEnv (filled green dot). The program that creates the mouse and also the cat, conversely, want the ZEnv (hence the hollow green dot).

In blue I represented the success channel. The mouse catch detector may return a unit when a capture happens, this gets transmitted upwards). The mouse program runs forever, and returns nothing, that's why it doesn't have any blue line going outwards. The canonical ZIO waits patiently for a ExitCode and won't accept much else.

In red we see the failures channel that, from the catch detectors propagate upwards. The cat, in turn, may propagate the error.

We have to run the cat and the mouse program, to do that we have to chase the canonical ZIO types.

Environment channel

There is not a lot to do for the green lines. ZEnv with ZEnv, everything matches. We won't worry about this.

Failure channel

The canonical ZIO doesn't want to deal with errors we have to stop them before they reach the boundaries of our program. One way to do this is to turn the error into a program that never fails (a ZIO with Nothing as error type); putStrLn is one such a program. We could log the error to screen and this would return a success Unit type. A message saying that the mouse escaped will be sufficient.

Success channel

The only type accepted here by the canonical ZIO is ExitCode. We said this before, but it's worth repeating that ExitCode is just a puffed up integer we give to the operating system. We can always return ExitCode.success, an alias for 0. This means we consider the cat both catching and missing the mouse a success.

The above three conditions are satisfied by compositeProgram.

val compositeProgram: URIO[zio.ZEnv, ExitCode] =
  for {
    lead <- makeMouseLead
    exitCode <- catProgram(lead)
      .raceFirst(mouseProgram(lead))
      .catchAll(exc => putStrLn(s"Cat was ${-exc.catLead} cm ahead when it lost the mouse."))
      .as(ExitCode.success)
  } yield exitCode

Check out the complete code.

Click "Run" a few times to see all possible outcomes.

Explanation of compositeProgram

Our idea to chase types is a very small for comprehension. In the first place we map over the makeMouseLead program to create the variable. Secondly we run both the mouse and the cat with .raceFirst which will concurrently spin them up and wait for the first to terminate either with a success or a failure. We know that the mouse program doesn't terminate so we are basically waiting the cat program. This however leaves us with a ZIO[ZEnv, MouseEscapedException, Unit].

To turn the error into a Nothing we use .catchAll and we pass an infallible program: putStrLn. To turn the Unit into an exit code we could simply .flatMap it to an ExitCode but when we are not interested in the result we can use .as which is a .flatMap that ignores the preceding result.

All these transformations are summarized in the below diagram which now looks complete.

Bonus section

But wait, there is more. Instead of using the .catchAll and .as(ExitCode.success) functions, ZIO gives us a nice shortcut that does something very similar and that will work in the majority of cases. You can just add .exitCode and this will either print an error (and exit with 1) or succeed (and exit with 0). Not precisely our behavior, we lose the nice error message and we exit with an error when the mouse escape, but good enough: our program is simpler and we save one line.

val compositeProgram: URIO[zio.ZEnv, ExitCode] =
    for {
      lead <- makeMouseLead
      exitCode <- catProgram(lead)
        .raceFirst(mouseProgram(lead))
        .exitCode
    } yield exitCode