63. Capstone: a Redis-compatible key-value server with persistence
This is the capstone of the C guide (§43)
and close to the Go guide's (§56),
written in Rust. It speaks enough of the RESP2 protocol for redis-cli to
work, and it combines most of the guide: Tokio tasks (31), framing with
limits (32, 53), enums and match (9), a std Mutex never held across
.await (31.5), expiry, atomic snapshots (41), and graceful shutdown (31.4).
// src/bin/kv.rs run: cargo run --bin kv -- serve then: redis-cli -p 6380 SET greeting hello
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
const MAX_ARGS: usize = 1024;
const MAX_BULK: usize = 8 * 1024 * 1024; // 8 MiB per argument
const MAX_LINE: usize = 64; // "*123\r\n" / "$456\r\n" headers are short
#[derive(Clone, Serialize, Deserialize)]
struct Entry { value: Vec<u8>, expires_at_ms: Option<u64> }
type Db = Arc<Mutex<HashMap<Vec<u8>, Entry>>>;
fn now_ms() -> u64 { SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64 }
enum Reply { Ok, Simple(&'static str), Bulk(Option<Vec<u8>>), Int(i64), Err(String) }
impl Reply {
fn encode(self) -> Vec<u8> {
match self {
Reply::Ok => b"+OK\r\n".to_vec(),
Reply::Simple(s) => format!("+{s}\r\n").into_bytes(),
Reply::Bulk(None) => b"$-1\r\n".to_vec(),
Reply::Bulk(Some(v)) => [format!("${}\r\n", v.len()).as_bytes(), &v, b"\r\n"].concat(),
Reply::Int(n) => format!(":{n}\r\n").into_bytes(),
Reply::Err(e) => format!("-ERR {e}\r\n").into_bytes(),
}
}
}
/// Reads one header line ("*3", "$5"), bounded so a client cannot send an endless line.
async fn read_header(r: &mut BufReader<TcpStream>, prefix: u8) -> io::Result<Option<usize>> {
let mut line = Vec::new();
let n = (&mut *r).take(MAX_LINE as u64).read_until(b'\n', &mut line).await?;
if n == 0 { return Ok(None); } // clean EOF between commands
let bad = || io::Error::new(io::ErrorKind::InvalidData, "protocol error");
if !line.ends_with(b"\r\n") || line.first() != Some(&prefix) { return Err(bad()); }
let num = std::str::from_utf8(&line[1..line.len() - 2]).map_err(|_| bad())?;
num.parse::<usize>().map(Some).map_err(|_| bad())
}
async fn read_command(r: &mut BufReader<TcpStream>) -> io::Result<Option<Vec<Vec<u8>>>> {
let Some(argc) = read_header(r, b'*').await? else { return Ok(None) };
if argc == 0 || argc > MAX_ARGS { return Err(io::Error::new(io::ErrorKind::InvalidData, "bad arg count")); }
let mut args = Vec::with_capacity(argc); // argc is bounded, so this is safe
for _ in 0..argc {
let len = read_header(r, b'$').await?.ok_or(io::ErrorKind::UnexpectedEof)?;
if len > MAX_BULK { return Err(io::Error::new(io::ErrorKind::InvalidData, "argument too large")); }
let mut buf = vec![0u8; len + 2];
r.read_exact(&mut buf).await?;
if !buf.ends_with(b"\r\n") { return Err(io::Error::new(io::ErrorKind::InvalidData, "missing CRLF")); }
buf.truncate(len);
args.push(buf);
}
Ok(Some(args))
}
fn execute(db: &Db, args: &[Vec<u8>]) -> Reply {
let cmd = String::from_utf8_lossy(&args[0]).to_ascii_uppercase();
let mut map = db.lock().unwrap(); // synchronous critical section
let live = |e: &Entry| e.expires_at_ms.is_none_or(|t| t > now_ms());
match (cmd.as_str(), args.len()) {
("PING", 1) => Reply::Simple("PONG"),
("PING", 2) | ("ECHO", 2) => Reply::Bulk(Some(args[1].clone())),
("GET", 2) => Reply::Bulk(map.get(&args[1]).filter(|e| live(e)).map(|e| e.value.clone())),
("SET", 3) | ("SET", 5) => {
let expires_at_ms = if args.len() == 5 {
let unit = String::from_utf8_lossy(&args[3]).to_ascii_uppercase();
let Some(n) = std::str::from_utf8(&args[4]).ok().and_then(|s| s.parse::<u64>().ok()).filter(|&n| n > 0) else {
return Reply::Err("invalid expire time".into());
};
match unit.as_str() {
"EX" => Some(now_ms().saturating_add(n.saturating_mul(1000))),
"PX" => Some(now_ms().saturating_add(n)),
_ => return Reply::Err("syntax error".into()),
}
} else { None };
map.insert(args[1].clone(), Entry { value: args[2].clone(), expires_at_ms });
Reply::Ok
}
("DEL", n) if n >= 2 => Reply::Int(args[1..].iter().filter(|k| map.remove(*k).is_some()).count() as i64),
("EXISTS", n) if n >= 2 => Reply::Int(args[1..].iter().filter(|k| map.get(*k).is_some_and(|e| live(e))).count() as i64),
("INCR", 2) => {
let current = map.get(&args[1]).filter(|e| live(e)).map(|e| e.value.clone());
let n = match current.as_deref().map(std::str::from_utf8) {
None => 0,
Some(Ok(s)) => match s.parse::<i64>() { Ok(n) => n, Err(_) => return Reply::Err("value is not an integer".into()) },
Some(Err(_)) => return Reply::Err("value is not an integer".into()),
};
let Some(next) = n.checked_add(1) else { return Reply::Err("increment would overflow".into()) };
map.insert(args[1].clone(), Entry { value: next.to_string().into_bytes(), expires_at_ms: None });
Reply::Int(next)
}
("DBSIZE", 1) => Reply::Int(map.values().filter(|e| live(e)).count() as i64),
("SAVE", 1) => {
let snapshot: Vec<(Vec<u8>, Entry)> = map.iter().filter(|(_, e)| live(e)).map(|(k, e)| (k.clone(), e.clone())).collect();
drop(map); // do not hold the lock during disk I/O
match serde_json::to_vec(&snapshot).map_err(io::Error::other).and_then(|b| atomic_write(Path::new(SNAPSHOT), &b)) {
Ok(()) => Reply::Ok,
Err(e) => Reply::Err(format!("save failed: {e}")),
}
}
_ => Reply::Err(format!("unknown command or wrong number of arguments for '{cmd}'")),
}
}
const SNAPSHOT: &str = "kv-snapshot.json";
fn atomic_write(path: &Path, data: &[u8]) -> io::Result<()> { // Ch 41, condensed
let tmp = path.with_extension("tmp");
{ let mut f = std::fs::File::create(&tmp)?; io::Write::write_all(&mut f, data)?; f.sync_all()?; }
std::fs::rename(&tmp, path)?;
std::fs::File::open(path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new(".")))?.sync_all()
}
fn load(path: &Path) -> HashMap<Vec<u8>, Entry> {
let Ok(bytes) = std::fs::read(path) else { return HashMap::new() };
serde_json::from_slice::<Vec<(Vec<u8>, Entry)>>(&bytes).map(|v| v.into_iter().collect()).unwrap_or_default()
}
async fn client(stream: TcpStream, db: Db) -> io::Result<()> {
let mut r = BufReader::new(stream);
while let Some(args) = read_command(&mut r).await? {
let reply = execute(&db, &args).encode(); // lock taken and released inside
r.get_mut().write_all(&reply).await?;
}
Ok(())
}
async fn serve(listener: TcpListener, db: Db) {
let sweeper_db = Arc::clone(&db);
tokio::spawn(async move { // active expiry, once per second
let mut tick = tokio::time::interval(Duration::from_secs(1));
loop { tick.tick().await; let t = now_ms(); sweeper_db.lock().unwrap().retain(|_, e| e.expires_at_ms.is_none_or(|x| x > t)); }
});
loop {
let Ok((stream, peer)) = listener.accept().await else { continue };
let db = Arc::clone(&db);
tokio::spawn(async move {
if let Err(e) = client(stream, db).await {
if e.kind() != io::ErrorKind::UnexpectedEof { eprintln!("{peer}: {e}"); }
}
});
}
}
async fn demo(addr: std::net::SocketAddr) -> io::Result<()> {
let mut s = TcpStream::connect(addr).await?;
let cmd = |parts: &[&str]| -> Vec<u8> {
let mut out = format!("*{}\r\n", parts.len()).into_bytes();
for p in parts { out.extend(format!("${}\r\n{p}\r\n", p.len()).into_bytes()); }
out
};
for parts in [&["PING"][..], &["SET", "greeting", "hello"], &["GET", "greeting"], &["INCR", "hits"], &["INCR", "hits"],
&["SET", "session", "abc", "PX", "100"], &["EXISTS", "greeting", "session", "nope"], &["INCR", "greeting"], &["SAVE"]] {
s.write_all(&cmd(parts)).await?;
let mut buf = vec![0u8; 512];
let n = s.read(&mut buf).await?;
println!("{:<40} -> {:?}", parts.join(" "), String::from_utf8_lossy(&buf[..n]));
}
tokio::time::sleep(Duration::from_millis(150)).await; // let `session` expire
s.write_all(&cmd(&["GET", "session"])).await?;
let mut buf = vec![0u8; 64];
let n = s.read(&mut buf).await?;
println!("{:<40} -> {:?}", "GET session (after 150ms)", String::from_utf8_lossy(&buf[..n]));
s.write_all(b"*99999999\r\n").await?; // hostile header: server drops the connection
println!("{:<40} -> closed: {}", "*99999999", s.read(&mut buf).await? == 0);
Ok(())
}
#[tokio::main]
async fn main() -> io::Result<()> {
let db: Db = Arc::new(Mutex::new(load(Path::new(SNAPSHOT))));
if std::env::args().nth(1).as_deref() == Some("serve") {
let listener = TcpListener::bind("127.0.0.1:6380").await?;
println!("kv listening on {}, {} keys loaded", listener.local_addr()?, db.lock().unwrap().len());
tokio::select! {
_ = serve(listener, Arc::clone(&db)) => {}
_ = tokio::signal::ctrl_c() => println!("shutting down"),
}
return Ok(());
}
let listener = TcpListener::bind("127.0.0.1:0").await?;
let addr = listener.local_addr()?;
tokio::spawn(serve(listener, db));
demo(addr).await?;
std::fs::remove_file(SNAPSHOT)
}
PING -> "+PONG\r\n"
SET greeting hello -> "+OK\r\n"
GET greeting -> "$5\r\nhello\r\n"
INCR hits -> ":1\r\n"
INCR hits -> ":2\r\n"
SET session abc PX 100 -> "+OK\r\n"
EXISTS greeting session nope -> ":2\r\n"
INCR greeting -> "-ERR value is not an integer\r\n"
SAVE -> "+OK\r\n"
GET session (after 150ms) -> "$-1\r\n"
127.0.0.1:55288: bad arg count <- the server's stderr
*99999999 -> closed: true
Compare it with the C version. The C capstone needs a hand-written hash
table, manual buffer management with a growing input buffer, careful
realloc error paths, and select() bookkeeping. Here the memory-safety
work is gone. The remaining code is protocol logic, limits, and
concurrency design, which is where the remaining bugs will be.
Stretch goals:
- Append-only log (Redis AOF, Bitcask in the Go guide): write each
mutating command to a file and replay it on startup. Compare
fsyncper command with once per second, the durability vs throughput trade-off from OS Ch 77. - Sharding: replace
Mutex<HashMap>with 16 shards selected by key hash and measure the throughput withredis-benchmark -p 6380 -t set,get -c 50. - Pipelining: the server already supports it, because it reads
commands in a loop. Batch the replies into one
writeper read buffer and measure again. - TLS: wrap the listener with the Ch 49 acceptor.
redis-cli --tlswill connect.