Combining ZIO and Akka to enable distributed FP in Scala

The ecosystem around functional programming in Scala is getting richer and richer: cats, zio, monix, http4s, doobie, refined, monocle, etc. There are tons of excellent FP libraries, covering almost everything you need. Everything, really? There’s one area that’s actually lacking: distributing work across multiple nodes in a cluster. You either have to rely on external infrastructure (e.g. Zookeeper, Kafka), or you can use Akka and its actor-based architecture but you’re leaving the FP world.

Is there a way to get the best of both worlds? Until a viable alternative emerges from the FP community, could we leverage Akka distributed features while keeping our code purely functional? The answer is yes, thanks to zio-akka-cluster, a ZIO wrapper around Akka Cluster that lets you use the distributed features of Akka without forcing you to use actors.

Building a distributed chat

Let’s say we want to implement a chat application. Users give their name, join a chatroom and can send messages that are propagated to other users. If your server is a single node, this is really simple but what if you want multiple nodes for redundancy (no downtime if a node is down) and scalability (add more nodes if your load increases)? If User A joins a chatroom X through node 1, and User B joins the same room through node 2, you want them both to see each other messages.

Akka has two features that help addressing this problem in an easy way:

• Akka Distributed PubSub lets you publish and subscribe to topics across a cluster of nodes. We will use that to listen to the messages published to a chatroom.
• Akka Cluster Sharding ensures that a given entity (identified by an ID) is alive in a single place across a cluster of nodes. You can send messages to this entity without knowing on which node it resides (this is called location transparency). We will use this concept to represent our chatroom: we want a given chatroom to be loaded in memory on a single node at a time.

Here’s a simple diagram to visualize what we are trying to achieve:

Let’s get started

The first thing we need when using Akka is an ActorSystem, which is the basic structure needed to use actors. Even though our application won’t use actors directly, it is still required. By default, Akka will look at the resources for a configuration file called application.conf, but you can also pass your configuration manually.

Creating an ActorSystem is a side effect, so we’ll need to wrap it in ZIO’s Task. It is also recommended to terminate the actor system in a clean way when shutting down your program, especially when you have a cluster: it will alert the other nodes that the node is leaving the cluster, so that operations like shard rebalancing can be performed right away. ZIO has a data structure called Managed that ensures a release action will be called when the resource usage is finished, whether the process ended successfully or with an error.

Chatroom business logic

A chatroom is able to process 3 types of events: someone joining, someone leaving and someone posting a message. We’ll first create an ADT to represent these messages. We will call our base message type ChatMessage.

We want each chatroom to be sharded, which means it will be loaded only once on the cluster. zio-akka-cluster lets you set it up easily by creating a Sharding[Msg, State], Msg being the type of message it can handle, and State the type of the entity state, which we can read and modified every time we handle a message. For our example, we’ll use ChatMessage as our message type and List[String] as our state: the list of connected users.

To create a Sharding object, we need to define a behavior using the following signature:

Msg => ZIO[Entity[State], Nothing, Unit]

This means that we need a function from Msg to a ZIO effect with the following properties:

• Provides an Entity[State]. This interface gives you access to the entity state, the entity ID and also lets you stop the entity. The state is held in ZIO’s Ref which means modifying it can be done in a purely functional way.
• Does not fail: you need to handle errors within the behavior, so that the sharded entity does not crash. In this example, we will use .ignore to ignore errors.
• Returns Unit.

We create the Sharding object by passing this behavior to Sharding.start, along with the name of the entity type, here Chat.

Our business logic will be the following: we will update the entity state when we receive Join and Leave events, and we’ll push an update to all participants to inform them. In case of a Message events, we’ll simply broadcast it to all participants.

For that, we’ll use a Publisher from zio-akka-cluster. It can be created by calling PubSub.createPublisher[Msg] and lets you publish messages of type Msg to a given topic that can be subscribed to from any other node in the cluster. We will use the chatroom name as the topic.

A proper implementation would probably have a more complex state (e.g. the list of all messages) and would use some persistence layer to save and recover the state when the entity is created.

Building the Chat Client

The client needs to do two things: send events to the chatroom, and subscribe to the chatroom events. For the sake of simplicity, I’m using the console to input and print messages.

To send events to the chatroom, we can re-use our Sharding object and use its send method, which requires the entity identifier (we will use the chatroom name) and the message of type ChatMessage.

To subscribe to the chatroom events, we will need to create a Subscriber. Note that you can use PubSub.createPubSub to get both a Publisher and a Subscriber at the same time. A Subscriber has a listen method that, for a given topic, returns a ZIO Queue that will be populated by the messages published to that topic. We’ll print all the received messages to the console.

Note how we use ZIO .forever to repeat an effect until it gets interrupted.

And that’s it! We can now run this snippet of code on multiple nodes, and chatrooms will be arbitrarily distributed across our cluster. All messages to a given chatroom will be redirected to the same entity on one of the nodes, and messages published by that chatroom can be received by clients located on any node.

You can find the full code in this repository, with instructions for running a 2-nodes cluster.

A couple comments to finish:

• Publisher, Subscriber and Sharding being traits, you can easily build local implementations for your unit tests so that you don’t rely on Akka at all when testing your application.
• Akka messages are serialized when they are sent across the network. By default, Java serialization is used but it is not recommended to use it in production. See Akka Documentation to see how to provide your own serializer.

In this example, we demonstrated how easy it was to distribute a simple application over a cluster with the help of zio-akka-cluster. We also showed some examples of using ZIO data structures such as Managed, Queue or Ref and how they make resource management and concurrency easy.

The library offers other features, such as the ability to listen to cluster events and detect when nodes are unreachable. I hope you will find it useful!