Stream messages over io_uring with PUSH/PULL and PUB/SUB

You will build two streams over TCP on localhost and run both on rzmq's io_uring backend. A PUSH/PULL pipeline carries 10,000 readings to a sink and a PUB socket streams price ticks to two subscribers. Along the way you pick the number of io_uring workers, switch TCP_CORK on and off and ask for zero-copy sends and multishot receives.

rzmq = "0"

Set up the project#

Start a new binary crate.

cargo new uring-streams
cd uring-streams

The io_uring backend is Linux only. Turn it on with the io-uring feature:

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

Streams are where io_uring pays off. A PUSH socket feeding a pipeline or a PUB socket feeding subscribers keeps a steady flow of messages going in one direction, so the backend can batch submissions to the kernel instead of making a system call per message.

Turn on io_uring for a PUSH/PULL stream#

Turning on the backend takes two things. initialize_uring_backend runs once before the first Context and sets up the backend's worker threads. Then each socket that should use it sets IO_URING_SESSION_ENABLED to 1i32 before bind or connect. Its TCP connections then run on io_uring. At the end, shutdown_uring_backend stops the workers after the context has terminated.

The first stream is a pipeline: a PUSH socket sends the readings 1 to 10,000 and a PULL sink adds them up. Replace src/main.rs with this:

use rzmq::socket::IO_URING_SESSION_ENABLED;
use rzmq::uring::{UringConfig, initialize_uring_backend, shutdown_uring_backend};
use rzmq::{Context, Msg, Socket, SocketType, ZmqError};

const READINGS: &str = "tcp://127.0.0.1:5660";
const READING_COUNT: u64 = 10_000;

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

async fn uring_socket(ctx: &Context, kind: SocketType) -> Result<Socket, ZmqError> {
    let socket = ctx.socket(kind)?;
    socket.set_option(IO_URING_SESSION_ENABLED, 1i32).await?;
    Ok(socket)
}

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

    let sink = uring_socket(&ctx, SocketType::Pull).await?;
    sink.bind(READINGS).await?;
    let source = uring_socket(&ctx, SocketType::Push).await?;
    source.connect(READINGS).await?;

    let producer = tokio::spawn(async move {
        for n in 1..=READING_COUNT {
            source.send(Msg::from_vec(n.to_string().into_bytes())).await?;
        }
        Ok::<(), ZmqError>(())
    });

    let mut sum = 0;
    for _ in 0..READING_COUNT {
        let reading = sink.recv().await?;
        sum += text(&reading).parse::<u64>().expect("a reading is a number");
    }
    producer.await.expect("producer panicked")?;
    println!("sink received {READING_COUNT} readings, sum {sum}");

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

The send and receive code is the same as on the default backend. The backend is a property of each socket's connections, so the PUSH/PULL rules don't change: PUSH queues readings for its connected peer and PULL hands them out in the order they were sent.

Run it:

$ cargo run
sink received 10000 readings, sum 50005000

Choose the number of workers#

The backend runs its I/O on UringWorker threads. Each worker has its own ring, its own buffer pools and its own set of file descriptors. A new connection is assigned to a worker round-robin and stays on that worker for its whole life, so the worker count caps how many threads share the connections.

UringConfig::num_workers sets the count. The default comes from rzmq::uring::default_uring_num_workers(), which takes half the available cores rounded up, subtracts 2 and clamps the result to between 1 and 8. That leaves room for the Tokio runtime's own threads.

Set it in the config you pass to initialize_uring_backend:

    initialize_uring_backend(UringConfig {
        num_workers: 2,
        ..Default::default()
    })?;

Each worker registers its own receive buffer ring, so the pinned memory for receive buffers grows with the worker count. The zero-copy send budget works the other way: default_send_buffer_count is a total that is split across the workers. More workers help when many connections are busy at once. A program with one or two streams gets nothing from more workers than it has connections.

The output doesn't change:

sink received 10000 readings, sum 50005000

Stream prices to subscribers with PUB/SUB#

A PUB socket sends every message to all subscribers whose subscriptions match. This stream publishes the ticks 1 to 1,000 to two subscribers, prices and audit, both subscribed to everything.

