Blog

The “suspend + receivers” style in Kotlin

02 Mar, 2023
Xebia Background Header Wave

The "suspend + receivers" style

The Arrow project, in particular starting with its 2.0 version, promotes a specific style of Kotlin to achieve composability of effects. Functional programming patterns inspire many concepts, but the specific implementation differs from Haskell’s or Scala’s typical way of doing things (monads, IO, and transformers). This post aims to clarify how this style works in Kotlin, compare it to other sibling languages, and discuss the main limitations.

This post assumes the reader is familiar with the transformer-style of effects. Good resources to understand that style are Scala with Cats, Practical FP in Scala, or the Book of Monads.

Warning: Context receivers are still experimental, per the 1.7.x series. There are still some rough edges when mixing them with suspend and inline, but the Kotlin team is clearly putting lots of effort toward making them stable.

Effects

Before diving into the technical portion, let’s discuss what we aim to achieve when talking about effects. The idea is to be explicit about the visible behavior of a function in its signature since that allows for more powerful compiler checks and analyses. By default, this is only the case for pure functions, which is a fancy way of saying "functions that only do computation." Take the following function:

fun add(x: Int, y: Int): Int = x + y

Here, the types are "honest;" we always get back a number, and such a number depends only on the input arguments. In Kotlin, we can easily sneak in other types of behaviors, which we often call side effects. We could print some value,

fun loggingAdd(x: Int, y: Int): Int {
  println("x = $x, y = $y")
  return x + y
}

or, going further, roll back the execution trace with an exception, but only at some random times.

fun crazyAdd(x: Int, y: Int): Int {
  if (Random.nextInt() == 42) throw WhatsHappeningException()
  return x + y
}

The goal is to shed light on those effects, making the potential behavior directly apparent. A function like the latter gets a type closer to

context(Raise<WhatsHappeningError>, Random)
fun crazyAdd(x: Int, y: Int): Int

This is again an extreme example, but think of those effects as network communication, access to the file system, or, more granularly, read the initial configuration or use the repository object for some part of the domain. The following signature is very explicit in telling that the function may or may not return a value (because of the nullable return type), it may use the user repository (so we expect some kind of database access), and there’s the possibility of an error (which is different from a non-existing value).

context(Raise<WhatsHappeningError>, UserRepository)
suspend fun getUserById(id: UserId): User?

The question is: how do we encode those effects using the type system? This is where the Kotlin and Haskell / Scala ways significantly differ.

Two Kotlin built-in features

To model and execute effects in a performant fashion, we recommend using coroutines and context receivers profusely in Kotlin. This section briefly introduces each feature and discusses its role in the overall pattern.

Coroutines are suspendable functions. From the developer’s point of view, these functions are marked with the suspend keyword:

suspend fun getUserById(id: UserId): User?

The Kotlin language knows how to combine several suspended functions without any intervention from the developer. In other words, coroutines are transparent for developers, except for the requirement of marking functions that use them. We can even mix suspended functions with regular functions, as with the .name field access here.

suspend fun getUserName(id: UserId): String? =
  getUserById(id)?.name

Although coroutines are often explained together with the concurrency mechanisms from kotlinx.coroutines, those are separate ideas. Coroutines allow fine-grained control over the computations within a suspend block. Concurrency is one form of exercising this control, deciding when and in which threads computations will happen. But this is not the only one:

  • Inikio is a library that leverages suspend to create "imperative" syntax for domain-specific languages;
  • Arrow provides computation builders where you can write functions over Either, Result, or nullable types without having to check for the error path at each step.

The ability to control computation in that way comes from the transformation performed by the compiler into continuation-passing style. Long story short, our functions above really take an additional argument — the continuation — which is "fed" the resulting value of the function.

fun getUserById(id: UserId, k: Continuation<User?>)

fun getUserName(id: UserId, k: Continuation<String?>) =
  getUserById(id) { user -> user?.let { k.resume(it.name) } }

At a conceptual level, you can think of Continuation<A> as simply a function (A) -> Unit. This means the caller can change how the function "returns," but changing the continuation passed as parameter k.

The second feature in the Kotlin language required for the proposed style is context receivers, which were introduced in version 1.6.20. For an in-depth discussion of this new feature, you can check this talk or its corresponding slides. Very briefly, context receivers can be thought of as implicit parameters, or as a sort of built-in dependency injection. Let’s say we define an interface for our user repository,

interface UserRepository {
  suspend fun getUserById(id: UserId): User?
}

Then we can make that repository available in another function by putting it as part of the context declaration.

