Build a task ventilator with PUSH/PULL and DEALER/ROUTER

For Rust developers who know Tokio and want to try ZeroMQ patterns without the C library

You will build a small job pipeline over TCP on localhost: a ventilator hands numbered tasks to two workers over PUSH/PULL, the workers send their answers back over DEALER sockets to a ROUTER collector and the whole thing runs and shuts down cleanly from one cargo run.

Before you start

You need Rust with Cargo. rzmq comes from crates.io:

rzmq = "0"

Set up the project#

Start a new binary crate.

cargo new ventilator
cd ventilator

Add rzmq and Tokio. rzmq runs on the Tokio runtime, so you need both. Your [dependencies] section should look like this:

Cargo.toml

[dependencies]
rzmq = "0"
tokio = { version = "1", features = ["full"] }

This tutorial uses the default features, which cover everything below. Every socket in it talks TCP on 127.0.0.1, so the same code would work with the pieces split across processes or machines. Here they all live in one process as Tokio tasks so you can run it in one go.

Send tasks from PUSH to PULL#

Every rzmq program starts with a Context. You create sockets from it and call term() on it at the end to close them all. A PUSH socket sends messages out to whoever is connected and a PULL socket receives them. The ventilator is the stable end, so it binds. The worker connects to it.

Replace src/main.rs with this:

src/main.rs

use rzmq::{Context, Msg, SocketType, ZmqError};

const TASKS: &str = "tcp://127.0.0.1:5557";

#[tokio::main]
async fn main() -> Result<(), ZmqError> {
    let ctx = Context::new()?;

    let ventilator = ctx.socket(SocketType::Push)?;
    ventilator.bind(TASKS).await?;

    let worker = ctx.socket(SocketType::Pull)?;
    worker.connect(TASKS).await?;

    for task in 1..=3u64 {
        ventilator.send(Msg::from_vec(task.to_string().into_bytes())).await?;
        println!("sent task {task}");
    }

    for _ in 0..3 {
        let msg = worker.recv().await?;
        let text = String::from_utf8_lossy(msg.data().unwrap_or_default());
        println!("worker got task {text}");
    }

    ctx.term().await?;
    Ok(())
}

A message is a Msg. Msg::from_vec takes ownership of a Vec<u8> and msg.data() hands the bytes back as an Option<&[u8]>. rzmq doesn't care what the bytes are, so this program just sends the task number as text.

Run it:

$ cargo run
sent task 1
sent task 2
sent task 3
worker got task 1
worker got task 2
worker got task 3

All three sends finish before the worker reads anything. PUSH queues messages for the connected peer and the PULL side holds them until you call recv, so the two ends don't have to move in lockstep.

Add a second worker#

PUSH spreads messages round-robin across its connected peers, so adding workers is the way you scale this pattern. Each worker now gets its own PULL socket and runs in its own Tokio task.

PUSH can only share work among peers that have finished connecting. If you start sending right after the connect calls, the first worker to finish its handshake can take every task before the second one shows up.

To avoid that, the ventilator watches its own socket events. monitor_default() returns a channel of SocketEvents. A HandshakeSucceeded event arrives once per peer that has finished connecting. The program waits for one per worker before it sends anything. Open the monitor before calling bind so no events are missed.

Replace src/main.rs:

src/main.rs

use rzmq::socket::SocketEvent;
use rzmq::{Context, Msg, SocketType, ZmqError};

const TASKS: &str = "tcp://127.0.0.1:5557";
const WORKERS: usize = 2;
const TASK_COUNT: u64 = 6;

