High-Performance Rust Services

The Tokio runtime model

Tokio's default runtime is a multi-threaded work-stealing executor. Conceptually:

• A small pool of OS threads (default: number of CPU cores).
• Each thread has its own local task queue.
• Tasks are futures the runtime polls. Each task is a state machine that yields at .await points.
• When a task yields (e.g., waiting on a socket), the thread picks up another task. Idle tasks don't occupy a thread.
• If one thread runs out of work, it steals from another thread's queue.

The result: a handful of OS threads can run hundreds of thousands of in-flight tasks, as long as most tasks are blocked on I/O most of the time. This is the right model for an inference gateway where most tasks are waiting on a model server.

Work-stealing in practice

Work-stealing means cores stay busy under uneven load. But it also means:

• Tasks can move between threads at .await points. So state captured by an async block must be Send.
• Tasks that never yield don't get stolen. A CPU-bound loop without an .await hogs its worker thread. Other tasks queued behind it on that local queue starve until a steal happens.
• tokio::task::yield_now() exists for the rare case you have a long compute and want to give the runtime a checkpoint.

For gateway code: avoid CPU-heavy work inside async tasks. If you must (tokenization, hashing, serialization of large payloads), see the next section.

Blocking vs cooperative tasks

The single most common Tokio mistake: calling something that blocks the thread inside an async task. Examples:

• std::thread::sleep — blocks the worker thread. Use tokio::time::sleep.
• Synchronous file I/O (std::fs::read). Use tokio::fs.
• Calling a reqwest::blocking client. Use the async one.
• CPU-bound work like crypto / compression / tokenization in a loop without yielding.
• Holding a std::sync::Mutex guard across .await — see 03.

The escape hatch for CPU-bound work:

let tokens = tokio::task::spawn_blocking(move || {
    // synchronous, CPU-heavy work runs on a dedicated blocking pool
    expensive_tokenize(&text)
}).await?;

spawn_blocking moves the work to a separate thread pool (default capacity: 512 threads, configurable) so it can't starve the async workers. Use it for genuinely synchronous work; don't use it as a hammer.

Backpressure

Backpressure is how a downstream tells the upstream "slow down, I'm full." Without it, a fast producer fills queues until OOM.

In async Rust the primary tool is the bounded channel. A bounded mpsc::Sender's send() awaits when the channel is full. That await propagates back: the caller awaits, can't accept more requests, the HTTP layer rejects new connections or returns 503.

Other backpressure tools:

• tower's Buffer + ConcurrencyLimit — bound the in-flight requests per service.
• tower's RateLimit — bound requests per unit time.
• tower's LoadShed — return errors immediately when the service is at capacity (vs blocking the caller).
• Semaphore (tokio::sync::Semaphore) — bound concurrent operations explicitly.

use std::sync::Arc;
use tokio::sync::Semaphore;

let permits = Arc::new(Semaphore::new(100)); // 100 concurrent model calls

let p = permits.clone();
tokio::spawn(async move {
    let _permit = p.acquire().await.unwrap();
    call_model(req).await
});

tower middleware

tower defines a Service<Request> trait — a function-like type that maps requests to futures of responses. Every layer is composable. The shape:

use tower::{ServiceBuilder, Service};
use std::time::Duration;

let svc = ServiceBuilder::new()
    .timeout(Duration::from_secs(5))
    .concurrency_limit(1000)
    .rate_limit(10_000, Duration::from_secs(1))
    .buffer(100)
    .layer(TracingLayer::new())
    .layer(AuthLayer::new())
    .service(inference_handler);

Layers you'll see on a production gateway:

• Timeout — every request must finish within N seconds.
• ConcurrencyLimit / LoadShed — bound in-flight requests.
• RateLimit — bound rate.
• Retry — retry with a policy.
• Trace — open a span per request.
• Auth — validate JWT, fish out claims.
• Compression — gzip/brotli.
• Custom: idempotency, request hedging, cost accounting.

tonic — Rust gRPC

tonic is the standard gRPC implementation for Rust. Built on hyper + tower. You define services in .proto, tonic-build generates Rust traits, you implement them.

use tonic::{Request, Response, Status};

#[tonic::async_trait]
impl Inference for InferenceService {
    async fn generate(
        &self,
        req: Request<GenerateReq>,
    ) -> Result<Response<GenerateResp>, Status> {
        let req = req.into_inner();
        let out = self.backend.call(&req).await
            .map_err(|e| Status::internal(e.to_string()))?;
        Ok(Response::new(out))
    }

