Rate limits, retries, and idempotency

Programs that talk to other programs over a network see failures that local code never sees. A server is busy and refuses a request. A connection drops halfway. A reply takes so long that the client gives up, without knowing whether the server did the work.

This chapter builds three tools for those situations:

• A rate limiter protects a service from too many requests, by refusing requests above a set rate.
• A retry policy tries a failed call again, waiting longer after each failure.
• An idempotency key makes it safe to send the same request twice, by making the second one do nothing new.

The three work together. A client that retries must not overload the server, so it backs off. And a retried request must not repeat a side effect, such as charging a card twice.

19.1 A token bucket

Suppose a service allows each client 5 requests per second. A simple counter that resets every second allows 5 requests at 0.99 s and 5 more at 1.01 s: 10 requests in 20 ms. A token bucket smooths this out.

Picture a bucket that holds up to 5 tokens. Every request takes one token out. If the bucket is empty, the request is refused. Tokens flow back in at a steady rate, say 1 per second, until the bucket is full again (figure 19.1).

Two numbers describe the bucket. The capacity is the largest burst allowed at once. The rate is the long-run number of requests per second.

The bucket does not need a timer that adds tokens. Each request computes how many tokens arrived since the last request, from the time that has passed. It adds them before it checks.

19.1.1 The bucket behind a lock

Listing 19.1 The bucket (lines 6 to 31). src/problems/rate_limiter.rs

use std::{sync::Mutex, time::Instant};

pub struct TokenBucket {
    /// Max tokens.
    capacity: f64,
    /// Tokens per second.
    rate: f64,
    state: Mutex<BucketState>,
}

struct BucketState {
    tokens: f64,
    last_refill: Instant,
}

impl TokenBucket {
    pub fn new(capacity: usize, rate_per_sec: f64) -> Self {
        assert!(capacity > 0);
        assert!(rate_per_sec > 0.0);
        Self {
            capacity: capacity as f64,
            rate: rate_per_sec,
            state: Mutex::new(BucketState {
                tokens: capacity as f64, // Start full
                last_refill: Instant::now(),
            }),
        }
    }
    // ...
}

The capacity and the rate never change, so they are plain fields. The token count and the time of the last refill change on every call, so they live in BucketState behind a Mutex. The token count is an f64, because a fraction of a token can arrive between two requests. The bucket starts full.

Instant is the standard library’s clock for measuring durations. It only moves forward, unlike the wall-clock time, which can jump when the system clock is corrected.

Listing 19.2 try_acquire (lines 32 to 49).

impl TokenBucket {
    // ...
    pub fn try_acquire(&self) -> bool {
        let mut s = self.state.lock().unwrap();
        let now = Instant::now();

        // Refill tokens based on elapsed time
        let elapsed = now - s.last_refill;
        let new_tokens = elapsed.as_secs_f64() * self.rate;
        s.tokens = (s.tokens + new_tokens).min(self.capacity);
        s.last_refill = now;

        // Try to take one
        if s.tokens >= 1.0 {
            s.tokens -= 1.0;
            true
        } else {
            false
        }
    }
}

try_acquire does all its work under the lock:

1. It computes the time since the last refill, now - s.last_refill, a Duration.
2. It converts that to seconds and multiplies by the rate: the number of new tokens.
3. It adds them and caps the total at the capacity with .min(self.capacity).
4. It records now as the last refill.
5. If at least one token is there, it takes it and returns true. Otherwise it returns false.

Work through the test by hand. The bucket starts with 5 tokens. Five calls happen within microseconds, so almost no tokens refill, and each call takes one. The sixth finds about 0.00001 tokens and returns false. After a one-second sleep, about 1.00001 tokens are there, so one call succeeds, and the next finds about 0.00001 again.

Listing 19.3 The complete file, with its test. src/problems/rate_limiter.rs

//! Thread-safe token bucket rate limiter.

use std::{sync::Mutex, time::Instant};

pub struct TokenBucket {
    /// Max tokens.
    capacity: f64,
    /// Tokens per second.
    rate: f64,
    state: Mutex<BucketState>,
}

struct BucketState {
    tokens: f64,
    last_refill: Instant,
}

impl TokenBucket {
    pub fn new(capacity: usize, rate_per_sec: f64) -> Self {
        assert!(capacity > 0);
        assert!(rate_per_sec > 0.0);
        Self {
            capacity: capacity as f64,
            rate: rate_per_sec,
            state: Mutex::new(BucketState {
                tokens: capacity as f64, // Start full
                last_refill: Instant::now(),
            }),
        }
    }

