Section C · Coding · Drill mode

Coding Problems Worked

10 problems with the shape interviewers actually ask for AI infra Rust roles. Multiple approaches per problem; drill on a timer.

P1 — Rate limiter (token bucket) in Rust

Prompt: Implement a token bucket rate limiter. fn try_acquire(n: u32) -> bool returns true if N tokens are available (and consumes them), false otherwise. Bucket refills at rate tokens/sec up to capacity.

Approach 1 — single-process, lazy refill

use std::sync::Mutex;
use std::time::Instant;

pub struct TokenBucket {
    inner: Mutex<Inner>,
    rate: f64,       // tokens per second
    capacity: f64,
}

struct Inner {
    tokens: f64,
    last: Instant,
}

impl TokenBucket {
    pub fn new(rate: f64, capacity: f64) -> Self {
        Self {
            inner: Mutex::new(Inner { tokens: capacity, last: Instant::now() }),
            rate, capacity,
        }
    }

    pub fn try_acquire(&self, n: u32) -> bool {
        let mut g = self.inner.lock().unwrap();
        let now = Instant::now();
        let elapsed = (now - g.last).as_secs_f64();
        g.tokens = (g.tokens + elapsed * self.rate).min(self.capacity);
        g.last = now;
        let need = n as f64;
        if g.tokens >= need {
            g.tokens -= need;
            true
        } else {
            false
        }
    }
}

Tradeoffs: One Mutex per bucket. Fine for moderate contention; hot bucket → contention → bottleneck. Refill is lazy (computed on access), no background timer.

Approach 2 — atomic, lock-free (advanced)

Pack (last_refill_nanos: u64, tokens_micro: u32) into an AtomicU64 and CAS-loop. Lock-free under contention. Worth mentioning, usually not worth implementing in 30 minutes.

Approach 3 — distributed (Redis)

Lua script that does the read-refill-decrement atomically in Redis. Use for cross-process or cross-pod rate limits. Latency cost (RTT per check) — cache the decision locally for short windows.

Variations interviewers ask
  • "Async version that awaits until tokens available" — uses tokio::sync::Notify or a sleep based on calculated wait time.
  • "Per-key rate limit" — sharded DashMap.
  • "Sliding window vs fixed window" — token bucket is conceptually sliding; fixed-window counters are simpler but burstier.

P2 — Async LRU cache

Prompt: Thread-safe async LRU cache with get and put. Bounded by capacity; least-recently-used evicted on insert when full.

Approach — wrap a sync LRU in a Mutex

The simplest correct answer:

use std::num::NonZeroUsize;
use lru::LruCache;
use parking_lot::Mutex;
use std::sync::Arc;

pub struct AsyncLru<K: Eq + std::hash::Hash, V: Clone> {
    inner: Arc<Mutex<LruCache<K, V>>>,
}

impl<K: Eq + std::hash::Hash, V: Clone> AsyncLru<K, V> {
    pub fn new(cap: usize) -> Self {
        Self { inner: Arc::new(Mutex::new(LruCache::new(NonZeroUsize::new(cap).unwrap()))) }
    }
    pub async fn get(&self, k: &K) -> Option<V> {
        self.inner.lock().get(k).cloned()
    }
    pub async fn put(&self, k: K, v: V) {
        self.inner.lock().put(k, v);
    }
}

Note: parking_lot::Mutex is sync (no .await in the critical section), which is what we want — the operation is microseconds.

Senior variant — coalescing get_or_compute

Three concurrent gets for the same missing key should produce one underlying compute, not three. Pattern: a DashMap<K, broadcast::Sender<V>> of in-flight computes.

use tokio::sync::broadcast;
use dashmap::DashMap;

pub async fn get_or_compute<F, Fut, K, V>(
    cache: &Arc<Mutex<LruCache<K, V>>>,
    inflight: &DashMap<K, broadcast::Sender<V>>,
    key: K,
    compute: F,
) -> V
where
    K: Eq + std::hash::Hash + Clone,
    V: Clone,
    F: FnOnce() -> Fut,
    Fut: std::future::Future<Output = V>,
{
    if let Some(v) = cache.lock().get(&key).cloned() { return v; }
    if let Some(tx) = inflight.get(&key) {
        let mut rx = tx.subscribe();
        drop(tx);
        return rx.recv().await.unwrap();
    }
    let (tx, _) = broadcast::channel(1);
    inflight.insert(key.clone(), tx.clone());
    let v = compute().await;
    cache.lock().put(key.clone(), v.clone());
    inflight.remove(&key);
    let _ = tx.send(v.clone());
    v
}

P3 — gRPC streaming aggregator