context(UserRepository)
suspend fun getUserName(id: UserId): String? =
  getUserById(id)?.name

This signature could also be written as suspend context(UserRepository) fun, but we’ve decided to write the contexts first to highlight them.

The example above shows that if you have a UserRepository in your context, then you can call any other function that requires that context — here getUserById — without further ceremony. This composability extends as much as you want; the getUserName may again be called by another function with the same context. At some point, we need to inject the actual implementation, though; this is done using the with function,

fun main() {
  createDbConnection().use { db ->
    with(DbUserRepository(db)) {
      println(getUserName(UserId(1)))
    }
  }
}

In this case, the implementation connects to a database, but we could also inject a different in-memory object for testing. This is where the dependency injection idea comes into play.

Context receivers are a generalization of an older Kotlin feature called extension receivers. Although its original use was to "extend" the functionality of a class while preserving the .function syntax for access members, they’ve also been used to describe scopes where certain additional functionality is available. Here’s the (simplified) signature of the async function that spawns a computation in a new lightweight thread:

fun <T> CoroutineScope.async(block: suspend CoroutineScope.() -> T): Deferred<T>

This signature states that async may only be called within a CoroutineScope block, and that the spawned computation represented by the block argument also gets access to another CoroutineScope, so it can launch other jobs itself. If we "modernize" the signature using context receivers,

context(CoroutineScope)
fun <T> async(block: context(CoroutineScope) suspend () -> T): Deferred<T>

we can see that those scopes are nothing more than effects! This notion is visible in several places in the Arrow library:

  • A ResourceScope brings the ability to acquire and release resources correctly.
  • STM introduces the concept of transactional variables, whose concurrent access is protected.

To perform their duty, resourceScope (the "runner" for ResourceScope) and atomically (the "runner" for STM) need control over the computation. That’s precisely where the coroutine system becomes necessary.

The emphasis on context receivers comes from their ability to compose, as opposed to the other extension receivers. If a function requires both resource handling and a user repository, we can easily express this using the former

context(UserRepository, ResourceScope)
fun whoKnowsWhatItDoes(name: String)

but there’s no (simple) way with the latter. We hope this composability encourages developers to create more and finer-grained scopes/effects to be as concrete as possible about the behavior of the functions.

No monads, no higher-kinded types

Traditionally, the ability to control execution in a similar way to what the coroutine system in Kotlin offers was offered by a monadic interface. More concretely, monads provide a function to combine one sequence of steps with the next step to be performed; this operation is known as flatMap or bind (>>= if you’re a Haskeller).

fun <A, B> M<A>.flatMap(next: (A) -> M<B>): M<B>

Since the implementor of each particular monad controls the behavior of flatMap, any additional control can be injected at that point. For example, the nullability monad short-circuits any other behavior when a null is produced in one of the steps.

fun <A, B> A?.flatMap(next: (A) -> B?): B? = 
  when (this) {
    null -> null
    else -> next(this)
  }

One important monad for both Haskell and Scala is the IO monad, which marks the code as "impure," that is, as having potential side effects like writing or reading from a device. Note that, in the case of Scala, there’s not a unified IO monad; we have Cats Effect and ZIO as prominent examples. In Kotlin the same can be achieved with suspend; to start running a suspended computation, we need to call it from main or from runBlocking, which amounts to the same guarantees as found in Haskell and Scala.

Alas, code written with monads tends to break under the weight of all the flatMaps in the way and the nesting associated with them.

doOneThing().flatMap { x ->
  doAnotherThing(x).flatMap { y ->
    // imagine if your code was even longer!
    return x to y
  }
}

For that reason, both Scala and Haskell provide specialized syntax to work with monads — for comprehensions for the former, and do notation for the latter. The same code can be written much more clearly as

for {
  x <- doOneThing()
  y <- doAnotherThing(x)
} yield (x, y)

This is a step in the right direction, but not perfect. In particular, monadic notation enforces a strong divide in style between monadic and non-monadic code. If you even need to introduce some effect to a function, you are forced to rewrite it from the "pure" subset of the language to the "monadic" one. This problem extends not only to blocks of code, but also to the operators available to them. The famous traverse operation

def traverse[F[_], A, B](x: List[A], f: A => F[B]): F[List[B]]

is nothing more than your regular list map, but where f may incorporate some effects. Once again, the solution is good. But, in a perfect world, you wouldn’t ask developers to learn two sets of abstractions if one could do.