    type GenerateStreamStream =
        std::pin::Pin<Box<dyn Stream<Item = Result<Token, Status>> + Send>>;

    async fn generate_stream(
        &self,
        req: Request<GenerateReq>,
    ) -> Result<Response<Self::GenerateStreamStream>, Status> {
        // bidirectional streaming for token-by-token output
        // ...
        unimplemented!()
    }
}

Key tonic features:

• Unary, server-streaming, client-streaming, bidi. Use server-streaming for token-by-token LLM output.
• Interceptors — middleware analog. Add auth, tracing.
• Deadline propagation — gRPC deadlines flow through the metadata. Honor them.
• Status codes — gRPC-native error model (UNAVAILABLE, DEADLINE_EXCEEDED, etc.).

hyper and axum

hyper is the low-level HTTP implementation. You almost never write hyper directly in app code.

axum is a popular tower-based web framework on top of hyper. Ergonomic handlers, extractors, layered middleware. For non-gRPC public-facing endpoints (REST, SSE streaming), axum is a strong default.

use axum::{Router, routing::post, Json};

async fn complete(Json(req): Json<CompleteReq>) -> Json<CompleteResp> {
    Json(call_model(req).await.unwrap())
}

let app = Router::new()
    .route("/v1/complete", post(complete))
    .layer(tower_http::trace::TraceLayer::new_for_http())
    .layer(tower_http::timeout::TimeoutLayer::new(Duration::from_secs(30)));

let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
axum::serve(listener, app).await.unwrap();

Allocators — jemalloc, mimalloc

Rust uses the system allocator by default. For long-running, allocation-heavy services, switching to jemalloc or mimalloc often wins 10-30% on throughput and meaningfully on tail latency. Both handle high-concurrency, multi-threaded allocation patterns better than glibc malloc.

// Cargo.toml: tikv-jemallocator = "0.5"

#[global_allocator]
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

For interview purposes: know the two names, know that #[global_allocator] lets you swap allocators, know that the gain is real but situational — benchmark before claiming a number.

Hot-path discipline — avoiding allocations

The fast path through the gateway gets executed once per request — millions of times per minute. Every allocation in there costs CPU and produces tail-latency jitter (allocator contention, page faults, etc.). Disciplines:

• Pre-allocate buffers and reuse them via a pool (bytes::BytesMut, object pools).
• Avoid format! in hot paths — it allocates. Use write! into a pre-sized buffer.
• Avoid .to_string() and .clone() on the request path; pass &str / &[u8] where you can.
• Use Cow<str> when sometimes-borrowed, sometimes-owned.
• Use SmallVec / tinyvec when most cases fit on the stack.
• Use bytes::Bytes — cheaply cloneable, ref-counted byte buffer; the standard "share a chunk of network data without copying."

Zero-copy with the bytes crate

An LLM token response can be hundreds of KB. Copying it three times through your service (deserialize → process → reserialize) adds real cost. The bytes crate gives you:

• Bytes — a cheaply cloneable, ref-counted, immutable byte slice. Multiple Bytes can point into the same underlying buffer.
• BytesMut — a writable buffer that can be split into Bytes handles.

use bytes::{Bytes, BytesMut, BufMut};

let mut buf = BytesMut::with_capacity(4096);
buf.put_slice(b"header");
buf.put_slice(&payload);
let frame: Bytes = buf.freeze(); // immutable, cheaply cloneable

let copy1 = frame.clone(); // refcount++, no data copy
let copy2 = frame.clone(); // refcount++, no data copy

hyper, tonic, and most of the Rust networking ecosystem already use Bytes under the hood. Stay in that type as long as possible.

Profiling Rust services

Tools you should know by name and rough purpose:

# Flamegraph a release build under load
cargo flamegraph --release --bin gateway

# tokio-console: enable in Cargo.toml + RUSTFLAGS="--cfg tokio_unstable"
TOKIO_CONSOLE_BIND=0.0.0.0:6669 ./target/release/gateway &
tokio-console http://127.0.0.1:6669

The senior-Rust perf checklist

Items to mentally run through when an interviewer asks "how would you optimize this service?":