    pub fn try_acquire(&self) -> bool {
        let mut s = self.state.lock().unwrap();
        let now = Instant::now();

        // Refill tokens based on elapsed time
        let elapsed = now - s.last_refill;
        let new_tokens = elapsed.as_secs_f64() * self.rate;
        s.tokens = (s.tokens + new_tokens).min(self.capacity);
        s.last_refill = now;

        // Try to take one
        if s.tokens >= 1.0 {
            s.tokens -= 1.0;
            true
        } else {
            false
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{thread, time::Duration};

    use super::*;

    #[test]
    fn it_works() {
        let bucket = TokenBucket::new(5, 1.0); // 5 tokens, refill 1 per sec

        // Use all 5
        for _ in 0..5 {
            assert!(bucket.try_acquire());
        }
        // 6th fails
        assert!(!bucket.try_acquire());

        // Wait 1 second
        thread::sleep(Duration::from_secs(1));
        // Should have 1 token back
        assert!(bucket.try_acquire());
        // And none left
        assert!(!bucket.try_acquire());
    }
}

$ cargo test --lib problems::rate_limiter
running 1 test
test problems::rate_limiter::tests::it_works ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 136 filtered out; finished in 1.00s

The test takes one second, because it sleeps to let a token refill. A test that depends on real time can also fail on a heavily loaded machine. Exercise 1 removes the dependency by passing the time in.

19.1.2 The bucket with atomics

Every call to the locked bucket takes the mutex. Under heavy load, threads wait for each other there, as section 14.1 described for maps. The second version replaces the lock with atomic integers.

Listing 19.4 The limiter (lines 5 to 37). src/bin/rate_limiter_atomic_token_bucket.rs

use std::{
    sync::atomic::{AtomicI64, AtomicU64, Ordering::Relaxed},
    thread,
    time::{Duration, Instant},
};

#[repr(align(64))]
pub struct RateLimiter {
    tokens: AtomicI64,
    last_refilled_at: AtomicU64, // secs since created_at
    capacity: i64,
    tokens_per_sec: i64,
    created_at: Instant,
}

impl RateLimiter {
    pub fn new(capacity: i64, tokens_per_sec: i64) -> Self {
        Self {
            tokens: AtomicI64::new(capacity),
            last_refilled_at: AtomicU64::new(0),
            capacity,
            tokens_per_sec,
            created_at: Instant::now(),
        }
    }

    pub fn try_acquire(&self) -> bool {
        let now = self.created_at.elapsed().as_secs();
        let last = self
            .last_refilled_at
            .fetch_max(now, Relaxed);
        if now > last {
            self.tokens
                .fetch_add((now - last) as i64 * self.tokens_per_sec, Relaxed);
            self.tokens
                .fetch_min(self.capacity, Relaxed);
        }

        if self.tokens.load(Relaxed) <= 0 {
            return false;
        }
        self.tokens.fetch_sub(1, Relaxed) > 0
    }
}

The design differs from the locked version in three ways:

• Tokens are whole numbers in an AtomicI64, and time is whole seconds since the limiter was created. Tokens arrive in steps, once per second, instead of continuously.
• last_refilled_at.fetch_max(now, Relaxed) records the current second and returns the previous value, in one atomic step. If several threads start a new second at once, only the first sees an older value, so only that thread adds tokens.
• #[repr(align(64))] gives the limiter its own cache lines, so it does not share a line with a neighbor, as section 15.4 described.

Then comes the check. If the count is 0 or less, the call returns false. Otherwise it takes a token with fetch_sub(1), and returns true if the count before the subtraction was positive.

Listing 19.5 The complete program. src/bin/rate_limiter_atomic_token_bucket.rs

use std::{
    sync::atomic::{AtomicI64, AtomicU64, Ordering::Relaxed},
    thread,
    time::{Duration, Instant},
};

#[repr(align(64))]
pub struct RateLimiter {
    tokens: AtomicI64,
    last_refilled_at: AtomicU64, // secs since created_at
    capacity: i64,
    tokens_per_sec: i64,
    created_at: Instant,
}

impl RateLimiter {
    pub fn new(capacity: i64, tokens_per_sec: i64) -> Self {
        Self {
            tokens: AtomicI64::new(capacity),
            last_refilled_at: AtomicU64::new(0),
            capacity,
            tokens_per_sec,
            created_at: Instant::now(),
        }
    }

    pub fn try_acquire(&self) -> bool {
        let now = self.created_at.elapsed().as_secs();
        let last = self
            .last_refilled_at
            .fetch_max(now, Relaxed);
        if now > last {
            self.tokens
                .fetch_add((now - last) as i64 * self.tokens_per_sec, Relaxed);
            self.tokens
                .fetch_min(self.capacity, Relaxed);
        }

        if self.tokens.load(Relaxed) <= 0 {
            return false;
        }
        self.tokens.fetch_sub(1, Relaxed) > 0
    }
}

fn main() {
    let limiter = RateLimiter::new(5, 5);

    for i in 1..=6 {
        println!("request {i}: {}", limiter.try_acquire());
    }

    thread::sleep(Duration::from_secs(1));

    for i in 1..=6 {
        println!("request {i}: {}", limiter.try_acquire());
    }
}

$ cargo run --bin rate_limiter_atomic_token_bucket
request 1: true
request 2: true
request 3: true
request 4: true
request 5: true
request 6: false
request 1: true
request 2: true
request 3: true
request 4: true
request 5: true
request 6: false

With one thread, the output is exactly right: 5 requests pass, the sixth fails, and after a second 5 more pass.

19.1.3 What happens with many threads

Each atomic operation is safe on its own. The problem is that try_acquire makes several of them, and other threads can act between them. Figure 19.2 shows one interleaving with one token left.

B is correctly refused, but its fetch_sub still ran. The count drops below zero, and the next refill starts from there, so the limiter lets through fewer requests than its rate. To measure this, I ran 8 threads calling try_acquire in a loop for 3.5 seconds, on a limiter with capacity 5 and rate 5. A correct limiter admits 20 requests: 5 at the start and 5 in each of the next three seconds. Three runs admitted 18, 19, and 17, and ended with the count at 0, −1, and −1.

A second gap goes the other way. The refill is a fetch_add followed by a separate fetch_min. Between the two, the count can be above the capacity, and other threads can take the extra tokens before the cap is applied. That would let through more than the capacity in one burst. My runs did not catch this window, but nothing in the code prevents it.

Both gaps close if the check, the take, and the refill happen as one atomic step. compare_exchange in a loop does that. It reads the count, computes the new count, and writes it only if no other thread changed the count in between. Exercise 2 asks you to write it.

19.2 Retrying failed calls

Many failures are temporary. A server restarts, a network link drops a packet, or a service is briefly overloaded. Trying again a moment later often works.

Retrying well takes three decisions:

1. Which errors to retry. A “not found” error will fail again. A “server busy” error may not.
2. How long to wait. Waiting a little longer after each failure gives a struggling server time to recover.
3. How many times. At some point, the caller should give up and report the error.

19.2.1 Which errors to retry

The retry code works with ApiError, the error enum from section 7.7. Here it is again. Each variant is one way a call to a service can fail, and carries the data that failure needs:

/// API error modeled as an enum with associated data.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ApiError {
    InvalidRequest(String),
    Unauthorized,
    RateLimited { retry_after: Duration },
    NotFound { resource_id: String },
    Server { status: u16, message: String },
}

impl ApiError {
    pub fn is_retryable(&self) -> bool {
        matches!(
            self,
            Self::RateLimited { .. }
                | Self::Server {
                    status: 500..=599,
                    ..
                }
        )
    }
    // ...
}

• InvalidRequest and NotFound describe a request that is wrong. Sending it again gives the same answer.
• Unauthorized needs new credentials, not a second try.
• RateLimited says the server is busy, and carries how long it asked the client to wait.
• Server is a failure on the server’s side, with the HTTP status code.

The enum’s own method, is_retryable, already answers the first question for each variant. matches! tests a value against a pattern and returns true or false. It returns true for RateLimited, and for Server with a status from 500 to 599.

The retry code should not be tied to ApiError, though. It should work with any error that can answer the question. That is a job for a trait:

Listing 19.6 The Retryable trait (lines 22 to 44). src/problems/retry.rs

use std::time::Duration;

use crate::problems::state_machine::ApiError;

/// An error that knows whether retrying could help.
pub trait Retryable {
    /// `true` if the operation is worth retrying.
    fn is_retryable(&self) -> bool;