A subscriber's subscription reaches the publisher a moment after its connection is up and a PUB socket drops messages that no subscriber has asked for yet. So before streaming, the publisher repeats a hello message every 10 ms. Each subscriber answers its first hello over a separate PUSH/PULL channel. Once both have answered, the publisher knows both are listening and starts the real stream. An end message closes it.

The publisher also turns on TCP_CORK, which the next section explains. The helper now takes a cork flag and sets TCP_CORK right after opting the socket into io_uring.

Replace src/main.rs:

use rzmq::socket::{IO_URING_SESSION_ENABLED, SUBSCRIBE, TCP_CORK};
use rzmq::uring::{UringConfig, initialize_uring_backend, shutdown_uring_backend};
use rzmq::{Context, Msg, Socket, SocketType, ZmqError};
use std::time::Duration;

const TICKS: &str = "tcp://127.0.0.1:5661";
const READY: &str = "tcp://127.0.0.1:5662";
const TICK_COUNT: u64 = 1_000;
const SUBSCRIBERS: [&str; 2] = ["prices", "audit"];

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

async fn option(socket: &Socket, id: i32) -> Result<i32, ZmqError> {
    let raw = socket.get_option(id).await?;
    Ok(i32::from_ne_bytes(raw.try_into().expect("option is an i32")))
}

async fn uring_socket(ctx: &Context, kind: SocketType, cork: bool) -> Result<Socket, ZmqError> {
    let socket = ctx.socket(kind)?;
    socket.set_option(IO_URING_SESSION_ENABLED, 1i32).await?;
    socket.set_option(TCP_CORK, if cork { 1i32 } else { 0i32 }).await?;
    Ok(socket)
}