#[tokio::main]
async fn main() -> Result<(), ZmqError> {
    let ctx = Context::new()?;

    let ventilator = ctx.socket(SocketType::Push)?;
    let events = ventilator.monitor_default().await?;
    ventilator.bind(TASKS).await?;

    let mut workers = Vec::new();
    for id in 1..=WORKERS {
        let tasks = ctx.socket(SocketType::Pull)?;
        tasks.connect(TASKS).await?;
        workers.push(tokio::spawn(async move {
            for _ in 0..TASK_COUNT / WORKERS as u64 {
                let msg = tasks.recv().await?;
                let text = String::from_utf8_lossy(msg.data().unwrap_or_default());
                println!("worker-{id} got task {text}");
            }
            Ok::<(), ZmqError>(())
        }));
    }

    let mut ready = 0;
    while let Ok(event) = events.recv().await {
        if let SocketEvent::HandshakeSucceeded { endpoint } = event {
            println!("worker connected from {endpoint}");
            ready += 1;
            if ready == WORKERS {
                break;
            }
        }
    }

    for task in 1..=TASK_COUNT {
        ventilator.send(Msg::from_vec(task.to_string().into_bytes())).await?;
    }

    for worker in workers {
        worker.await.expect("worker task panicked")?;
    }
    ctx.term().await?;
    Ok(())
}

PUSH deals messages round-robin across connected peers, so six tasks split three and three between two workers. Each worker takes its three and returns. main waits for both workers and only then calls ctx.term(), which closes the sockets once the work is done.

Run it:

$ cargo run
worker connected from tcp://127.0.0.1:56363
worker connected from tcp://127.0.0.1:56364
worker-1 got task 2
worker-1 got task 4
worker-1 got task 6
worker-2 got task 1
worker-2 got task 3
worker-2 got task 5

Your port numbers will be different. The order of the lines and which worker gets the odd tasks can also change from run to run, but each worker gets three.

Send results back over DEALER and ROUTER#

The workers need somewhere to send their answers. A ROUTER socket fits well for the collector. Every message it receives starts with an extra frame naming the peer it came from. When it sends, it reads the first frame to pick which peer gets the message. Each worker connects to it with a DEALER socket, which is the asynchronous partner for ROUTER and can send messages made of several frames.

Each worker gives its DEALER a name with the ROUTING_ID option before connecting. That name is what shows up as the first frame on the ROUTER side. A worker sends its result as two frames, the task number and its square, using send_multipart. The collector gets three frames from recv_multipart (name, task and square) and replies ok by sending the name frame back with the reply behind it. The worker waits for that reply before pulling its next task.

The worker also sleeps for 20 * n milliseconds to fake some work, so bigger tasks take longer. As in the previous step, each worker handles its share of the tasks and returns.

Replace src/main.rs with the full program:

src/main.rs

use rzmq::socket::{ROUTING_ID, SocketEvent};
use rzmq::{Context, Msg, Socket, SocketType, ZmqError};
use std::collections::BTreeMap;
use std::time::Duration;

const TASKS: &str = "tcp://127.0.0.1:5557";
const RESULTS: &str = "tcp://127.0.0.1:5558";
const WORKERS: usize = 2;
const TASK_COUNT: u64 = 8;

fn text(msg: &Msg) -> String {
    String::from_utf8_lossy(msg.data().unwrap_or_default()).into_owned()
}

async fn worker(tasks: Socket, results: Socket) -> Result<(), ZmqError> {
    for _ in 0..TASK_COUNT / WORKERS as u64 {
        let msg = tasks.recv().await?;
        let n: u64 = text(&msg).parse().expect("task is a number");
        tokio::time::sleep(Duration::from_millis(20 * n)).await;

        results
            .send_multipart(vec![
                Msg::from_vec(n.to_string().into_bytes()),
                Msg::from_vec((n * n).to_string().into_bytes()),
            ])
            .await?;
        results.recv().await?;
    }
    Ok(())
}

async fn collector(results: Socket) -> Result<BTreeMap<String, u64>, ZmqError> {
    let mut per_worker = BTreeMap::new();
    let mut total = 0;
    for _ in 0..TASK_COUNT {
        let frames = results.recv_multipart().await?;
        let (who, n, square) = (text(&frames[0]), text(&frames[1]), text(&frames[2]));
        println!("{who}: {n} squared is {square}");

        total += square.parse::<u64>().expect("result is a number");
        *per_worker.entry(who).or_insert(0) += 1;

        results
            .send_multipart(vec![frames[0].clone(), Msg::from_static(b"ok")])
            .await?;
    }
    println!("sum of squares: {total}");
    Ok(per_worker)
}