    /// A server-provided minimum delay, such as HTTP `Retry-After`.
    fn retry_after(&self) -> Option<Duration> {
        None
    }
}

impl Retryable for ApiError {
    fn is_retryable(&self) -> bool {
        ApiError::is_retryable(self)
    }

    fn retry_after(&self) -> Option<Duration> {
        match self {
            ApiError::RateLimited { retry_after } => Some(*retry_after),
            _ => None,
        }
    }
}

Retryable is a trait that an error type implements to answer the first question. is_retryable has no default, so each error type must decide. retry_after has a default body that returns None, so an error type only needs to implement it if the server can send a hint.

The trait is implemented for ApiError. Its is_retryable passes the question to the enum’s own method. retry_after returns the delay carried by a RateLimited error. That delay usually comes from the HTTP header Retry-After, which a server sends to say when to try again.

The line ApiError::is_retryable(self) calls the enum’s own method, not the trait’s. Both have the same name. The path ApiError::is_retryable finds the inherent method first, so the trait method does not call itself forever.

19.2.2 How long to wait

A fixed delay has a problem. If a server fails, all its clients fail at the same moment. With a fixed delay they all retry at the same moment too. The server is hit by a wave of retries at the moment it recovers. This is called a thundering herd.

Two techniques spread the retries out:

• Exponential backoff doubles the delay after each failure: 100 ms, 200 ms, 400 ms, and so on, up to a cap (figure 19.3). The longer the outage, the less often each client tries.
• Jitter adds randomness to the delay, so clients that failed together do not retry together (figure 19.4). Full jitter picks a random delay between 0 and the backoff.

Listing 19.7 The policy (lines 46 to 102).

/// How randomness is applied to the backoff.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Jitter {
    /// Delay is exactly the capped exponential backoff.
    None,
    /// Delay is `unit_random * capped_backoff` (AWS "full jitter").
    Full,
}

/// Backoff configuration.
#[derive(Debug, Clone, Copy)]
pub struct RetryPolicy {
    /// Total attempts, including the first one. Must be at least 1.
    pub max_attempts: u32,
    /// Delay before the first retry.
    pub base_delay: Duration,
    /// Upper bound on a single delay.
    pub max_delay: Duration,
    pub jitter: Jitter,
}

impl RetryPolicy {
    pub fn new(
        max_attempts: u32,
        base_delay: Duration,
        max_delay: Duration,
        jitter: Jitter,
    ) -> Self {
        assert!(max_attempts > 0, "max_attempts must be > 0");
        Self {
            max_attempts,
            base_delay,
            max_delay,
            jitter,
        }
    }

    /// Capped exponential backoff for a zero-based retry attempt:
    /// `min(base * 2^attempt, max_delay)`.
    pub fn exponential_delay(&self, attempt: u32) -> Duration {
        let shift = attempt.min(31);
        let factor = 1u32 << shift;
        self.base_delay
            .saturating_mul(factor)
            .min(self.max_delay)
    }

    /// Backoff for `attempt` with `unit_random` in `[0, 1)` applied when jitter
    /// is enabled.
    pub fn delay_for(&self, attempt: u32, unit_random: f64) -> Duration {
        let capped = self.exponential_delay(attempt);
        match self.jitter {
            Jitter::None => capped,
            Jitter::Full => capped.mul_f64(unit_random.clamp(0.0, 1.0)),
        }
    }
}

RetryPolicy holds the four settings. max_attempts counts the first try as well, so 3 means one try and two retries.

exponential_delay computes base * 2^attempt, capped at max_delay. Three details keep it from overflowing:

• 1u32 << shift computes 2 to the power shift by shifting the bit 1 left. attempt.min(31) keeps the shift within a u32, because shifting by 32 or more is an error.
• Duration::saturating_mul returns the largest Duration instead of overflowing.
• .min(self.max_delay) applies the cap.

delay_for applies jitter. It takes a random number in [0, 1) as an argument, and multiplies the delay by it with Duration::mul_f64. clamp(0.0, 1.0) guards against a random source that returns something out of range.

19.2.3 The retry loop

Listing 19.8 retry (lines 104 to 137).

/// Retry `operation` until it succeeds, fails permanently, or the policy is
/// exhausted.
///
/// `unit_random` must return values in `[0, 1)`; pass `|| 0.0` for deterministic
/// tests or a real RNG in production.
pub fn retry<T, E, F, R>(policy: &RetryPolicy, mut operation: F, mut unit_random: R) -> Result<T, E>
where
    F: FnMut(u32) -> Result<T, E>,
    E: Retryable,
    R: FnMut() -> f64,
{
    let mut attempt = 0;
    loop {
        match operation(attempt) {
            Ok(value) => return Ok(value),
            Err(error) => {
                let exhausted = attempt + 1 >= policy.max_attempts;
                if exhausted || !error.is_retryable() {
                    return Err(error);
                }

                let backoff = policy.delay_for(attempt, unit_random());
                // Honor the server's hint, but never retry sooner than our own
                // backoff says is healthy.
                let delay = error
                    .retry_after()
                    .map_or(backoff, |hint| hint.max(backoff));

                std::thread::sleep(delay);
                attempt += 1;
            }
        }
    }
}

