Rust Tokio Shared State
Shared state is data that every connection task can reach. In Mini-Redis, the shared state is one map of keys.
Share the map behind a lock
type Db = Arc<Mutex<HashMap<String, Bytes>>>;
let db: Db = Arc::new(Mutex::new(HashMap::new()));Every task gets a handle
let (socket, _) = listener.accept().await?;
let db = db.clone(); // a handle, not a copy of the map
tokio::spawn(async move { process(socket, db).await });Lock briefly, in a small scope
The guard is dropped at the end of the block โ before any await can happen.
{
let mut map = db.lock().unwrap();
map.insert("name".to_string(), Bytes::from("Ferris"));
} // guard dropped here
let value = db.lock().unwrap().get("name").unwrap().clone();
println!("shared value = {}", String::from_utf8_lossy(&value));Or hide the lock in a struct
If the lock only appears inside plain methods, no guard can ever cross an await.
struct Db2 {
map: Mutex<HashMap<String, String>>,
}
impl Db2 {
fn set(&self, key: String, value: String) { // not async, on purpose
self.map.lock().unwrap().insert(key, value);
}
}Or let one task own the state
Strategy 2: requests are messages, and each message carries its own reply channel.
enum Command {
Get { key: String, resp: oneshot::Sender<Option<String>> },
Set { key: String, value: String, resp: oneshot::Sender<()> },
}Shard when one lock gets hot
type ShardedDb = Arc<Vec<Mutex<HashMap<String, Bytes>>>>;
// a key's shard: db[hash(key) % db.len()]Example
// Cargo.toml: tokio = { version = "1", features = ["full"] }
use std::collections::HashMap;
use tokio::sync::{mpsc, oneshot};
enum Command {
Get { key: String, resp: oneshot::Sender<Option<String>> },
Set { key: String, value: String, resp: oneshot::Sender<()> },
}
#[tokio::main]
async fn main() {
let (tx, mut rx) = mpsc::channel::<Command>(8);
let manager = tokio::spawn(async move {
let mut db: HashMap<String, String> = HashMap::new();
while let Some(cmd) = rx.recv().await {
match cmd {
Command::Set { key, value, resp } => {
db.insert(key, value);
let _ = resp.send(());
}
Command::Get { key, resp } => {
let _ = resp.send(db.get(&key).cloned());
}
}
}
});
let (resp_tx, resp_rx) = oneshot::channel();
tx.send(Command::Set {
key: "name".to_string(),
value: "Ferris".to_string(),
resp: resp_tx,
})
.await
.unwrap();
resp_rx.await.unwrap();
println!("set done");
let (resp_tx, resp_rx) = oneshot::channel();
tx.send(Command::Get {
key: "name".to_string(),
resp: resp_tx,
})
.await
.unwrap();
println!("name = {}", resp_rx.await.unwrap().unwrap());
drop(tx);
manager.await.unwrap();
println!("manager task finished");
}Key points
- Use the std Mutex for short, non-async critical sections.
- The Tokio Mutex is for guards that truly must live across an .await.
- A std MutexGuard is not Send: hold it across an .await and spawn refuses the future โ scope the lock so the guard drops before the await
- Choose strategy 2 when working on the state involves awaiting (I/O) โ or when contention on one lock becomes the bottleneck
- Bytes values clone cheaply (a counter bump, not a copy) โ which is why the map stores Bytes, not Vec<u8>
- "Handle" is Tokio's word for a cloned access point to shared state.
- Cloning a Db or a Sender gives another handle.
- Shared state stays in one place, and every task reaches it through a handle.
In this example, the manager task owns the HashMap โ no Mutex anywhere, because only one task ever touches the map. Main asks for a Set by sending a Command message. The message carries a oneshot sender. The manager stores the pair. It answers through that sender. So main can print set done only after the store really happened. The Get works the same way and returns Ferris. Finally, main drops its channel sender. The manager's recv loop then sees the channel close and ends, and main confirms it: manager task finished. This is the tutorial's second strategy, the one it recommends whenever state work involves waiting.
Note
It never gives a copy of the state.
Try it Yourself
Edit the code, then press Run.
Running is simulated ยท Rust Playground
Exercise
Why does holding a std MutexGuard across an .await break tokio::spawn?