Thread-Safe Memoization in Rust with RwLock and OnceCell Sharding
use std::collections::HashMap;
use std::hash::Hash;
use std::sync::{Arc, RwLock};
pub struct Memoizer<K, V> {
store: RwLock<HashMap<K, Arc<V>>>,
}
impl<K, V> Memoizer<K, V>
where
K: Eq + Hash + Clone,
{
pub fn new() -> Self {
Memoizer {
store: RwLock::new(HashMap::new()),
}
}
pub fn get_or_compute<F>(&self, key: K, compute: F) -> Arc<V>
where
F: FnOnce() -> V,
{
if let Some(hit) = self.store.read().unwrap().get(&key) {
return Arc::clone(hit);
}
let mut guard = self.store.write().unwrap();
// Another thread may have filled this slot while we waited.
if let Some(hit) = guard.get(&key) {
return Arc::clone(hit);
}
let value = Arc::new(compute());
guard.insert(key, Arc::clone(&value));
value
}
pub fn len(&self) -> usize {
self.store.read().unwrap().len()
}
}use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::sync::Arc;
use crate::memoizer::Memoizer;
pub struct ShardedMemoizer<K, V> {
shards: Vec<Memoizer<K, V>>,
}
impl<K, V> ShardedMemoizer<K, V>
where
K: Eq + Hash + Clone,
{
pub fn with_shards(n: usize) -> Self {
let mut shards = Vec::with_capacity(n);
for _ in 0..n {
shards.push(Memoizer::new());
}
ShardedMemoizer { shards }
}
fn shard_for(&self, key: &K) -> &Memoizer<K, V> {
let mut hasher = DefaultHasher::new();
key.hash(&mut hasher);
let idx = (hasher.finish() as usize) % self.shards.len();
&self.shards[idx]
}
pub fn get_or_compute<F>(&self, key: K, compute: F) -> Arc<V>
where
F: FnOnce() -> V,
{
self.shard_for(&key).get_or_compute(key, compute)
}
pub fn entries(&self) -> usize {
self.shards.iter().map(|s| s.len()).sum()
}
}mod memoizer;
mod shard;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use shard::ShardedMemoizer;
fn slow_fib(n: u64) -> u64 {
thread::sleep(Duration::from_millis(5)); // pretend this is expensive
match n {
0 => 0,
1 => 1,
_ => slow_fib(n - 1).wrapping_add(slow_fib(n - 2)),
}
}
fn main() {
let cache: Arc<ShardedMemoizer<u64, u64>> = Arc::new(ShardedMemoizer::with_shards(8));
let handles: Vec<_> = (0..8)
.map(|t| {
let cache = Arc::clone(&cache);
thread::spawn(move || {
for n in 20..30 {
let key = n + (t % 3);
let value = cache.get_or_compute(key, || slow_fib(key));
println!("thread {t}: fib({key}) = {value}");
}
})
})
.collect();
for h in handles {
h.join().unwrap();
}
println!("distinct cached entries: {}", cache.entries());
}This snippet shows how an expensive, pure computation can be cached behind a thread-safe memoizer so that many threads can share results without recomputing and without corrupting shared state. The core idea is memoization: the first call for a given input runs the real work and stores the answer keyed by that input; every later call for the same input returns the stored value. In a single-threaded program a plain HashMap suffices, but under concurrency the map needs synchronization, and the design must avoid two classic failures — data races on the map and the thundering herd, where many threads compute the same missing value at once.
In memoizer.rs, Memoizer<K, V> wraps an RwLock<HashMap<K, Arc<V>>>. Values are stored as Arc<V> so the lock can be released the instant a clone of the pointer is taken; callers then hold a cheap reference-counted handle rather than pinning the whole map. get_or_compute follows a read-then-write pattern: it first takes a shared read lock and returns early on a hit, letting arbitrarily many readers proceed in parallel. Only on a miss does it acquire the exclusive write lock. Crucially it re-checks the entry after upgrading, because another thread may have inserted the value in the gap between dropping the read lock and taking the write lock. This double-checked insertion keeps the map correct even though the compute closure runs while the write lock is held.
Holding the write lock across the closure is a deliberate trade-off. It guarantees that a key is computed at most once (no duplicated expensive work), at the cost of serializing all misses through one lock. That is the right call when the computation dominates and duplicate work is the thing to avoid; it is the wrong call if the closure is slow and keys are highly independent, since one slow miss blocks unrelated misses.
To relieve that contention, shard.rs introduces ShardedMemoizer, which fans keys across N independent Memoizer instances chosen by hashing the key with DefaultHasher. Two different keys usually land in different shards and so contend on different locks, turning one global bottleneck into many small ones. shard_for uses the modulo of the hash to pick a shard, and get_or_compute simply delegates. This is the same technique that concurrent hash maps use internally.
In main.rs, a ShardedMemoizer<u64, u64> is shared via Arc across eight threads that all request an overlapping range of slow_fib values. Because the cache is populated cooperatively, each distinct input is computed once regardless of how many threads ask for it. The pitfalls worth noting: the compute closure must be pure and deterministic for memoization to be sound, V should be cheap to wrap in Arc, and this design caches forever — there is no eviction, so it fits bounded key spaces or process-lifetime lookup tables rather than unbounded user input.
Share this code
Here's the card — post it anywhere.