retry is generic over four types. T and E are the success and error types. F is the operation, a closure that receives the attempt number. R produces random numbers.

The operation is FnMut, not FnOnce, because the loop calls it several times, and it may change captured state, like the test’s attempt counter.

On each error, the loop checks two reasons to stop: the attempts are used up, or the error is not retryable. Either way it returns the last error. Otherwise it computes its own backoff, and then applies the server’s hint:

let delay = error.retry_after().map_or(backoff, |hint| hint.max(backoff));

map_or(default, f) returns default for None, and f(hint) for Some(hint). So without a hint, the delay is the backoff. With a hint, it is the larger of the two. The client never retries sooner than the server asked, and never sooner than its own backoff.

std::thread::sleep(delay) blocks the thread for the delay. In async code, chapter 22, the wait would be an async timer instead, so the thread can do other work.

Injecting randomness. The random number comes from the unit_random closure the caller passes in. Production code passes a real random number generator. Tests pass || 0.0, which makes every delay predictable. Passing in a dependency this way, instead of creating it inside the function, is called dependency injection. It makes the function testable without a real clock or a real random source.

Listing 19.9 The complete file, with tests. src/problems/retry.rs

//! Retry with exponential backoff, full jitter, and `Retry-After`.
//!
//! Every network client retries. Done badly it makes an outage worse: synchronized
//! retries from thousands of clients form a thundering herd, and retrying a
//! non-idempotent `POST` can create duplicate workloads. This module models the
//! policy explicitly.
//!
//! Two details:
//!
//! - **Jitter.** Full jitter picks a random delay in `[0, capped_backoff]` so
//!   clients spread out instead of retrying in lockstep.
//! - **`Retry-After`.** When the server tells you when to come back, wait at
//!   least that long; never retry sooner than the server asked.
//!
//! The random source is injected as a closure so the delay math is deterministic
//! under test.

use std::time::Duration;

use crate::problems::state_machine::ApiError;

/// An error that knows whether retrying could help.
pub trait Retryable {
    /// `true` if the operation is worth retrying.
    fn is_retryable(&self) -> bool;

    /// A server-provided minimum delay, such as HTTP `Retry-After`.
    fn retry_after(&self) -> Option<Duration> {
        None
    }
}

impl Retryable for ApiError {
    fn is_retryable(&self) -> bool {
        ApiError::is_retryable(self)
    }

    fn retry_after(&self) -> Option<Duration> {
        match self {
            ApiError::RateLimited { retry_after } => Some(*retry_after),
            _ => None,
        }
    }
}

/// How randomness is applied to the backoff.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Jitter {
    /// Delay is exactly the capped exponential backoff.
    None,
    /// Delay is `unit_random * capped_backoff` (AWS "full jitter").
    Full,
}

/// Backoff configuration.
#[derive(Debug, Clone, Copy)]
pub struct RetryPolicy {
    /// Total attempts, including the first one. Must be at least 1.
    pub max_attempts: u32,
    /// Delay before the first retry.
    pub base_delay: Duration,
    /// Upper bound on a single delay.
    pub max_delay: Duration,
    pub jitter: Jitter,
}

impl RetryPolicy {
    pub fn new(
        max_attempts: u32,
        base_delay: Duration,
        max_delay: Duration,
        jitter: Jitter,
    ) -> Self {
        assert!(max_attempts > 0, "max_attempts must be > 0");
        Self {
            max_attempts,
            base_delay,
            max_delay,
            jitter,
        }
    }

    /// Capped exponential backoff for a zero-based retry attempt:
    /// `min(base * 2^attempt, max_delay)`.
    pub fn exponential_delay(&self, attempt: u32) -> Duration {
        let shift = attempt.min(31);
        let factor = 1u32 << shift;
        self.base_delay
            .saturating_mul(factor)
            .min(self.max_delay)
    }

    /// Backoff for `attempt` with `unit_random` in `[0, 1)` applied when jitter
    /// is enabled.
    pub fn delay_for(&self, attempt: u32, unit_random: f64) -> Duration {
        let capped = self.exponential_delay(attempt);
        match self.jitter {
            Jitter::None => capped,
            Jitter::Full => capped.mul_f64(unit_random.clamp(0.0, 1.0)),
        }
    }
}

