Warp WebSockets and Tokio Actors

I use warp a lot in my projects and frequently need to expose some service through websockets. The warp project has a good example on how to use web sockets, and it is pretty easy to use. However, I think using actors here makes things a little cleaner and easier to grok.

Here is a small example on how to use actors with Warp websockets. To follow along you have to include the following dependencies to your Cargo.toml:

[dependencies]
tiny-tokio-actor = "0.2"
tokio = { version = "1", features = ["full"] }
futures = "0.3"
tokio-stream = "0.1"
uuid = { version = "0.8", features = ["v4"] }
warp = "0.3"
env_logger = "0.9"
dotenv = "0.15.0"

The Echo Actor

To create the actor that will echo back any message it gets, I use a small library I created on top of tokio called tiny-tokio-actor. The actor is very simple:

To copy/paste use the complete code at the bottom

Besides the actor we also create a struct called ServerEvent for which we implement the SystemEvent trait. This will be used by the actor system which we will see later. For now let’s focus on the actor.

The actor holds one property, and that is an unbounded sender for websocket messages. The actor will be sending back the echo messages over this sender.

Next we create the message that our EchoActor will receive:

As you can see, this message is just a wrapper around a warp websocket message. The actor handler that handles this message simply logs a debug message, and sends it back using the sender property of the actor.

Our echo service is done! All we need to do now is wire it to the rest of warp.

Warp Server

To run our server we will need to create a main function with the tokio runtime. In this main function we will also start our actor system:

Next we create the warp server that can serve a websocket connection:

We have the warp route take our actor system, the remote address of the connecting client, the warp websocket handler, and pass it all to the start_echo function.

In the start_echo function we will create an EchoActor for the connected client, using the client’s address as part of the actor name. As our actor takes a UnboundedSender<warp::ws::Message> we will need to construct the necessary channel that will link to the outgoing websocket sink. To get the sink, we split the websocket connection:

We first split the websocket into a ws_out sink, and a ws_in stream to receive messages from the connected websocket client. The ws_out sink is used to send messages back to the client.

We create an unbounded channel that we link to the ws_out sink. The sender of this channel will be used to send messages to over the ws_out sink back to the websocket client.

Next we create the EchoActor and pass in the sender of the channel we just created. We run this actor on our actor system using a unique name generated from the remote client’s address (hostname & port). If for some reason we cannot retrieve the remote address we just use a randomly generated UUID.

Now we can take messages from the ws_in stream and send them to our EchoActor:

We loop over messages received from the ws_in stream until we get an error. If an error happens we log it and break out of the loop. When the loop ends, we stop our EchoActor too.

Thats it! When you run the server you should be able to connect to it with something like websocat, e.g.:

$ websocat ws://127.0.0.1:9000/echo

The complete code:

This code is also part of the tiny-tokio-actor’s project in the examples folder.