Tune the adaptive I/O throttle on a PUSH/PULL stream

For Rust developers running rzmq under uneven traffic who want to see and change how each connection shares its time between sending and receiving

You will stream 1,000 readings from a PUSH source to a PULL sink over TCP on localhost and read the adaptive throttle's state on both ends. Then you give the sink its own throttle settings, see what happens when they arrive too late and switch the throttle off on the source.

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 throttled-stream
cd throttled-stream

Add rzmq and Tokio:

Cargo.toml

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

The programs use Tokio's single-threaded current_thread runtime, which keeps the numbers they print the same from run to run.

Read the throttle's settings#

Every rzmq connection on the default backend runs an adaptive I/O throttle. It keeps a balance per connection: each incoming message adds credit_per_message and each outgoing message subtracts it. It learns the normal balance over time. When the live balance drifts too far from it, the throttle sometimes yields to the Tokio scheduler so the direction that is falling behind gets to run. A connection that mostly sends can't starve its own receiving that way.

A socket's throttle settings live in the ADAPTIVE_THROTTLE option as an AdaptiveThrottleSocketConfig. get_option returns it encoded and FromBytes decodes it. Replace src/main.rs:

src/main.rs

use rzmq::socket::types::FromBytes;
use rzmq::socket::{ADAPTIVE_THROTTLE, AdaptiveThrottleSocketConfig};
use rzmq::{Context, SocketType, ZmqError};

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<(), ZmqError> {
    let ctx = Context::new()?;
    let socket = ctx.socket(SocketType::Pull)?;

    let config = AdaptiveThrottleSocketConfig::from_bytes(&socket.get_option(ADAPTIVE_THROTTLE).await?)?;
    println!("{config:#?}");

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

$ cargo run
AdaptiveThrottleSocketConfig {
    enabled: true,
    credit_per_message: 5,
    healthy_balance_width: 1024000,
    max_imbalance: 6553600,
    yield_after_n_consecutive: 256,
    nudge_interval_ops: 100,
    adaptive_learning_rate: 0.05,
    strategy: PowerCurve {
        exponent: 2.0,
    },
    priority: Auto,
    priority_boost_factor: 5.0,
}

These are the defaults every socket starts with:

credit_per_message: how far one message moves the balance.
healthy_balance_width: the half-width of the zone around the learned balance where the throttle never yields.
max_imbalance: how far past that zone the chance of yielding reaches 100%.
yield_after_n_consecutive: forces a yield after that many operations in one direction.
nudge_interval_ops and adaptive_learning_rate: how often and how fast the learned balance follows the live one.
strategy: the curve from imbalance to the chance of yielding, ThrottleStrategy::PowerCurve { exponent } or ThrottleStrategy::Linear.
priority: the direction to favour. ThrottlePriority::Auto picks Egress for a connection the socket accepted and Ingress for one it opened. Egress, Ingress and None set it outright.
priority_boost_factor: multiplies the chance of yielding while doing the non-priority direction during an imbalance.
enabled: whether the throttle runs at all.

Watch the throttle on a stream#

ADAPTIVE_THROTTLE_STATS is read-only. It returns one ThrottleStats per connection, with the priority the connection settled on and its live balance. A stream only moves one way, so the balance on each end runs away from zero in its own direction.

The source binds and the sink connects. Replace src/main.rs:

src/main.rs

use rzmq::socket::types::FromBytes;
use rzmq::socket::{ThrottleStats, ADAPTIVE_THROTTLE_STATS};
use rzmq::{Context, Msg, Socket, SocketType, ZmqError};

const READINGS: &str = "tcp://127.0.0.1:5680";
const READING_COUNT: u32 = 1_000;

async fn throttle_stats(socket: &Socket) -> Result<Vec<ThrottleStats>, ZmqError> {
    Vec::<ThrottleStats>::from_bytes(&socket.get_option(ADAPTIVE_THROTTLE_STATS).await?)
}

async fn report(name: &str, socket: &Socket) -> Result<(), ZmqError> {
    for stats in throttle_stats(socket).await? {
        println!(
            "{name}: priority {:?}, balance {}",
            stats.priority, stats.current_balance
        );
    }
    Ok(())
}

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

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

    for n in 1..=READING_COUNT {
        source.send(Msg::from_vec(n.to_string().into_bytes())).await?;
    }
    for _ in 0..READING_COUNT {
        sink.recv().await?;
    }
    println!("sink received {READING_COUNT} readings");

    report("source", &source).await?;
    report("sink", &sink).await?;

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

sink received 1000 readings
source: priority Egress, balance -5000
sink: priority Ingress, balance 5000

The source accepted its connection, so Auto resolved to Egress. The sink opened its connection and got Ingress. Each side moved 1,000 messages at the default credit_per_message of 5, outgoing on the source and incoming on the sink. The stats are read after the sink has every reading, so both connections have finished their work and the numbers are settled.

Give the sink its own settings#

To change the settings, build an AdaptiveThrottleSocketConfig and set it as ADAPTIVE_THROTTLE. Start from ..Default::default() and name only the fields you want to change. Here the sink counts each message as 1, yields after 64 operations in a row, uses the linear curve and favours neither direction. Replace src/main.rs:

src/main.rs

use rzmq::socket::types::FromBytes;
use rzmq::socket::{
    ADAPTIVE_THROTTLE, ADAPTIVE_THROTTLE_STATS, AdaptiveThrottleSocketConfig, ThrottlePriority,
    ThrottleStats, ThrottleStrategy,
};
use rzmq::{Context, Msg, Socket, SocketType, ZmqError};

const READINGS: &str = "tcp://127.0.0.1:5680";
const READING_COUNT: u32 = 1_000;

async fn throttle_config(socket: &Socket) -> Result<AdaptiveThrottleSocketConfig, ZmqError> {
    AdaptiveThrottleSocketConfig::from_bytes(&socket.get_option(ADAPTIVE_THROTTLE).await?)
}

async fn report(name: &str, socket: &Socket) -> Result<(), ZmqError> {
    let stats = Vec::<ThrottleStats>::from_bytes(&socket.get_option(ADAPTIVE_THROTTLE_STATS).await?)?;
    for s in stats {
        println!("{name}: priority {:?}, balance {}", s.priority, s.current_balance);
    }
    Ok(())
}

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

    let source = ctx.socket(SocketType::Push)?;
    source.bind(READINGS).await?;

    let sink = ctx.socket(SocketType::Pull)?;
    let tuned = AdaptiveThrottleSocketConfig {
        credit_per_message: 1,
        yield_after_n_consecutive: 64,
        strategy: ThrottleStrategy::Linear,
        priority: ThrottlePriority::None,
        ..Default::default()
    };
    sink.set_option(ADAPTIVE_THROTTLE, &tuned).await?;
    let config = throttle_config(&sink).await?;
    println!(
        "sink throttle: credit {}, yield after {}, {:?}, priority {:?}",
        config.credit_per_message, config.yield_after_n_consecutive, config.strategy, config.priority
    );
    sink.connect(READINGS).await?;

    for n in 1..=READING_COUNT {
        source.send(Msg::from_vec(n.to_string().into_bytes())).await?;
    }
    for _ in 0..READING_COUNT {
        sink.recv().await?;
    }
    println!("sink received {READING_COUNT} readings");

    report("source", &source).await?;
    report("sink", &sink).await?;

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

sink throttle: credit 1, yield after 64, Linear, priority None
sink received 1000 readings
source: priority Egress, balance -5000
sink: priority None, balance 1000

Reading the option back returns what you set. The sink's connection now runs with priority None and its balance grew by 1 per reading. The source kept its defaults.

Socket::with_throttle_config sets the same option while you build the socket:

src/main.rs

let sink = ctx
    .socket(SocketType::Pull)?
    .with_throttle_config(tuned)
    .await?;

Set it before bind or connect#

A connection takes its throttle settings when it starts. Once a socket has called bind or connect, setting ADAPTIVE_THROTTLE returns ZmqError::InvalidState. Add this after sink.connect:

src/main.rs

    match sink.set_option(ADAPTIVE_THROTTLE, 0i32).await {
        Err(ZmqError::InvalidState(reason)) => println!("too late: {reason}"),
        other => println!("unexpected: {other:?}"),
    }

sink throttle: credit 1, yield after 64, Linear, priority None
too late: ADAPTIVE_THROTTLE must be set before bind or connect
sink received 1000 readings
source: priority Egress, balance -5000
sink: priority None, balance 1000

The sink keeps the settings it had when it connected.

Switch the throttle off#

ADAPTIVE_THROTTLE also takes 1i32 or 0i32. That changes enabled and leaves the other fields as they are. A disabled throttle does no work on the connection's hot path, which suits a socket that only ever moves data one way. Turn it off on the source before it binds:

src/main.rs

    let source = ctx.socket(SocketType::Push)?;
    source.set_option(ADAPTIVE_THROTTLE, 0i32).await?;
    println!("source throttle enabled: {}", throttle_config(&source).await?.enabled);
    source.bind(READINGS).await?;

source throttle enabled: false
sink throttle: credit 1, yield after 64, Linear, priority None
too late: ADAPTIVE_THROTTLE must be set before bind or connect
sink received 1000 readings
source: priority Egress, balance 0
sink: priority None, balance 1000

The source's connection is still listed, but its balance stays at 0 because the throttle no longer counts anything. The readings arrive the same either way. The throttle only decides when a connection yields to the scheduler.

The finished program:

src/main.rs

use rzmq::socket::types::FromBytes;
use rzmq::socket::{
    ADAPTIVE_THROTTLE, ADAPTIVE_THROTTLE_STATS, AdaptiveThrottleSocketConfig, ThrottlePriority,
    ThrottleStats, ThrottleStrategy,
};
use rzmq::{Context, Msg, Socket, SocketType, ZmqError};

const READINGS: &str = "tcp://127.0.0.1:5680";
const READING_COUNT: u32 = 1_000;

async fn throttle_config(socket: &Socket) -> Result<AdaptiveThrottleSocketConfig, ZmqError> {
    AdaptiveThrottleSocketConfig::from_bytes(&socket.get_option(ADAPTIVE_THROTTLE).await?)
}

async fn report(name: &str, socket: &Socket) -> Result<(), ZmqError> {
    let stats = Vec::<ThrottleStats>::from_bytes(&socket.get_option(ADAPTIVE_THROTTLE_STATS).await?)?;
    for s in stats {
        println!("{name}: priority {:?}, balance {}", s.priority, s.current_balance);
    }
    Ok(())
}

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

    let source = ctx.socket(SocketType::Push)?;
    source.set_option(ADAPTIVE_THROTTLE, 0i32).await?;
    println!("source throttle enabled: {}", throttle_config(&source).await?.enabled);
    source.bind(READINGS).await?;

    let sink = ctx.socket(SocketType::Pull)?;
    let tuned = AdaptiveThrottleSocketConfig {
        credit_per_message: 1,
        yield_after_n_consecutive: 64,
        strategy: ThrottleStrategy::Linear,
        priority: ThrottlePriority::None,
        ..Default::default()
    };
    sink.set_option(ADAPTIVE_THROTTLE, &tuned).await?;
    let config = throttle_config(&sink).await?;
    println!(
        "sink throttle: credit {}, yield after {}, {:?}, priority {:?}",
        config.credit_per_message, config.yield_after_n_consecutive, config.strategy, config.priority
    );
    sink.connect(READINGS).await?;

    match sink.set_option(ADAPTIVE_THROTTLE, 0i32).await {
        Err(ZmqError::InvalidState(reason)) => println!("too late: {reason}"),
        other => println!("unexpected: {other:?}"),
    }

    for n in 1..=READING_COUNT {
        source.send(Msg::from_vec(n.to_string().into_bytes())).await?;
    }
    for _ in 0..READING_COUNT {
        sink.recv().await?;
    }
    println!("sink received {READING_COUNT} readings");

    report("source", &source).await?;
    report("sink", &sink).await?;

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

Sockets on the io_uring backend#

io_uring sockets do not support the adaptive throttle. A socket with IO_URING_SESSION_ENABLED set runs its TCP connections on the io_uring backend, which has no throttle. ADAPTIVE_THROTTLE settings have no effect on those connections and ADAPTIVE_THROTTLE_STATS does not list them.

Shut down cleanly#

The program ends with ctx.term():

src/main.rs

    ctx.term().await?;

The sink has read every reading by then, so nothing is left in flight. term() closes every socket made from the context and waits for their background tasks.

Where to go next#

The PUSH/PULL tutorial covers high-water marks and send timeouts. The DEALER/ROUTER tutorial builds a connection that carries traffic both ways.

The rzmq usage guide on GitHub has the adaptive throttle chapter and the rest of the socket options. The full API reference is on docs.rs.