/// Retry `operation` until it succeeds, fails permanently, or the policy is
/// exhausted.
///
/// `unit_random` must return values in `[0, 1)`; pass `|| 0.0` for deterministic
/// tests or a real RNG in production.
pub fn retry<T, E, F, R>(policy: &RetryPolicy, mut operation: F, mut unit_random: R) -> Result<T, E>
where
    F: FnMut(u32) -> Result<T, E>,
    E: Retryable,
    R: FnMut() -> f64,
{
    let mut attempt = 0;
    loop {
        match operation(attempt) {
            Ok(value) => return Ok(value),
            Err(error) => {
                let exhausted = attempt + 1 >= policy.max_attempts;
                if exhausted || !error.is_retryable() {
                    return Err(error);
                }

                let backoff = policy.delay_for(attempt, unit_random());
                // Honor the server's hint, but never retry sooner than our own
                // backoff says is healthy.
                let delay = error
                    .retry_after()
                    .map_or(backoff, |hint| hint.max(backoff));

                std::thread::sleep(delay);
                attempt += 1;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn policy(max_attempts: u32, jitter: Jitter) -> RetryPolicy {
        RetryPolicy::new(
            max_attempts,
            Duration::from_millis(100),
            Duration::from_millis(1_000),
            jitter,
        )
    }

    #[test]
    fn exponential_backoff_doubles_until_capped() {
        let policy = policy(6, Jitter::None);
        assert_eq!(policy.exponential_delay(0), Duration::from_millis(100));
        assert_eq!(policy.exponential_delay(1), Duration::from_millis(200));
        assert_eq!(policy.exponential_delay(2), Duration::from_millis(400));
        assert_eq!(policy.exponential_delay(3), Duration::from_millis(800));
        // Capped from here on.
        assert_eq!(policy.exponential_delay(4), Duration::from_millis(1_000));
        assert_eq!(policy.exponential_delay(50), Duration::from_millis(1_000));
    }

    #[test]
    fn full_jitter_scales_the_delay() {
        let policy = policy(6, Jitter::Full);
        assert_eq!(policy.delay_for(2, 0.0), Duration::ZERO);
        assert_eq!(policy.delay_for(2, 0.5), Duration::from_millis(200));
        // Values above 1.0 are clamped rather than overflowing.
        assert_eq!(policy.delay_for(2, 9.9), Duration::from_millis(400));
    }

    #[test]
    fn retries_then_succeeds() {
        let policy = RetryPolicy::new(5, Duration::ZERO, Duration::ZERO, Jitter::None);
        let mut attempts = 0;
        let result: Result<u32, ApiError> = retry(
            &policy,
            |_| {
                attempts += 1;
                if attempts < 3 {
                    Err(ApiError::Server {
                        status: 503,
                        message: "temporarily unavailable".into(),
                    })
                } else {
                    Ok(42)
                }
            },
            || 0.0,
        );

        assert_eq!(result, Ok(42));
        assert_eq!(attempts, 3);
    }

    #[test]
    fn non_retryable_errors_return_immediately() {
        let policy = RetryPolicy::new(5, Duration::ZERO, Duration::ZERO, Jitter::None);
        let mut attempts = 0;
        let result: Result<(), ApiError> = retry(
            &policy,
            |_| {
                attempts += 1;
                Err(ApiError::NotFound {
                    resource_id: "item-0".into(),
                })
            },
            || 0.0,
        );

        assert!(result.is_err());
        assert_eq!(attempts, 1);
    }

    #[test]
    fn exhausted_policy_returns_the_last_error() {
        let policy = RetryPolicy::new(3, Duration::ZERO, Duration::ZERO, Jitter::None);
        let mut attempts = 0;
        let result: Result<(), ApiError> = retry(
            &policy,
            |_| {
                attempts += 1;
                Err(ApiError::Server {
                    status: 500,
                    message: "boom".into(),
                })
            },
            || 0.0,
        );

        assert!(result.is_err());
        assert_eq!(attempts, 3);
    }

    #[test]
    fn retry_after_is_honored_as_a_floor() {
        let policy = policy(2, Jitter::None);
        let error = ApiError::RateLimited {
            retry_after: Duration::from_secs(5),
        };
        let backoff = policy.delay_for(0, 0.0);
        let chosen = error
            .retry_after()
            .map_or(backoff, |hint| hint.max(backoff));
        assert_eq!(chosen, Duration::from_secs(5));
    }
}

The tests use a policy with zero delays, so they run instantly. retries_then_succeeds fails twice with a 503 and then succeeds: three attempts in total. non_retryable_errors_return_immediately returns NotFound, and must stop after one attempt.

$ cargo test --lib problems::retry
running 6 tests
test problems::retry::tests::full_jitter_scales_the_delay ... ok
test problems::retry::tests::exponential_backoff_doubles_until_capped ... ok
test problems::retry::tests::retry_after_is_honored_as_a_floor ... ok
test problems::retry::tests::non_retryable_errors_return_immediately ... ok
test problems::retry::tests::exhausted_policy_returns_the_last_error ... ok
test problems::retry::tests::retries_then_succeeds ... ok

test result: ok. 6 passed; 0 failed; 0 ignored; 0 measured; 131 filtered out

19.3 Idempotency keys

Retries have a danger. Suppose a client asks a server to charge a card, and the reply is lost. The client cannot tell whether the charge happened. If it retries, the card may be charged twice.

An operation is idempotent if doing it twice has the same effect as doing it once. Reading a record is idempotent. Charging a card is not, unless the server makes it so.

The usual way is an idempotency key. The client picks a unique key for each logical request, such as an order number, and sends it with every attempt. The server records the result under that key. When the same key arrives again, the server returns the recorded result and does not repeat the work (figure 19.5).

The store is a HashMap from key to result, behind a Mutex so that several threads can share it:

use std::{collections::HashMap, hash::Hash, sync::Mutex};

pub struct Idempotent<K, V> {
    store: Mutex<HashMap<K, V>>,
}

Listing 19.10 The complete program. src/bin/idempotent_operation.rs

use std::{collections::HashMap, hash::Hash, sync::Mutex};

pub struct Idempotent<K, V> {
    store: Mutex<HashMap<K, V>>,
}

impl<K, V> Idempotent<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    pub fn new() -> Self {
        Self {
            store: Mutex::new(HashMap::new()),
        }
    }

    pub fn execute<F>(&self, key: K, f: F) -> Result<V, Box<dyn std::error::Error>>
    where
        F: FnOnce() -> Result<V, Box<dyn std::error::Error>>,
    {
        let mut store = self
            .store
            .lock()
            .map_err(|_| "lock poisoned")?;
        if let Some(v) = store.get(&key) {
            return Ok(v.clone());
        }
        let v = f()?;
        store.insert(key, v.clone());
        Ok(v)
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let idem = Idempotent::new();

    let charge = || -> Result<String, Box<dyn std::error::Error>> {
        println!("charging card...");
        Ok("txn_42".into())
    };

    assert_eq!(idem.execute("order-1", charge)?, "txn_42");
    assert_eq!(idem.execute("order-1", charge)?, "txn_42");
    Ok(())
}

Idempotent<K, V> stores results in a HashMap behind a Mutex. execute takes a key and a closure f that does the work:

1. It locks the store. A poisoned lock becomes an error: map_err(|_| "lock poisoned")? turns the PoisonError into a &str, and ? converts that into a Box<dyn Error>.
2. If the key is present, it returns a clone of the stored result.
3. Otherwise it runs f(). If f fails, ? returns the error, and nothing is stored, so a later call can try again.
4. It stores a clone of the result and returns it.

In main, charge is called with the same key twice. The closure captures nothing, so it is Copy, and passing it to execute twice copies it. The line “charging card…” appears once:

$ cargo run --bin idempotent_operation
charging card...

The second file is the same program with a type alias, type RuntimeError = Box<dyn std::error::Error>;, which shortens the signatures:

Listing 19.11 The complete program. src/bin/idempotent_operation_with_error_progagation.rs

use std::{collections::HashMap, hash::Hash, sync::Mutex};

type RuntimeError = Box<dyn std::error::Error>;

pub struct Idempotent<K, V> {
    store: Mutex<HashMap<K, V>>,
}

impl<K, V> Idempotent<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    pub fn new() -> Self {
        Self {
            store: Mutex::new(HashMap::new()),
        }
    }

    pub fn execute<F>(&self, key: K, f: F) -> Result<V, RuntimeError>
    where
        F: FnOnce() -> Result<V, RuntimeError>,
    {
        let mut store = self
            .store
            .lock()
            .map_err(|_| "lock poisoned")?;
        if let Some(v) = store.get(&key) {
            return Ok(v.clone());
        }
        let v = f()?;
        store.insert(key, v.clone());
        Ok(v)
    }
}

fn main() -> Result<(), RuntimeError> {
    let idem = Idempotent::new();

    let charge = || -> Result<String, RuntimeError> {
        println!("charging card...");
        Ok("txn_42".into())
    };

    assert_eq!(idem.execute("order-1", charge)?, "txn_42");
    assert_eq!(idem.execute("order-1", charge)?, "txn_42");
    Ok(())
}

$ cargo run --bin idempotent_operation_with_error_progagation
charging card...

Holding the lock while f() runs makes the store correct. Two calls with the same key cannot both find it missing and both run the work. It is also the store’s limit. Every call waits for the one lock, even calls with different keys, so one slow operation delays all others. Section 19.4 removes that limit. A real service also keeps the results in a database, so they survive a restart.

19.4 One execution per key: single flight

The store in section 19.3 has two problems under load.

• It serializes every key. It holds one lock while the work runs. A slow charge for order-1 makes a call for order-2 wait, although the two have nothing in common.
• It cannot be split naively. Suppose the lock is released during the work. Two callers with the same key can then both find it missing, and both charge the card.

The same shape appears in caches. When a popular entry expires, hundreds of requests can miss at once, and each recomputes the value or queries the database. That burst is a cache stampede, a thundering herd aimed at one key.

The fix is to record, for each key, that work is in progress. The first caller for a key marks it, releases the lock, and runs the work. A caller that finds the mark waits for the result instead of starting its own. Callers with other keys never wait. This is called single flight: at most one execution per key is in the air at a time.

Each key moves through three states (figure 19.6). A success stores the value for every later caller. A failure, or a panic, removes the mark, so the next caller can try again.

19.4.1 The state

Each key maps to a Slot, and all the slots share one map behind one mutex. A Condvar wakes the waiters when any slot changes:

use std::{
    collections::HashMap,
    hash::Hash,
    panic::{self, AssertUnwindSafe},
    sync::{
        Barrier,
        Condvar,
        Mutex,
        MutexGuard,
        atomic::{AtomicUsize, Ordering::Relaxed},
    },
    thread,
    time::{Duration, Instant},
};

type Error = Box<dyn std::error::Error + Send + Sync>;

/// The state of one key.
enum Slot<V> {
    InProgress,
    Done(V),
}

/// Results by key, with at most one execution in flight per key.
pub struct SingleFlight<K, V> {
    slots: Mutex<HashMap<K, Slot<V>>>,
    finished: Condvar,
}

The mutex is held only to read or change the map, never while the work runs. So one lock is enough for every key. notify_all wakes every waiter, for any key. Each one checks its own key and goes back to sleep if nothing it cares about changed. That costs a few spurious wakeups, and keeps the code to one Condvar.

19.4.2 The claim

The caller that runs the work holds a claim on its key. If the work fails or panics, someone has to remove the InProgress mark, or the waiters sleep forever. Drop is the place that runs in both cases:

/// The claim of the caller that is running the work for `key`. If it is
/// dropped while it still holds the key, after an error or a panic, it clears
/// the `InProgress` mark and wakes the waiters, so one of them can run the work.
struct Claim<'a, K: Eq + Hash, V> {
    flight: &'a SingleFlight<K, V>,
    key: Option<K>,
}

impl<K: Eq + Hash, V> Drop for Claim<'_, K, V> {
    fn drop(&mut self) {
        if let Some(key) = self.key.take() {
            self.flight.lock().remove(&key);
            self.flight.finished.notify_all();
        }
    }
}

