Make your own IO with context-scoped data on top of ZIO

Photo by Anas Alshanti on Unsplash

Scala is a programming language that brings OOP and FP together. The usual story of adopting Scala is to start using it as a “better Java”. But as time goes by, and once you get enough confidence there is a tendency to adopt Function Programming in all its glory and pain.

You’re no longer afraid of implicits, know what does the typeclass mean and maybe even tried to use Shapeless for CSV codecs. So it’s time to start a new project and you ready want to make it pure functional from a day one.

The first question you should answer is what approach you should use for building a new shiny microservice? OOP/Spring Framework is not an option as you an experienced FP-developer, but what alternatives do we have?

As you already may know, there are dozens of them:

• Cake Pattern or Cakeless
• Free Monad
• Akka
• Final Tagless
• FP-pure code with DI like MacWire/Google Guice/DIStage/ReaderT

The community seems to be settled on using Final Tagless, so you made a choice. You choose the most popular IO-monad from Cats-Effect, and the sky is blue and everything is going to be great.

But… How to convenience your teammates? What if you want to adopt FP in a big company with thousands of employees? The explain/train strategy simply doesn’t scale. There are a ton of tricks and pitfalls — you should understand how error management works, when to use Context Shift to block/unblock or why it could be useful for fairness, how Bracket/Resource can leak if you are not careful enough, how to use IO.Par for doing Parallel operations, and so on, and so on.

That’s why ZIO is promising. It’s a great project which is easy-to-use for newcomers. It even has great interoperability with Cats! But the vendor lock issue is still a thing. In an ideal world, we want to minimize our dependencies and have all of them under control. So if you want to, you can switch them easily, restrict some parts of its API or extend it by providing required API for your specific problems.

Here I will try to explore the idea of having your own IO on top of ZIO (you can replace ZIO with Monix IO or Cats IO, btw). At the end it has the same benefits as ZIO, but it’s completely yours and more manageable.

Real-World Problem

Let’s imagine I want to store some request/session scope data easily. It could be authorization info, request transaction id, etc. Basically, it could be any data that you usually store in thread-locals if you’re from Java world. The idea is not new and I see others trying to do the same — either by putting this thread-local data to implicit parameters and declare it everywhere or by relying on monad transformer or Kleisli algebras.

Note that even if it’s still possible to store service dependencies in a request-scoped environment, we will not do this. It’s better to use the right tool for dependencies and our industry already has a good solution for it — it’s an old good Dependency Injection framework. I personally recommend to try DIStage.

But the problem with implicits/monad transformers/Kleisli is that it’s too invasive and not straightforward to use. My motivation is to make it easy for developers which are new to Scala, but also flexible as a Final Tagless approach if you want to.

In this post we will see how to do this. The plain is:

• Create our own abstraction on top of ZIO
• Extend our abstraction by providing nice API for working with request-scoped data like authorization and transaction id
• Add an interoperability with Cats and Final Tagless approach
• Explore what are the disadvantages and the possibilities of future improvements

Creating our own abstraction on top of ZIO

First, we need to define a simple trait with basic operations. It has 2 type parameters: Ctx for environment and A is for value:

Can we use right now? Nope, there is no data type with our trait:

The implementation of map is trivial:

The flatMap is tricky, since ZIO’s flatMap expects RIO not Env:

We need a method that will return a ZIO for a given Env. Let’s make it simple for now and just define it inside of our Env trait. I know it’s too much coupling, but we will try to solve it later.

The updated Env:

And now we can implement flatMap in our class by calling f(a).run:

We already can build simple programs with code above:

This is just a description of a program. And to actually run it, we have to convert it to ZIO and run using ZIO Runtime like this:

As you can see, our API leaks a lot of ZIO details, so in order to abstract away from it we want to make some changes:

• make EnvCtx constructor private and provide utility create methods in companion object
• include run methods directly into our trait, so it can be easily run with default runtime
• in order to abstract from ZIO exit status, we will need to convert it to our own Exit