async fn subscriber(ctx: Context, name: &'static str) -> Result<(&'static str, u64, u64), ZmqError> {
    let ticks = uring_socket(&ctx, SocketType::Sub, false).await?;
    ticks.set_option(SUBSCRIBE, "").await?;
    ticks.connect(TICKS).await?;
    let ready = uring_socket(&ctx, SocketType::Push, false).await?;
    ready.connect(READY).await?;

    let mut answered = false;
    let (mut count, mut sum) = (0, 0);
    loop {
        let msg = ticks.recv().await?;
        match text(&msg).as_str() {
            "hello" if !answered => {
                ready.send(Msg::new()).await?;
                answered = true;
            }
            "hello" => {}
            "end" => return Ok((name, count, sum)),
            tick => {
                count += 1;
                sum += tick.parse::<u64>().expect("a tick is a number");
            }
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), ZmqError> {
    initialize_uring_backend(UringConfig {
        num_workers: 2,
        ..Default::default()
    })?;
    let ctx = Context::new()?;

    let publisher = uring_socket(&ctx, SocketType::Pub, true).await?;
    publisher.bind(TICKS).await?;
    let ready = uring_socket(&ctx, SocketType::Pull, false).await?;
    ready.bind(READY).await?;
    println!("publisher TCP_CORK = {}", option(&publisher, TCP_CORK).await?);

    let mut subscribers = Vec::new();
    for name in SUBSCRIBERS {
        subscribers.push(tokio::spawn(subscriber(ctx.clone(), name)));
    }

    let answers = {
        let ready = ready.clone();
        tokio::spawn(async move {
            for _ in 0..SUBSCRIBERS.len() {
                ready.recv().await?;
            }
            Ok::<(), ZmqError>(())
        })
    };
    let mut hello = tokio::time::interval(Duration::from_millis(10));
    while !answers.is_finished() {
        hello.tick().await;
        publisher.send(Msg::from_static(b"hello")).await?;
    }
    answers.await.expect("answer task panicked")?;

    for n in 1..=TICK_COUNT {
        publisher.send(Msg::from_vec(n.to_string().into_bytes())).await?;
    }
    publisher.send(Msg::from_static(b"end")).await?;

    for handle in subscribers {
        let (name, count, sum) = handle.await.expect("subscriber panicked")?;
        println!("{name}: {count} ticks, sum {sum}");
    }

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

Run it:

$ cargo run
publisher TCP_CORK = 1
prices: 1000 ticks, sum 500500
audit: 1000 ticks, sum 500500

Both subscribers got every tick. main joins them in the order they were spawned, so the lines always come out in this order.

Switch TCP_CORK on and off#

TCP_CORK takes 1i32 or 0i32 and turns TCP corking on or off for an io_uring socket. With corking on, the kernel holds small writes back and fills TCP segments before sending them, which batches a stream of small messages into fewer packets. The publisher above sends a thousand short ticks back to back, which is the case corking is for.

Corking is the wrong choice when every message has to go out the moment it is sent, such as a request that someone is waiting on. The subscribers and the ready channel in this program only send one short answer each, so they keep it off.

To see the setting change, pass false for the publisher:

    let publisher = uring_socket(&ctx, SocketType::Pub, false).await?;

publisher TCP_CORK = 0
prices: 1000 ticks, sum 500500
audit: 1000 ticks, sum 500500

Corking changes how bytes are grouped into TCP segments. It never changes which messages arrive or their order, so the ticks and sums are the same either way. Put true back before moving on.

Ask for zero-copy sends and multishot receives#

Three more options tune how an io_uring socket moves data. Each one is a request that the backend configuration has to support:

• IO_URING_SNDZEROCOPY asks for zero-copy sends. It takes effect when UringConfig::default_send_zerocopy is true, which also sets up the send buffer pool.
• IO_URING_ZC_SEND_THRESHOLD is the smallest payload in bytes that goes out as a zero-copy send. Smaller payloads use a normal copy-based send. The default is 16384 and it only matters when IO_URING_SNDZEROCOPY is on.
• IO_URING_RCVMULTISHOT asks for multishot receives, where one receive request keeps delivering data as it arrives. It depends on UringConfig::default_recv_multishot, which is true in the default config.

Import the new options:

use rzmq::socket::{
    IO_URING_RCVMULTISHOT, IO_URING_SESSION_ENABLED, IO_URING_SNDZEROCOPY,
    IO_URING_ZC_SEND_THRESHOLD, SUBSCRIBE, TCP_CORK,
};

Turn on the zero-copy pool in the backend config:

    initialize_uring_backend(UringConfig {
        num_workers: 2,
        default_send_zerocopy: true,
        ..Default::default()
    })?;

Set zero-copy on the publisher and multishot on each subscriber's SUB socket, before bind and connect:

    let publisher = uring_socket(&ctx, SocketType::Pub, true).await?;
    publisher.set_option(IO_URING_SNDZEROCOPY, 1i32).await?;
    publisher.set_option(IO_URING_ZC_SEND_THRESHOLD, 16384i32).await?;
    publisher.bind(TICKS).await?;

    let ticks = uring_socket(&ctx, SocketType::Sub, false).await?;
    ticks.set_option(IO_URING_RCVMULTISHOT, 1i32).await?;
    ticks.set_option(SUBSCRIBE, "").await?;
    ticks.connect(TICKS).await?;

A tick is a few bytes, well under the 16384-byte threshold, so ticks still go out as copy-based sends. Zero-copy pays off for large payloads such as snapshots or file chunks, where copying the data would cost more than setting up the zero-copy send. The output is the same as before:

publisher TCP_CORK = 1
prices: 1000 ticks, sum 500500
audit: 1000 ticks, sum 500500

Shut down cleanly#

Every program here ends the same way:

    ctx.term().await?;
    shutdown_uring_backend().await?;

The streams end on their own: the sink stops after the last reading and each subscriber returns on end. ctx.term() then closes every socket made from the context and waits for their background tasks. shutdown_uring_backend() runs last and joins the backend's worker threads. initialize_uring_backend runs once per process, so it belongs at the top of main before the first context.

Where to go next#

The rzmq usage guide on GitHub covers every UringConfig field, the rest of the socket options and outbound write batching with SNDBATCH_COUNT and SNDBATCH_BYTES. The PUB/SUB tutorial covers topic filtering and UNSUBSCRIBE. The PUSH/PULL tutorial covers high-water marks and send timeouts.

The full API reference, including the rzmq::uring module and UringConfig, is on docs.rs.