impl<K: Eq + Hash, V> SingleFlight<K, V> {
    /// The map, even if a thread panicked while holding the lock: every
    /// update leaves the map consistent, so its contents are still valid.
    fn lock(&self) -> MutexGuard<'_, HashMap<K, Slot<V>>> {
        self.slots
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }
}

key is an Option. The success path takes the key out of the claim before the claim is dropped, so Drop does nothing. On an error or a panic, the key is still inside, and Drop removes the mark and wakes the waiters.

lock recovers a poisoned mutex with into_inner, as in section 16.6. The map is consistent after every update, because each update is one insert or one remove, so the data is safe to keep using.

19.4.3 execute

execute loops while the key is in progress, then claims it, runs the work outside the lock, and stores the result:

impl<K, V> SingleFlight<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    pub fn new() -> Self {
        Self {
            slots: Mutex::new(HashMap::new()),
            finished: Condvar::new(),
        }
    }

    /// Return the stored result for `key`, wait for the execution in flight,
    /// or run `work` if neither exists.
    pub fn execute<F>(&self, key: K, work: F) -> Result<V, Error>
    where
        F: FnOnce() -> Result<V, Error>,
    {
        let mut slots = self.lock();
        loop {
            match slots.get(&key) {
                Some(Slot::Done(value)) => return Ok(value.clone()),
                Some(Slot::InProgress) => {
                    slots = self
                        .finished
                        .wait(slots)
                        .unwrap_or_else(|poisoned| poisoned.into_inner());
                }
                None => break,
            }
        }
        slots.insert(key.clone(), Slot::InProgress);
        drop(slots);

        let mut claim = Claim {
            flight: self,
            key: Some(key),
        };
        let value = work()?;

        let key = claim
            .key
            .take()
            .expect("claim holds its key");
        self.lock()
            .insert(key, Slot::Done(value.clone()));
        self.finished.notify_all();
        Ok(value)
    }
}