Updated version:

Implementing Exit.fromZIOExit is trivial:

Our program looks better now:

That’s cool, but we want to work with some context, so let’s try to use one.

Adding API for working with request-scoped data

First, let me remind you the ZIO API.

ZIO has 3 type parameters — one for the result, one for error and one for the environment. The idea behind the 3rd parameter is to have some kind of Cake pattern build into IO-container. You can restrict your functions by declaring that they will return a ZIO which will rely on Console service for example and then in order to run you have to provide a corresponding environment with Console service implementation. If you combine your function with other function which relies on Clock service, the result function will require Console with Clock and so on. You can have as many layers as you want to, but you have to provide an implementation of all of them in your main method.

Here we will exploit the power of R for a different purpose. We will put a ZIO Fiber Ref inside of our R and rely on ZIO runtime to preserve our value between ZIO fibers, so it would be possible to override the ref. If you’re no familiar with semantics of FiberRefs in ZIO, the FiberRefSpec has some easy-to-read examples.

Here will use the idiomatic ZIO way to define service as part of R. First, we define a trait that returns our new service.

The service has the following API:

UIO[A] is a type alias for ZIO[Any, Nothing, A]

We need some means to create our new service as well:

Here we create a new FiberRef and just proxy all methods to update it via our service.

Another idiomatic part of working with ZIO is to add helpers methods to service companion object, so they can be accessed easily:

The full code of Service for working with environment:

Once we have a ZIO service prepared, we can get back to our Env type and use them:

Let’s build a simple program which illustrates how it’s possible to work with the environment with EnvCtx:

As a result, the user of EnvCtx knows nothing about how our environment stored — if there is a fiber ref or something else, ZIO or Monix — it doesn’t matter.

Looks good for me. But that’s not all.

Interoperability with Cats

The one thing which is good to have is interoperability with Cats. Why’s it important? Because of the code of others. There are plenty of useful libraries that are polymorphic by effect type, which means you can use them directly by providing your custom effect type. Thanks for Final Tagless approach, all we need to do is implement typeclass instances of Cats (default-choice library for doing FT in Scala).

Let’s start from simple ones. Let’s create empty object catz with all implicits. Timer and ContextShift instances are built based on ZIO Cats interop like this:

Now we have to implement all Cats typeclasses. ConcurrentEffect inherits from all of them, so we can start top to bottom. There is a lot of code but it’s just wrappers.

where EnvConcurrentEffect

Let’s build simple Final Tagless program and see if we can parametrize it with our EnvCtx:

Can you spot an error in this code? We forgot our Runtime which is needed for ConcurrentEffect! We can define the implicit in-place by providing Runtime and invoke unsafe run on it to actually get it. But this looks cumbersome:

Here we have to repeat ourselves 2 times — first, when we derive ConcurrentEffect and second time when we run the program. We can do a little bit better by introducing an implicit dependency on service which provides an empty environment and derive Runtime based on that automatically.

Unsafe method to create Context[Ctx] will be useful:

Final solution with all above:

To eliminate repeated ExampleEnv(“init”), we can introduce specific methods to run unsafe with an empty environment.

What other options do we have? If the environment is Option[Ctx] we can derive provideEnv automatically be providing a None as an empty state. I like the former approach better but didn’t have enough time to experiment with it.

Summary

Wrapping and making your own app-specific monad is not so hard (thanks to ZIO). You can add a lot of useful stuff there which is common for your company/application, like passing the NewRelic transaction id, add custom metric methods, adding custom thread-pools — whatever you wish. Explore new ideas and have fun!

P.S. Careful reader may remember that at the beginning of the post, I said adding def run: RIO[Ctx, A] method to trait Env is a short-cut. The main idea to fix this is to introduce a different abstraction called Runner, so if you want to change the underlying IO-container, you will do this only in one place. But this is a long post already, I'm going to get back to this later. Feel free to ping me if you have any questions:)

You can check code examples in Github.