1. Measure first. Flamegraph under representative load. No guessing.
2. Find blocking calls. tokio-console. Anything blocked on a thread is suspect.
3. Look for hot allocations. Heap profile. Eliminate per-request String / Vec growth.
4. Check serialization. JSON is expensive; if internal RPCs are JSON, move them to protobuf / bincode.
5. Check the channel queue depth. Are bounded channels saturating? That's a capacity tell.
6. Inspect locks. Hot mutexes are tail-latency killers; consider parking_lot, RwLock, sharding, or actor patterns.
7. Try jemalloc/mimalloc. Measure, don't claim.
8. Look at the codegen. cargo asm, check #[inline] opportunities only after a profile says it matters.
9. Last: consider PGO and LTO. Release builds with lto = "fat" and codegen-units = 1.

✓ Check yourself

Answer from memory before opening each one.

1. A gateway handler tokenises a 200 KB prompt in a tight CPU loop with no await. Latency for unrelated requests rises. Why does that happen, and what is the fix?

   Show answer

   A task that never yields keeps its worker thread, and other tasks queued on that thread wait behind it. Move the synchronous work to the blocking pool with spawn_blocking so the async workers stay free, and avoid using it as a hammer for work that is already async. See Work-stealing in practice and Blocking vs cooperative tasks.

2. Traffic spikes and the gateway's memory climbs until the pod is killed, though CPU never saturates. Where do you look first, and what does the structure of a fix look like?

   Show answer

   Look for an unbounded channel or an unlimited number of spawned tasks between the front door and the model server. Bound the queue and cap in-flight work with a semaphore or concurrency limit, so that a full system makes callers wait or receive a 503 early instead of growing memory. See Backpressure.

3. In a tower stack, what is the difference between a concurrency limit and load shedding, and which do you want in front of a model server?

   Show answer

   A concurrency limit makes extra callers wait for capacity, while load shedding rejects them at once when the service is at capacity. Waiting hides overload as latency, so for a latency-bound service in front of a model server you generally want shedding, with a small bounded queue if bursts are short. See Backpressure and tower middleware.

4. An LLM response of several hundred KB passes through deserialise, transform and reserialise steps. How do you cut the copying, and what property of the type makes it safe?

   Show answer

   Hold the data as a Bytes value and clone it by reference count instead of copying. Bytes is immutable and cheap to clone, and many handles can point into one buffer, so stages share it without a copy. Build with BytesMut and freeze once. See Zero-copy with the bytes crate.

5. Swapping the global allocator to jemalloc is proposed as a fix for high p99. How do you respond?

   Show answer

   Treat it as a hypothesis. The gain is real for allocation-heavy, multi-threaded services but depends on the workload, so measure throughput and tail latency before and after under representative load, and compare against what a flamegraph says about allocation cost. See Allocators and the perf checklist.

6. Someone proposes rewriting a slow Rust service's hot path with unsafe code. Walk through your order of investigation, and say what evidence would justify the rewrite.

   Show answer

   Profile first with a flamegraph under load, then look for blocking calls, per-request allocations, JSON on internal RPCs, saturated channels and hot locks, and only then codegen. A rewrite is justified when the profile shows the cost concentrated in code a safe change cannot fix. See The senior-Rust perf checklist.

Takeaway

• Tokio runs many cheap tasks on a few worker threads, and the one rule that matters is never to block a worker thread.
• CPU-heavy or synchronous work belongs on the blocking pool via spawn_blocking, because a task that never yields starves the tasks queued behind it.
• Backpressure comes from bounded channels, semaphores and tower limits. Overload should surface early as a wait or a 503, not as memory growth.
• tower layers compose timeout, limits, retry, tracing and auth around one service, and tonic and axum are built on the same model.
• Hot-path cost comes from allocations and copies, so reuse buffers, avoid format! and clones, and share data with the bytes crate.
• Allocators, PGO and LTO are measured wins, not claims. Profile with a flamegraph and tokio-console before changing anything.

The second pillar is the layer these services call: model servers, batching, streaming and orchestration. → Deep dive: secondary topic

Further reading

• Tokio tutorial — Tokio project. The runtime model and task semantics behind this chapter.
• Async: What is blocking? — Alice Ryhl, 2020. A precise account of why blocking in async code hurts and how to fix it.
• tokio::task::spawn_blocking — docs.rs. Exact behaviour and limits of the blocking pool.
• tower — docs.rs. Service, Layer and the limit, load-shed, retry and timeout middleware.
• bytes — docs.rs. Bytes and BytesMut, and how cloning shares a buffer.
• The Rust Performance Book — Nicholas Nethercote. Profiling, allocation and build settings for Rust programs.
• tokio-console — Tokio project. The task inspector used to find tasks stuck on await.