The loop handles all three states:

• Done(value): return a clone. This is the idempotency of section 19.3.
• InProgress: wait on the Condvar. wait releases the lock while it sleeps, and takes it back on wakeup. The loop then checks the key again, because the wakeup may be for another key, or spurious.
• No entry: leave the loop and claim the key.

After the claim, drop(slots) releases the lock, and work() runs with no lock held. If work returns an error, ? returns it, and the claim’s Drop removes the mark. If work succeeds, the code takes the key out of the claim, stores Done, and wakes the waiters.

19.4.4 Measured

main starts its callers together with a Barrier, so they really overlap, and times them:

/// Run `call` on `threads` threads that start together. Returns the results
/// and the time from the start to the last return.
fn together<T: Send>(threads: usize, call: impl Fn(usize) -> T + Sync) -> (Vec<T>, Duration) {
    let start_line = Barrier::new(threads);
    let start = Instant::now();
    let results = thread::scope(|scope| {
        let handles: Vec<_> = (0..threads)
            .map(|i| {
                let (start_line, call) = (&start_line, &call);
                scope.spawn(move || {
                    start_line.wait();
                    call(i)
                })
            })
            .collect();
        handles
            .into_iter()
            .map(|handle| handle.join().expect("caller thread"))
            .collect()
    });
    (results, start.elapsed())
}

The work is a fake charge that sleeps for 100 ms and counts how many times it ran. The second scene also runs the four keys through a store like section 19.3’s: one lock, held while the work runs.

1. Eight concurrent calls with one key
  executions 1, all got txn-order-1: true, 103 ms

2. Four concurrent calls with four keys
  single flight:                 executions 4, 103 ms
  one lock held during the work: executions 4, 417 ms

3. A failure is not stored
  first: Err("card declined"), retry: Ok("txn-order-5"), executions 1

Eight callers with one key caused one charge, and all eight got its result in 103 ms. Four keys ran in parallel in 103 ms. The one-lock store ran them one after another, in 417 ms. A declined charge was not stored, so the retry ran the work, and that run succeeded.

The fourth scene makes the work panic while a second caller waits for the same key:

4. The work panics while another caller waits

thread '<unnamed>' panicked at src/bin/single_flight.rs:209:21:
payment provider crashed
  caller 0: Err("panicked")
  caller 1: Ok(Ok("txn-order-6"))

The panicked at lines are the default panic message, written to standard error. Caller 0’s panic dropped its claim, and Drop removed the mark and woke caller 1. Caller 1 found no entry, claimed the key, and ran the work itself. Without the claim’s Drop, the mark would stay InProgress, and caller 1 would wait forever. Animation 19.2 shows both runs.

A production version adds two things. Waiters need a deadline, wait_timeout, so a hung charge cannot hold them forever. And the results need an expiry, or a store that only grows becomes a memory leak.

Listing 19.12 Single flight: one execution per key in flight. src/bin/single_flight.rs

//! Single flight: concurrent calls with the same key share one execution.
//!
//! The first caller marks the key in progress and runs the work without
//! holding the lock. Later callers with that key wait for its result, while
//! calls with other keys run in parallel. A result is kept, so a retry after
//! success gets the stored result. A failure is not kept, so a retry runs the
//! work again. A panic in the work clears the mark, so waiters are not stranded.
//!
//! Run with: cargo run --bin single_flight

use std::{
    collections::HashMap,
    hash::Hash,
    panic::{self, AssertUnwindSafe},
    sync::{
        Barrier,
        Condvar,
        Mutex,
        MutexGuard,
        atomic::{AtomicUsize, Ordering::Relaxed},
    },
    thread,
    time::{Duration, Instant},
};

type Error = Box<dyn std::error::Error + Send + Sync>;

/// The state of one key.
enum Slot<V> {
    InProgress,
    Done(V),
}

/// Results by key, with at most one execution in flight per key.
pub struct SingleFlight<K, V> {
    slots: Mutex<HashMap<K, Slot<V>>>,
    finished: Condvar,
}

/// The claim of the caller that is running the work for `key`. If it is
/// dropped while it still holds the key, after an error or a panic, it clears
/// the `InProgress` mark and wakes the waiters, so one of them can run the work.
struct Claim<'a, K: Eq + Hash, V> {
    flight: &'a SingleFlight<K, V>,
    key: Option<K>,
}

impl<K: Eq + Hash, V> Drop for Claim<'_, K, V> {
    fn drop(&mut self) {
        if let Some(key) = self.key.take() {
            self.flight.lock().remove(&key);
            self.flight.finished.notify_all();
        }
    }
}

impl<K: Eq + Hash, V> SingleFlight<K, V> {
    /// The map, even if a thread panicked while holding the lock: every
    /// update leaves the map consistent, so its contents are still valid.
    fn lock(&self) -> MutexGuard<'_, HashMap<K, Slot<V>>> {
        self.slots
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
    }
}