Prompt: A gRPC bidi-streaming endpoint. Clients stream a sequence of numbers, the server emits running sum + count after each. Cancellation when client closes.

use tonic::{Request, Response, Status, Streaming};
use tokio_stream::wrappers::ReceiverStream;
use futures::StreamExt;

#[tonic::async_trait]
impl Agg for AggService {
    type StreamSumStream = ReceiverStream<Result<Update, Status>>;

    async fn stream_sum(&self, req: Request<Streaming<Num>>)
        -> Result<Response<Self::StreamSumStream>, Status>
    {
        let mut input = req.into_inner();
        let (tx, rx) = tokio::sync::mpsc::channel(16);

        tokio::spawn(async move {
            let mut sum: i64 = 0;
            let mut count: u64 = 0;
            while let Some(item) = input.next().await {
                match item {
                    Ok(n) => {
                        sum += n.value as i64;
                        count += 1;
                        if tx.send(Ok(Update { sum, count })).await.is_err() {
                            // client gone — drop the stream
                            break;
                        }
                    }
                    Err(e) => { let _ = tx.send(Err(e)).await; break; }
                }
            }
        });

        Ok(Response::new(ReceiverStream::new(rx)))
    }
}

Things they probe: What happens if the client disconnects mid-stream? (The send fails; spawned task exits; input.next() would return None.) What about backpressure? (The bounded channel blocks the producer if the consumer is slow.) How do you cancel cleanly? (Drop the spawned task by dropping tx; the receiver task's recv returns None.)

P4 — Request dedup / coalescing

Prompt: Given a function fn compute(key: K) -> V that's expensive, build a wrapper such that concurrent calls for the same key produce one underlying compute, all callers receive the same result.

This is the singleflight pattern. The senior variant in P2's get_or_compute is the answer. Variations:

  • "What if compute fails?" — propagate the error to all subscribers; don't cache the failure (or cache with shorter TTL).
  • "What if compute is cancelled by the original caller?" — pick: cancel all subscribers, or have a subscriber take over the work. Production choice is usually "let it complete" by spawning the compute as a detached task.
  • "How do you guard against unbounded inflight set growth?" — TTL + cap; expire stale entries.

P5 — Channel-based worker pool

Prompt: N workers, one input channel, graceful shutdown.

use tokio::sync::mpsc;
use std::sync::Arc;

pub struct Pool { tx: mpsc::Sender<Job> }

pub struct Job(pub Box<dyn FnOnce() -> futures::future::BoxFuture<'static, ()> + Send>);

impl Pool {
    pub fn new(workers: usize, cap: usize) -> Self {
        let (tx, rx) = mpsc::channel::<Job>(cap);
        let rx = Arc::new(tokio::sync::Mutex::new(rx));
        for _ in 0..workers {
            let rx = rx.clone();
            tokio::spawn(async move {
                loop {
                    let job_opt = { rx.lock().await.recv().await };
                    match job_opt {
                        Some(Job(f)) => f().await,
                        None         => break,
                    }
                }
            });
        }
        Self { tx }
    }

    pub async fn submit<F, Fut>(&self, f: F) -> Result<(), &'static str>
    where
        F: FnOnce() -> Fut + Send + 'static,
        Fut: std::future::Future<Output = ()> + Send + 'static,
    {
        let job = Job(Box::new(move || Box::pin(f())));
        self.tx.send(job).await.map_err(|_| "pool closed")
    }
    // Drop pool's tx -> channel closes -> workers exit -> graceful shutdown
}

P6 — Lock-free counter

Prompt: Multi-thread counter, max throughput. Methods: incr, get.

use std::sync::atomic::{AtomicU64, Ordering};

pub struct Counter { n: AtomicU64 }

impl Counter {
    pub fn new() -> Self { Self { n: AtomicU64::new(0) } }
    pub fn incr(&self) { self.n.fetch_add(1, Ordering::Relaxed); }
    pub fn get(&self) -> u64 { self.n.load(Ordering::Relaxed) }
}

Senior follow-up: "This contends under heavy load — every fetch_add is a cache-line bounce. How would you scale it?" Answer: sharded counters. One counter per CPU (or per thread), summed lazily on get.

pub struct ShardedCounter { shards: Vec<AtomicU64> }

impl ShardedCounter {
    pub fn new(n: usize) -> Self {
        Self { shards: (0..n).map(|_| AtomicU64::new(0)).collect() }
    }
    pub fn incr(&self) {
        // pick a shard by thread id hash; cheap and contention-distributing
        let idx = thread_local_index() % self.shards.len();
        self.shards[idx].fetch_add(1, Ordering::Relaxed);
    }
    pub fn get(&self) -> u64 {
        self.shards.iter().map(|s| s.load(Ordering::Relaxed)).sum()
    }
}
fn thread_local_index() -> usize { /* thread_local! cell with rand index */ 0 }

Memory ordering note: Relaxed is correct for monotonic counters. Use Acquire/Release when the counter synchronizes with other data.

P7 — Retry-with-backoff combinator

Prompt: Higher-order async function: retry an operation up to N times with exponential backoff + jitter. Only retry errors deemed retryable.

use rand::Rng;
use std::time::Duration;
use tokio::time::sleep;

pub trait Retryable { fn is_retryable(&self) -> bool; }

pub async fn with_backoff<F, Fut, T, E>(
    mut op: F,
    max_attempts: u32,
    base_ms: u64,
    cap_ms: u64,
) -> Result<T, E>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<Output = Result<T, E>>,
    E: Retryable,
{
    let mut delay = base_ms;
    for attempt in 1..=max_attempts {
        match op().await {
            Ok(v) => return Ok(v),
            Err(e) if attempt == max_attempts => return Err(e),
            Err(e) if !e.is_retryable() => return Err(e),
            Err(_) => {
                let jitter: u64 = rand::thread_rng().gen_range(0..delay.max(1));
                sleep(Duration::from_millis(delay + jitter)).await;
                delay = (delay * 2).min(cap_ms);
            }
        }
    }
    unreachable!()
}

Follow-ups: "What if the operation has side effects?" → caller passes an idempotency key. "What if we want a global retry budget?" → wrap with a budget check (see 08). "Full jitter vs decorrelated jitter?" → both reduce thundering-herd; decorrelated jitter is the AWS-recommended default.

P8 — Structured cancellation cascade

Prompt: An agent run launches three sub-tasks (search, summarize, persist). If the run is cancelled, all three must stop and clean up. If any one fails, others stop too.

use tokio_util::sync::CancellationToken;
use tokio::task::JoinSet;

pub async fn agent_run(token: CancellationToken) -> Result<Outcome, Error> {
    let mut set = JoinSet::new();

    {
        let t = token.child_token();
        set.spawn(async move { search(t).await });
    }
    {
        let t = token.child_token();
        set.spawn(async move { summarize(t).await });
    }
    {
        let t = token.child_token();
        set.spawn(async move { persist(t).await });
    }

    let mut outcomes = Vec::new();
    while let Some(res) = set.join_next().await {
        match res {
            Ok(Ok(v)) => outcomes.push(v),
            Ok(Err(e)) => {
                token.cancel();
                while let Some(_) = set.join_next().await {}
                return Err(e);
            }
            Err(join_err) => {
                token.cancel();
                while let Some(_) = set.join_next().await {}
                return Err(Error::Join(join_err));
            }
        }
    }
    Ok(assemble(outcomes))
}

Key points to articulate: child tokens are linked to the parent; cancelling the parent cascades; on any task failure we cancel and drain the rest.

P9 — Bounded channel from scratch

Prompt: Implement a bounded MPSC channel with send (awaits when full) and recv (awaits when empty). Don't use tokio::sync::mpsc directly.

Implementation sketch (real code is non-trivial): a Mutex<VecDeque<T>> + two Notify primitives (one for "not full", one for "not empty"). On send: lock, if full await not_full, push, notify not_empty, unlock. On recv: lock, if empty await not_empty, pop, notify not_full, unlock.

Why this is asked

Tests whether you understand the cooperative-cancellation pitfalls (you cannot hold a sync mutex across an await), the difference between Notify wakeups (one waiter) vs broadcast (all waiters), and the spurious-wakeup contract.

P10 — Graceful shutdown

Prompt: A Tokio HTTP server that, on SIGTERM, stops accepting new connections, finishes in-flight requests within 30 seconds, then exits.

use tokio::signal::unix::{signal, SignalKind};
use tokio_util::sync::CancellationToken;
use std::time::Duration;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let shutdown = CancellationToken::new();
    let server_shutdown = shutdown.clone();

    let server = tokio::spawn(async move {
        // axum's `with_graceful_shutdown` takes a future that resolves on shutdown
        let app = axum::Router::new().route("/health", axum::routing::get(|| async { "ok" }));
        let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
        axum::serve(listener, app)
            .with_graceful_shutdown(server_shutdown.cancelled_owned())
            .await
            .ok();
    });

    let mut term = signal(SignalKind::terminate())?;
    let mut intr = signal(SignalKind::interrupt())?;
    tokio::select! {
        _ = term.recv() => tracing::info!("SIGTERM"),
        _ = intr.recv() => tracing::info!("SIGINT"),
    }
    shutdown.cancel();

    // Give in-flight requests 30s; force-exit otherwise.
    match tokio::time::timeout(Duration::from_secs(30), server).await {
        Ok(_) => {}
        Err(_) => tracing::warn!("force-exit after timeout"),
    }
    Ok(())
}

Probes: SIGTERM vs SIGKILL; what about in-flight gRPC streams; how to drain Kafka consumers cleanly; what does k8s' preStop hook do (sleep so the load balancer notices the pod is going away before SIGTERM hits).

How to drill

  1. Set a 25-minute timer per problem. Code on paper or in a text editor without LSP.
  2. State the approach in 60 seconds before coding. Make the tradeoffs explicit.
  3. Once compiling (or "would compile" on paper), state two failure modes and one performance concern.
  4. Re-do the problem the next day. The structure is what you're memorizing, not the syntax.
The senior tell

Talking through the tradeoffs before coding wins more points than perfect syntax. "I'm choosing approach 1 because contention is low; if you tell me the bucket is hot I'd switch to sharded atomics" is what they want to hear.

✓ Check yourself

Answer from memory before opening each one.

  1. A token bucket is computed lazily on each call, with no background timer. What does it compute on access, and what do you change when one bucket becomes hot?

    Show answer

    It adds elapsed time times the rate to the tokens, caps at capacity and then tries to subtract the request. A single mutex per bucket becomes the bottleneck under contention, so move to a packed atomic with a compare-and-swap loop, or shard by key. Lazy refill means no timer thread per bucket. See P1 rate limiter.

  2. Fifty identical requests arrive while the first one is still computing, and the caller that started the compute is cancelled. What are your choices for the other 49?

    Show answer

    Cancel all of them, or let a waiting caller take over the work. The usual production choice is to run the compute as a detached task so it completes for everyone. Also decide that a failure goes to all waiters but is not cached, and bound the in-flight map. See P4 request dedup.

  3. A shared AtomicU64 counter slows down as core counts grow, though each increment is one instruction. Why, and what is the fix?

    Show answer

    Every increment bounces the cache line between cores. Use per-thread or per-CPU shards that are summed on read. Relaxed ordering suffices for a monotonic counter, but use acquire and release if the counter publishes other data. See P6 lock-free counter.

  4. Your retry combinator wraps a POST that creates an order. What does it need beyond backoff and jitter?

    Show answer

    An idempotency key passed in by the caller, so a repeated attempt does not create a second order, and a check that the error is retryable at all. A fleet-wide retry budget can also cap the extra load. See P7 retry combinator.

  5. An agent run spawns three sub-tasks and one fails. What does the parent do, and why does it wait for the others after cancelling?

    Show answer

    It cancels the shared token, which cascades to the child tokens, and then drains the join set before returning the error. Waiting lets each sub-task run its cleanup, so nothing is left running behind the caller. See P8 cancellation cascade.

  6. A pod receives SIGTERM while the load balancer still sends it traffic. Your graceful shutdown stops accepting connections immediately and clients see resets. What do you change?

    Show answer

    Keep serving briefly after the signal, using a preStop sleep, so the load balancer notices the pod is leaving before the server stops accepting. Then stop accepting, drain in-flight requests within the deadline and exit. See P10 graceful shutdown.

Takeaway

  • A token bucket refills lazily on access, with one mutex per bucket until it is hot, then atomics or sharding; a distributed limit needs an atomic Redis script and local caching.
  • Request coalescing keeps an in-flight map so concurrent callers share one compute, and the compute should finish even if the first caller leaves.
  • Sharded counters remove cache-line bouncing, and Relaxed ordering is enough for a counter that publishes nothing else.
  • A retry combinator retries only retryable errors with capped, jittered backoff, and needs an idempotency key when the operation has side effects.
  • Cancel with a token tree: cancel the parent on any failure and drain the join set so cleanup completes.
  • Graceful shutdown is stop accepting, drain within a deadline and exit, and in Kubernetes a preStop delay gives the load balancer time to notice.
  • Talk through the approach and trade-offs before coding, and drill the structure rather than the syntax.

With the coding rounds covered, the next chapter turns to the data that flows through the platform. → Data & pipelines

Further reading

  • Graceful Shutdown — Tokio documentation. Detecting shutdown, telling tasks to stop and waiting for them.
  • CancellationToken — docs.rs. Child tokens and cascading cancel used in P8 and P10.
  • JoinSet — docs.rs. Spawning a group of tasks and draining them.
  • Exponential Backoff and Jitter — AWS Architecture Blog, 2015. The comparison of full, equal and decorrelated jitter behind P7.
  • Pod Lifecycle — Kubernetes documentation. How termination, preStop hooks and grace periods interact.
  • Atomics — The Rustonomicon. Memory orderings, and when Relaxed is enough.