Context Passing in Tagless Final

Context passing is the thing that all programmers face regardless of a programming language they use. In this article I’d like to discuss the ways how this can be solved in Scala using Cats Effect, ZIO, cats-mtl and the tagless final encoding.

Context and its way through

Before going further with coding we should answer a few questions - What is a context? and Why do we need to pass it?.

What is a context?

In this article I would use context in the meaning of preloaded structured immutable information about something we need across a stack of calls made in our business logic(function).

Context should have the next properties:

  • Info should be loaded to memory - in case if a part of the context is stored in external memory, this info should be loaded from it in advance BEFORE executing the logic.
    Opposite would be requesting the info from the external memory only when it is needed in the business logic.
  • Info should be extracted/parsed/validated - info should be parsed into a domain specific data structure and validated.
    Opposite would be having this context info represented as a string or JSON or any other low-level data representation.
  • Info should be immutable during execution of an operation - data structure and info in it is immutable after it is prepared. It does not represent the latest changes that could be made in the external memory during execution of the operation.
  • The operation itself is considered as short-lived

Why do we need to pass it?

Here are some practical examples of a context:

  • Information about a user account that is used for a http request - like its id for tracking user actions in the system.
  • Session information of the account like access roles assigned to it - for checking if certain actions are allowed across the business logic.
  • Information about the http request itself - a request id, a request’s timestamp and other info useful for logging and tracing.

Explicit context passing

In functional Scala we have quite a lot of different tools to be used. But let’s start with the simplest one we have - an explicit value. We pass a context as an argument of a function or a constructor parameter in this case.

object Example {
  // Passing `context` as an explicit argument
  def passContext1(arg1: Arg1, arg2: Arg2, context: Context): Result = ???

  // Passing `context` as an implicit argument
  def passContext2(arg1: Arg1, arg2: Arg2)(using context: Context): Result = ???
}

// Passing `context` as an explicit constructor parameter
class PassContext3(arg1: Arg1, arg2: Arg2, context: Context) {
  def doSomething(): Result
}

// Passing `context` as an implicit constructor parameter
class PassContext4(arg1: Arg1, arg2: Arg2)(using Context) {
  def doSomething(): Result
}

This is the simplest way of passing a context in Scala. The disadvantage here is the fact that the context value must be passed through all layers of a call stack even if it is used on in a few of them. All methods have to have Context parameter defined explicitly in their signatures whether as an explicit or an implicit parameter.

ThreadLocal

Another way is the old way of Java - ThreadLocal(or InheritableThreadLocal).

object ContextHolder {
  private val context = new ThreadLocal[Context]()
  def setContext(in: Context): Unit = context.set(in)
  def getContext(): Context = context.get()
  def dropContext(): Unit = context.remove()
}

class Controller(service: Service) { 
  // Here is happy-path logic only
  def doSomething(args: Args): Result = {
    val context = ??? // get context from somewhere
    ContextHolder.setContext(context)
    service.doSomethingElse()
    ContextHolder.dropContext()
  }
}

class Service {
  def doSomethingElse(): Result = {
    val context = ContextHolder.getContext()
    //... context dependent logic is here
    ???
  }
}

I can’t recommend to use it in Scala code nowadays as it is way too unsafe in terms of null values and memory leaks. Also it is necessary to be mindful about releasing/replacing the old context value with the new one. Another thing is understanding how the logic is mapped on threads. For example, in cats-effect it might be quite difficult to understand what particular thread is used for executing one or another line of logic. The reason is that different parts of a program can be run on different threads in an execution context.

Cats’ IOLocal

IOLocal is an alternative for ThreadLocal made specifically for cats-effect. The idea behind it is the same as for ThreadLocal but it can work properly around cat’s Fiber-s. For more details check this example.

ZIO’s FiberRef

Zio also has its own implementation of a fiber local values - FiberRef. More information about it and how to use it for passing a context around can be found here.

Context passing in Tagless Final

But there is, I think, the better way of doing what we are doing with our context - use type classes and the Tagless Final approach 😎

Services

First of all, let’s define our types that we are going to use in this example.

  • Context - is a representation of the context we are going to pass across our calls’ stack.
  • Input - is a type for input params. For simplicity let’s use this one type for all inputs.
  • Output - is the same but for outputs of our functions.

Then we need to bring a type class from cats-mtl - Ask[F[_], Ctx]. It makes possible of getting a value of Ctx out from F[_].

Code

  type Context
  type Input
  type Output

  trait Layer1[F[_]]:
    def operation1(in: Input): F[Output]

  trait Layer2[F[_]]:
    def operation2(in: Input): F[Output]

  trait Layer3[F[_]]:
    def operation3(in: Input): F[Output]

  trait Layer4[F[_]]:
    def operation4(in: Input): F[Output]

Now we need to bring a type class from cats-mtl - Ask[F[_], Ctx]. It makes possible of getting a value of Ctx out from F[_]. To improve readability we can define such a type alias that we can use later as a single argument type class.