impl<K, V> SingleFlight<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    pub fn new() -> Self {
        Self {
            slots: Mutex::new(HashMap::new()),
            finished: Condvar::new(),
        }
    }

    /// Return the stored result for `key`, wait for the execution in flight,
    /// or run `work` if neither exists.
    pub fn execute<F>(&self, key: K, work: F) -> Result<V, Error>
    where
        F: FnOnce() -> Result<V, Error>,
    {
        let mut slots = self.lock();
        loop {
            match slots.get(&key) {
                Some(Slot::Done(value)) => return Ok(value.clone()),
                Some(Slot::InProgress) => {
                    slots = self
                        .finished
                        .wait(slots)
                        .unwrap_or_else(|poisoned| poisoned.into_inner());
                }
                None => break,
            }
        }
        slots.insert(key.clone(), Slot::InProgress);
        drop(slots);

        let mut claim = Claim {
            flight: self,
            key: Some(key),
        };
        let value = work()?;

        let key = claim
            .key
            .take()
            .expect("claim holds its key");
        self.lock()
            .insert(key, Slot::Done(value.clone()));
        self.finished.notify_all();
        Ok(value)
    }
}

/// Run `call` on `threads` threads that start together. Returns the results
/// and the time from the start to the last return.
fn together<T: Send>(threads: usize, call: impl Fn(usize) -> T + Sync) -> (Vec<T>, Duration) {
    let start_line = Barrier::new(threads);
    let start = Instant::now();
    let results = thread::scope(|scope| {
        let handles: Vec<_> = (0..threads)
            .map(|i| {
                let (start_line, call) = (&start_line, &call);
                scope.spawn(move || {
                    start_line.wait();
                    call(i)
                })
            })
            .collect();
        handles
            .into_iter()
            .map(|handle| handle.join().expect("caller thread"))
            .collect()
    });
    (results, start.elapsed())
}

fn main() -> Result<(), Error> {
    let charges = AtomicUsize::new(0);
    let charge = |order: &str| -> Result<String, Error> {
        charges.fetch_add(1, Relaxed);
        thread::sleep(Duration::from_millis(100));
        Ok(format!("txn-{order}"))
    };

    println!("1. Eight concurrent calls with one key");
    let flight = SingleFlight::new();
    let (results, elapsed) = together(8, |_| flight.execute("order-1", || charge("order-1")));
    let agreed = results
        .iter()
        .all(|result| matches!(result, Ok(txn) if txn == "txn-order-1"));
    println!(
        "  executions {}, all got txn-order-1: {agreed}, {:.0} ms",
        charges.swap(0, Relaxed),
        elapsed.as_secs_f64() * 1000.0
    );

    println!("\n2. Four concurrent calls with four keys");
    let orders = ["order-1", "order-2", "order-3", "order-4"];
    let flight = SingleFlight::new();
    let (_, elapsed) = together(4, |i| flight.execute(orders[i], || charge(orders[i])));
    println!(
        "  single flight:                 executions {}, {:.0} ms",
        charges.swap(0, Relaxed),
        elapsed.as_secs_f64() * 1000.0
    );

    // The store of section 19.3: one lock, held while the work runs.
    let one_lock: Mutex<HashMap<&str, String>> = Mutex::new(HashMap::new());
    let (_, elapsed) = together(4, |i| -> Result<String, Error> {
        let mut store = one_lock
            .lock()
            .map_err(|_| "lock poisoned")?;
        if let Some(txn) = store.get(orders[i]) {
            return Ok(txn.clone());
        }
        let txn = charge(orders[i])?;
        store.insert(orders[i], txn.clone());
        Ok(txn)
    });
    println!(
        "  one lock held during the work: executions {}, {:.0} ms",
        charges.swap(0, Relaxed),
        elapsed.as_secs_f64() * 1000.0
    );

    println!("\n3. A failure is not stored");
    let flight = SingleFlight::new();
    let declined = flight.execute("order-5", || Err("card declined".into()));
    let retried = flight.execute("order-5", || charge("order-5"));
    println!(
        "  first: {declined:?}, retry: {retried:?}, executions {}",
        charges.swap(0, Relaxed)
    );

    println!("\n4. The work panics while another caller waits");
    let flight = SingleFlight::new();
    let (results, _) = together(2, |i| {
        if i == 1 {
            thread::sleep(Duration::from_millis(20));
        }
        panic::catch_unwind(AssertUnwindSafe(|| {
            flight.execute("order-6", || {
                if i == 0 {
                    thread::sleep(Duration::from_millis(50));
                    panic!("payment provider crashed");
                }
                charge("order-6")
            })
        }))
        .map_err(|_| "panicked")
    });
    for (i, result) in results.iter().enumerate() {
        println!("  caller {i}: {result:?}");
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn concurrent_calls_with_one_key_run_once() {
        let runs = AtomicUsize::new(0);
        let flight = SingleFlight::new();
        let (results, _) = together(8, |_| {
            flight.execute("k", || {
                runs.fetch_add(1, Relaxed);
                thread::sleep(Duration::from_millis(30));
                Ok(7)
            })
        });
        assert_eq!(runs.load(Relaxed), 1);
        assert!(
            results
                .iter()
                .all(|result| matches!(result, Ok(7)))
        );
    }

    #[test]
    fn a_failure_lets_the_next_call_run() {
        let flight = SingleFlight::new();
        assert!(
            flight
                .execute("k", || Err("no".into()))
                .is_err()
        );
        assert_eq!(flight.execute("k", || Ok(1)).unwrap(), 1);
        assert_eq!(flight.execute("k", || Ok(2)).unwrap(), 1);
    }
}

Chapter 20 moves down to the network itself: IP addresses, sockets, and a server that echoes back what it receives.

Exercises

1. Change TokenBucket::try_acquire to take the current time as an argument, try_acquire_at(&self, now: Instant). Rewrite the test without thread::sleep.
2. Rewrite the atomic limiter’s try_acquire with a compare_exchange loop that refills, checks, and takes a token in one atomic step. Run the 8-thread test from section 19.1.3 on it.
3. Add a max_elapsed: Duration to RetryPolicy, and make retry stop when the total time spent would pass it.
4. Add a deadline to SingleFlight::execute: a caller that waits longer than a given Duration for a key in progress returns an error. Use Condvar::wait_timeout, and keep the remaining time across spurious wakeups.