Building a Reactive, Distributed Messaging Server in Scala and Akka with WebSockets

If the users of your app have clients of their own they want to communicate with, you will need a messaging service. Here is how to implement one using the WebSocket API — a standard API for client-server communications — on top of Akka HTTP, streams, and typed actors in Scala

Motivation

The use case for this server is an app where the primary user of the app (the User) wants to communicate with a set of clients (the Client/s). The Clients do not use this app, and the User wants to be able to respond with real messages as opposed to computer-generated ones. Here are some examples:

• A tech support or home maintenance business wants to coordinate a good time to send a technician
• A voter outreach app needs to send messages to many phones and be able to respond to their replies

While sending messages from the User to Clients is easy with standard REST HTTP requests, receiving the replies requires either (A) continuously polling the server to check for new messages or (B) a two way connection. Since (B) is much more efficient we will use this approach.

We are going to build this server using WebSockets, a standard web API for opening a two-way interactive communication session between the server and client. For our server we will use Akka HTTP, a library for building web servers built on top of Akka Streams that supports both Java and Scala. We will be using Scala.

In this article we will focus on the actual WebSocket routing and behavior and abstract over the details of the actual client communications protocol by calling it ClientMessagingService . In my own app I use the Twilio API to send SMS messages to clients, but this can be anything.

Tools

We will be using the following tools:

• WebSocket is a standard web API. Mozilla describes it as a “technology that makes it possible to open a two-way interactive communication session between the user’s browser and a server.” This means that the server can send messages from clients as they come without waiting for the client to ask to see them. That is the chief advantage of using WebSocket as opposed to stateless HTTP requests. WebSocket, like HTTP, supports SSL/TLS encryption (http:// is to ws:// as https:// is to wss://) for security.
• Akka¹ HTTP is a library for building reactive web servers in Java and Scala. It is built on top of Akka Streams, a library for processing asynchronous distributed streams, which is itself built on top of Akka Actors, a library for building reactive distributed systems using the actor model. Akka HTTP is designed to be highly reactive, that is, resilient, responsive, elastic, and message-driven. It supports both Java and Scala, but here we will use Scala. Akka has a typed and untyped implementation. Using Akka typed guarantees strongly typed, explicitly configured protocols for actors and streams, so we will be using Akka typed. The Akka ecosystem is maintained by Lightbend.
• Scala is a multi-paradigm, statically-typed language with strong first class support for functional programming concepts (e.g., higher-order functions, immutable data structures, etc). Scala typically runs on the JVM (although JavaScript compilation is officially supported and LLVM is experimental). Scala was created and is supported by EPFL.

System Design

1. WebServer: the runnable App that exposes the routes and sets the rules for spinning the server up and down
2. ChatSessionMap: a map that keeps track of each User’s chat session. It has an entry for each User that sends it a connection request.
3. ChatSession: the abstraction of the User’s connected chat session, contains instructions of how to handle messages that come in either from the User over the WebSocket connection or from the Clients over whatever is specified by ClientMessageService. There is one instance per connected User.
4. ChatSessionActor: the actor that is embedded in the ChatSession. It’s a reactive agent that responds to messages sent to it from the User or Client. The actor model is a design pattern for concurrent/distributed computing; you can learn more about it and about Akka, the main actor library for Scala, here: https://doc.akka.io/docs/akka/current/typed/guide/actors-intro.html. There is one instance per chat session.
5. WebSocket: an anonymous actor that surfaces inside the session’s stream (more on that later) and materializes messages sent to it.
6. ClientMessageService: a service to communicate with Clients. Out of scope for this article, but should be able to communicate with a client based on some protocol. I am building one on top of the Twilio API for SMS communications.

Programming the Server

Note to beginners: If you need an introduction to Scala, I suggest: https://docs.scala-lang.org/getting-started/intellij-track/building-a-scala-project-with-intellij-and-sbt.html.

For an Akka HTTP introduction please refer to: https://doc.akka.io/docs/akka-http/current/introduction.html.

For Akka Streams here: https://doc.akka.io/docs/akka/current/stream/index.html?_ga=2.44841895.2097526044.1596922028-1131448099.1591641463.

Ensure you have all the required dependencies in your build file:

• Akka HTTP
• Akka HTTP Core
• Akka Streams Typed
• Akka Actor Typed

In steps, we will walk through developing the message server in a bottom-up approach.

Step I: Akka HTTP server

The first step is to create a basic web server to test out Akka HTTP:

This is a basic web server that welcomes you when you try to access it. When you run it and hit curl http://localhost:8080 on your command line it will return “Welcome to messaging service.

Now that we have a working server, let’s see how to handle WebSocket requests

Step II: Handling WebSocket Connections

Akka HTTP has first class support for WebSocket. It provides a method handleWebSocketMessages that abstracts over the handling of WebSocket connections, messages, and disconnections by treating the whole process as a stream with one inlet and one outlet. The first message to come through the inlet of the stream is the connection request², and when the stream stops, it disconnects. The outlet of the stream goes back to the client.

As a first example for how to handle WebSocket messages we will create a WebSockets route and call it/affirm . Its job is to affirm everything we say: so if we send it “Hello!”, it would send to the client “You said Hello!”.

Note: this would be a good time to install websocat, a command-line tool for connecting to WebSocket server. We can use it to make sure our server behaves as intended.

You can run the server and test the WebSockets\ endpoint with websocat -t ws://localhost:8080/affirm , which will open a connection where you can send and receive messages.

We did two things:

1. We added the WebSockets endpoint /affirm by defining the path and calling handleWebSocketMessages
2. We composed the new endpoint with the endpoint we created earlier using the ~ operator, which is used to compose two endpoints.

This is a good time to talk about Akka Streams in more depth As we mentioned before, handleWebSocketMessages is a method that requires an Akka Streams blueprint to know how to handle messages, and that this stream needs one inlet and one outlet. The core of Akka Streams is a system for constructing computational graphs for streaming data out of the following three components:

• Source[T]: describes a source of stream of type T (like a spring)
• Sink[T]: describes a way for streaming data to disappear or to exit the stream (like a lake)
• Flow[T, R]: describes the lay of a stream with incoming and exiting data, like a section of a river

There are other components, but they are mostly for gluing together these three in fan-in and fan-out operations, for example, joining several streams together into a single stream (fan-in) or separating one stream into separate ones (fan-out).

To handle the WebSocket connection we need to define a Flow[Message, Message] component, (which can be shortened to just Flow[Message] ). As you can see in the imports section, Message here refers to a WebSocket message; the inlet is a message from the client (“Hello!”) and the outlet is a message to the client (“You said Hello!”). In this flow, we just append “You said ” to every message that comes in.

Step III: Sketching out the Program

We built the basic concepts; now we start implementing the entire system. In this step, we will build out all the components except for the Flow[Message] that handleWebSocketMessages takes as a parameter, because it requires a deeper consideration of Akka Streams.

ChatSessionMap

First, in a file ChatSessionMap.scala, we create a basic map containing all sessions, where each session corresponds to a User:

This simple object just contains a record of every chat session to make sure connect requests to establish a WebSocket connections are handled correctly.

ChatSessionActor

The ChatSessionActor is a typed Akka actor responsible for doing the heavy lifting: it receives messages intended to the Client and sends them to ClientMessageService to deal with, also receives messages from this service and routes them to an anonymous typed actor serving as a stand-in for the WebSocket:

In Akka Typed, actors are implemented as a function with a return type of Behavior[T], where the type parameter is the type of messages they can respond to. In this case, the actor responds to messages from the user (UserMessage) by sending to the Client through ClientMessageService; to messages from the client by forwarding to the WebSocket actor, passed through as a parameter; and to Connected by writing the WebSocket reference to the one received via message.

Here is the protocol file:

Step IV: ChatSession and Route

The core and most complex component of the server is the ChatSession object. A ChatSession contains the following:

• A session actor which performs the work of routing messages around the system
• A webflow (Flow[Message]) which forms the argument to handleWebSocketMessages and is responsible for passing info in and out of the web socket

The session object is quite complicated, so it’s useful to remember these two points of Akka stream:

• A stream component is not a data transformation, rather a blueprint for a data operation
• A flow does not need to be connected, it just needs to expose an inlet and outlet. Using waterways as an example, a valid flow would have a river as an outlet, and inside the flow, flow to a lake. An unrelated spring could have its outlet as the flow’s outlet

Because the session actor exists outside an actor system, we have to ask the system for it, so we get it as a future.

To construct the flow, we use the Flow.fromGraph method, which takes a graph instance as a parameter. We pass in an ActorRef parameter; parameters inside streams are resources that live inside the stream. In our previous example of a flow consisting of a river flowing to a lake and a spring flowing out, the spring is a resource belonging to the flow. In this case we pass an actor as a source, meaning messages sent to that actor materialize inside the stream.

Here is a construction of the flow, component by component:

1. We define a Flow[Message, CoreChatEvent] item that maps any WebSocket message into a message from the user according to the protocol
2. We define a Flow[WebSocketEvent, Message] that takes a message sent to the WebSocket actor reference and transforms it into a WebSocket message
3. We define an internal ActorSink[CoreChatEvent] by assigning its reference to that of the session actor, meaning data conforming to the CoreChatEvent protocol can get routed into this sink as a message
4. We materialize the internal ActorSource[WebSocketEvent] actor reference resource, which acts as a reference to the WebSocket stream, and map it into a connection request so we can send it to the session actor via a Connect request as per its protocol
5. We use Merge[CoreChatEvent](2) , a fan-in block allowing two inlets respecting the protocol to merge into one outlet
6. We route all the components together using the ~> operator: messages from the WebSocket flow and the internal actor source flow into the merge, which flows into the session actor sink reference. The original internal actor source reference routes to the flow that transforms messages from its protocol (WebSocketEvents ) back to the WebSocket protocol.
7. Finally, we expose the inlet of webSocketSource as the flow’s inlet and the outlet of webSocketSink as the outlet to create a flow.

For a visual aid, refer back to the system diagram above. Note that because the actors shown in the cycle are out of the stream, they are references from inside the stream as source and sinks, but the connections between them are done in the session actor.

Finally we define an additional route for message/{userId} with pattern-matching on the argument:

You may have to import additional dependencies, and don’t forget to compose the new messageRoute with the existing routes in the HTTP server binding the same way we did with /affirm.

If we run the server now, we can test it by running websocat -t ws://127.0.0.1:8080/message/{id-goes-here} . Unless we implemented the service that messages the Clients, our messages will go nowhere and we cannot get replies, but the print statements show that the flow works as intended. If you want to see replies you can program an actor to send messages to the session actors on a schedule.

Conclusion

This is skeleton code for a distributed messaging server using Akka, Scala, and the WebSocket API. It works by routing messages through a stream to an actor that handles their distribution.

While it would work as is under certain assumptions (e.g., the User is always connected), there are a number of necessary improvements:

• Changing the string processing of name and number to JSON serialization to make it easier to parse on the front end
• Adding authentication requirements instead of authenticating by path argument
• Securing the server with TLS encryption
• Refactor the chat sessions map to be immutable and cluster-able so it doesn’t block horizontal scaling
• Delete inactive chat sessions so they don’t clog up the map
• Add a persistence layer so messages sent when the User is not connected do not get lost
• Add logging and telemetry to empower data-driven decisions about the evolution of the chat system, such as which parts are performance slowdowns or take up too much memory

[1] Akka is a palindrome of “AK” which stands for Actor Kernel

[2] The connection request is a standard HTTP request with specific headers that request that the connection be upgraded to a WebSocket