Code

  import cats.mtl.Ask
  final type AskCtx[F[_]] = Ask[F, Context]

Next step is to implement these layers. Our goal here is to pass the context value between all layers but use it only in Layer2 and Layer4.

Impl1 does nothing but calls the next layer. It does not use the context.

Code

final class Impl1[F[_]](layer2: Layer2[F]) extends Layer1[F]:
  override def operation1(in: Input): F[Output] = layer2.operation2(in)

Impl2 extracts a context value and verifies it.

Code

final class Impl2[F[_]](
    verify: Context => F[Boolean],
    layer3: Layer3[F]
)(using A: AskCtx[F], MT: MonadError[F, Throwable])
    extends Layer2[F]:
  override def operation2(in: Input): F[Output] = for {
    context <- A.ask
    result  <- verify(context)
    output <-
      if (result)
        layer3.operation3(in)
      else
        // Better way of working with errors
        // is going to be shown in the `error-handling` part
        MT.raiseError(new RuntimeException(s"Could not verify context: $context"))
  } yield output

Impl3 does nothing as well.

Code

  final class Impl3[F[_]](layer4: Layer4[F]) extends Layer3[F]:
    override def operation3(in: Input): F[Output] = layer4.operation4(in)

And Impl4 does “the real work”.

Code

final class Impl4[F[_]: Monad](
    doRealWork: Input => F[Output],
    println: String => F[Unit]
)(using
    A: AskCtx[F]
) extends Layer4[F]:
  override def operation4(in: Input): F[Output] = for
    context <- A.ask
    output  <- doRealWork(in)
    _       <- println(s"{Context: $context} -> (Input: $in) -> (Output: $output)")
  yield output

As you can see the layers that don’t require the context for their logic has no knowledge about it - it is not passed there neither as a method argument nor a constructor parameter. Also this code can work both with Cats and ZIO.

Using Cats-Effect

IO[..] in cats-effect does not have a channel for passing a context value. So that we use Kleisli[IO, Context, ..] to add Context to IO.

Code

object CatsExample extends IOApp.Simple with Example:
  final case class User(name: String)
  override type Context = User
  override type Input   = String
  override type Output  = Int
  
  // Add Context to IO using Kleisli
  private type CtxIO[A] = Kleisli[IO, Context, A]
  
  private def doRealWork(in: Input): CtxIO[Output] = in.length.pure
  private def println(in: String): CtxIO[Unit]     = Console[CtxIO].println(in)
  private def verify(ctx: Context): CtxIO[Boolean] = ctx.name.nonEmpty.pure
  override def run: IO[Unit] =
    // type definitions are for clarity only
    val layer4: Layer4[CtxIO] = Impl4(doRealWork, println)
    val layer3: Layer3[CtxIO] = Impl3(layer4)
    val layer2: Layer2[CtxIO] = Impl2(verify, layer3)
    val layer1: Layer1[CtxIO] = Impl1(layer2)

    val validCtx: Context   = User("valid_name")
    val invalidCtx: Context = User("")
    val input: Input        = "non_empty_string"

    for
      _ <- Console[IO].println("Let's go!")
      ctxProgram = layer1.operation1(input)
      result <- ctxProgram.run(validCtx)
      //      result <- ctxProgram.run(invalidCtx)
      _ <- Console[IO].println(s"result: $result")
    yield ()

Using ZIO

ZIO[..] already have a channel for a context, so there is no need in using Kleisli[..]. Just construct a program and run it by passing a context value wrapped with ZLayer.succeed(..).

Code

object ZIOExample extends ZIOAppDefault with Example:
  final case class User(name: String)
  override type Context = User
  override type Input   = String
  override type Output  = Int

  private type CtxIO[A] = ZIO[Context, Throwable, A]
  private def doRealWork(in: Input): CtxIO[Output] =
    Exit.succeed(in.length)
  private def println(in: String): CtxIO[Unit] =
    Console.printLine(in)
  private def verify(ctx: Context): CtxIO[Boolean] =
    Exit.succeed(ctx.name.nonEmpty)

  override def run: ZIO[Any & ZIOAppArgs & Scope, Any, Any] =
    // type definitions are here for clarity only
    val layer4: Layer4[CtxIO] = Impl4[CtxIO](doRealWork, println)
    val layer3: Layer3[CtxIO] = Impl3(layer4)
    val layer2: Layer2[CtxIO] = Impl2(verify, layer3)
    val layer1: Layer1[CtxIO] = Impl1(layer2)

    val validCtx: Context   = User("valid_name")
    val invalidCtx: Context = User("")
    val input: Input        = "non_empty_string"

    for
      _ <- Console.print("Let's go!")
      // Construct a program that requires a context value
      ctxProgram = layer1.operation1(input)
      // Pass the context value to run the program
      result <- ctxProgram.provide(ZLayer.succeed(validCtx))
//    result <- ctxProgram.provide(ZLayer.succeed(invalidCtx))
      _ <- Console.printLine(s"result: $result")
    yield ()
comments powered by Disqus