The signature of traverse shows another essential part of the monadic style of effects, namely higher-kinded types. Since monads are types that take other types as parameters — like Option or Result — to provide utility functions that work over different monads, one must have a way to refer to those types. In Scala, this is made explicit by the number of underscores after the type name, like F[_] above; in Haskell, this is not made explicit in the signature, but is still checked by the compiler. Even decades after their inception, higher-kinded types are still considered an advanced feature and haven’t made their way into mainstream programming languages.

Kotlin doesn’t require higher-kinded types to provide similar abstraction because the same functions that work over "regular" functions, like map, keep working when used in the coroutine world. Kotlin language designers have found, in my opinion, a great point in the design space in which they can offer the same abilities, but without the cost of such a complex feature.

Composition of effects

Monads have some inherent problems with composition, which in practice means that you cannot easily create a new one from the effects of two of them. There are several techniques to overcome this problem, but the main one is monad transformers. Instead of being a monad itself, a transformer "wraps" another monad, adding some additional capability. For example, the composition of "we keep a stateful DbConnection" with "we require a Config in context" with "we may perform I/O actions" is represented as follows:

StateT[DbConnection, ReaderT[Config, IO], User]

Note that a transformer like StateT has a higher-kind than a regular State monad. The latter takes one "regular" type as a parameter, whereas the former takes a monad (a type with a hole) as a parameter. This is one of the most challenging mountains to climb in order to understand how effects work in Haskell and Scala.

We propose to move away from this style in Kotlin, and simply enlarge the context required for a function. Context receivers compose without any further instruction from the developer; in that sense, they work similarly to algebraic effects, like those in effectful or polysemy. More concretely, you have no fixed order of injection of contexts, as opposed to monad transformers.

Errors everywhere!

One interesting application of this style is an ergonomic interface for errors. In Arrow 2.0 Raise is the name of the scope/effect of short-circuiting with an error of type E. A function with signature

context(Raise<E>) () -> A

can be executed into an Either, a Result (is E is Throwable), and many other similar types. However, nothing stops you from having more than one Raise in your context if you want to be completely explicit about errors.

context(Raise<DbConnectionError>, Raise<MalformedQuery>)
suspend fun queryUsers(q: UserQuery): List<User>

Context receivers play well with variance, so injecting a receiver for Raise can simultaneously work for both elements in the context. Compare this with EitherT[DbConnectionError, EitherT[MalformedQuery, IO], List[User]], where you need to handle each of the EitherT independently and in the order given by the monadic stack.

Limitations

Continuations are really powerful; in fact you can use Cont to emulate any other monad. However, this doesn’t speak of the ergonomics of using such emulation, so that argument is not really valid for this discussion. Furthermore, Kotlin imposes some limitations on how coroutines can be used, limiting the expressivity of the proposed style.

Kotlin’s continuations are one-shot, meaning they can only be called once. This is enough to model many interesting effects. For those like ResourceScope or CoroutineScope, you want to call the continuation normally. For those like errors, you may not call them, but that’s also allowed. This limitation kicks in, though, if you try to implement the list (also known as non-determinism) monad in Kotlin. To give a specific example, in Arrow Analysis, we needed one such abstraction, and the solution requires an explicit flatMap whenever you may run the same code more than once.

We shouldn’t be too quick to flag one-shot continuations as a mistake. If you are guaranteed to call your code at most once, you can use all kinds of mutable state to optimize how they work. The cost of allowing several calls means actually creating a whole closure at each suspend call, which is not always negligible. In my opinion, the Kotlin designers found a great point in the design space again, knowing that any performance hit would be taken badly by the JVM community they are targeting.

The lack of higher-kinded types in Kotlin means that some patterns cannot be abstracted. We already discussed how some of them are no longer needed — like traverse — but others are still relevant. In particular, higher-kinded types enable abstracting away the typical structure of a monad in the Operational or Free monads. This is impossible in Kotlin, so we must resort to code generation, as done with the Inikio KSP plug-in.

Conclusion

In this post, we describe a style of programming in Kotlin based on continuations + context receivers. This style is idiomatic for Kotliners and improves upon the "monadic style" available in other functional languages. Repeat after me: no more flatMap!

Xebia Functional đź’™ Kotlin

We, the functional programming division at Xebia, are great fans of Kotlin, exploring the many possibilities it brings to the back-end scene. We’re proud maintainers of Arrow, a set of companion libraries to Kotlin’s standard library, coroutines, and compiler; and provide Kotlin training to become an expert Kotliner. If you’re interested in talking to us, you can use our contact form, or join us on the Kotlin Slack.

Questions?

Get in touch with us to learn more about the subject and related solutions

Explore related posts