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::Notifyor 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" }));
axum::Server::bind(&"0.0.0.0:8080".parse().unwrap())
.serve(app.into_make_service())
.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
- Set a 25-minute timer per problem. Code on paper or in a text editor without LSP.
- State the approach in 60 seconds before coding. Make the tradeoffs explicit.
- Once compiling (or "would compile" on paper), state two failure modes and one performance concern.
- Re-do the problem the next day. The structure is what you're memorizing, not the syntax.
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.