#[tokio::main]
async fn main() -> Result<(), ZmqError> {
    let ctx = Context::new()?;

    let sink = ctx.socket(SocketType::Router)?;
    sink.bind(RESULTS).await?;
    let collecting = tokio::spawn(collector(sink));

    let ventilator = ctx.socket(SocketType::Push)?;
    let events = ventilator.monitor_default().await?;
    ventilator.bind(TASKS).await?;

    let mut workers = Vec::new();
    for id in 1..=WORKERS {
        let tasks = ctx.socket(SocketType::Pull)?;
        tasks.connect(TASKS).await?;

        let results = ctx.socket(SocketType::Dealer)?;
        results.set_option(ROUTING_ID, format!("worker-{id}")).await?;
        results.connect(RESULTS).await?;

        workers.push(tokio::spawn(worker(tasks, results)));
    }

    let mut ready = 0;
    while let Ok(event) = events.recv().await {
        if let SocketEvent::HandshakeSucceeded { .. } = event {
            ready += 1;
            if ready == WORKERS {
                break;
            }
        }
    }

    for n in 1..=TASK_COUNT {
        ventilator.send(Msg::from_vec(n.to_string().into_bytes())).await?;
    }

    let per_worker = collecting.await.expect("collector panicked")?;
    for (who, count) in &per_worker {
        println!("{who} handled {count} tasks");
    }

    for handle in workers {
        handle.await.expect("worker panicked")?;
    }
    ctx.term().await?;
    Ok(())
}

set_option takes an option ID from rzmq::socket and any value that implements rzmq's ToBytes trait. That includes String, &str, i32, bool and byte slices.

Run it:

$ cargo run
worker-1: 1 squared is 1
worker-2: 2 squared is 4
worker-1: 3 squared is 9
worker-2: 4 squared is 16
worker-1: 5 squared is 25
worker-2: 6 squared is 36
worker-1: 7 squared is 49
worker-2: 8 squared is 64
sum of squares: 204
worker-1 handled 4 tasks
worker-2 handled 4 tasks

As before, the two worker names can swap places between runs. The collector always sees eight results that add up to 204, split four and four.

Shut down cleanly#

Every task in this program knows when its work is done, so shutdown is a matter of waiting in the right order:

src/main.rs

let per_worker = collecting.await.expect("collector panicked")?;
for (who, count) in &per_worker {
    println!("{who} handled {count} tasks");
}

for handle in workers {
    handle.await.expect("worker panicked")?;
}
ctx.term().await?;

The collector returns after it has seen all eight results. Each worker returns after its share of tasks and the reply to its last result. main waits for all of them and then calls Context::term(), which closes every socket made from the context and waits for their background tasks to finish. Any error along the way is a real error, so it is passed up with ? instead of being ignored.

For a longer-running service that doesn't know its workload up front, you would add a stop signal instead, for example a tokio::sync::watch channel checked with tokio::select! next to recv. Dropping a pending recv future in rzmq doesn't lose the message; it stays queued for the next call. That makes it safe to race recv against a shutdown signal in select!.

Where to go next#

The rest of the socket types follow the same shape: create from the context, set options, bind or connect, then send and receive. PUB/SUB takes a SUBSCRIBE option with a topic prefix. The SNDHWM and RCVHWM options cap how many messages queue up per peer.

On Linux, rzmq has an optional io_uring backend for TCP. You turn it on with a feature flag and then opt in per socket with the IO_URING_SESSION_ENABLED option:

Cargo.toml

[dependencies]
rzmq = { version = "0", features = ["io-uring"] }

The feature is Linux-only, so this tutorial leaves it out.

The usage guide covers socket options, security (PLAIN, CURVE and Noise_XX), monitoring and io_uring setup in the rzmq usage guide on GitHub.

The full API reference is on docs.rs.