codeloopback
Guide contents

Lesson 12 of 12

Key-Value Server

We build kvd, a small in-memory key-value server with a Redis-like text protocol, one piece at a time. Every piece is tested before we move on, and the last example is the complete program: a server and several clients talking over real TCP.

What it does:

Text
SET user:1 ana          -> OK
GET user:1              -> ana
SET session xyz PX 50   -> OK          (expires after 50 ms)
INCR visits             -> 1
DEL user:1              -> 1
KEYS                    -> session,visits
STATS                   -> commands=6 connections=1 keys=2 expired=0

And the course topics each part uses:

Part Topics
protocol: parsing commands enums, pattern matching on slices, FromStr, thiserror, unit tests
store: data with expiry structs, HashMap, Option, traits for an injected clock, Result
responses Display, conversions
shared database Arc, Mutex, atomics
server Tokio, TcpListener, one task per connection, buffered line I/O
background work interval, select!, watch for shutdown
persistence serde, serde_json, tokio::fs, anyhow
testing #[test], #[tokio::test], fakes

The building blocks are Tokio, serde, and cargo test, and the whole server is a few hundred lines.

1. Project layout

A real project would split into modules like this; the examples below keep each part in a single file so they can run on their own.

Text
kvd/
├── Cargo.toml
└── src/
    ├── main.rs        start the server, handle Ctrl-C
    ├── lib.rs         pub mod protocol; pub mod store; pub mod server;
    ├── protocol.rs    Command, ParseError, Response
    ├── store.rs       Store, Clock, StoreError
    └── server.rs      Db, Stats, handle_connection, run
TOMLCargo.toml
[package]
name = "kvd"
version = "0.1.0"
edition = "2024"

[dependencies]
tokio = { version = "1", features = ["full"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
thiserror = "2"
anyhow = "1"

2. The command type

Start with the domain: every request is one of a fixed set of commands, so it's an enum. Each variant holds exactly the data it needs.

Rust
#[derive(Debug, Clone, PartialEq)]
pub enum Command {
    Get(String),
    Set { key: String, value: String, ttl_ms: Option<u64> },
    Del(String),
    Incr(String),
    Keys,
    Stats,
}

fn main() {
    let cmds = vec![
        Command::Set { key: "user:1".into(), value: "ana".into(), ttl_ms: None },
        Command::Get("user:1".into()),
        Command::Set { key: "session".into(), value: "xyz".into(), ttl_ms: Some(50) },
        Command::Keys,
    ];
    for c in &cmds {
        println!("{c:?}");
    }
}
Output
Set { key: "user:1", value: "ana", ttl_ms: None }
Get("user:1")
Set { key: "session", value: "xyz", ttl_ms: Some(50) }
Keys

3. Parsing commands with FromStr

The wire protocol is one command per line, words separated by spaces. Implementing FromStr lets the server write line.parse::<Command>(). Matching on (name, args.as_slice()) handles every command and its argument count in one match (Pattern Matching), and thiserror gives the error type its messages (thiserror & anyhow).

Rust
use std::str::FromStr;
use thiserror::Error;

#[derive(Debug, Clone, PartialEq)]
pub enum Command {
    Get(String),
    Set { key: String, value: String, ttl_ms: Option<u64> },
    Del(String),
    Incr(String),
    Keys,
    Stats,
}

#[derive(Debug, Error, PartialEq)]
pub enum ParseError {
    #[error("empty command")]
    Empty,
    #[error("unknown command '{0}'")]
    Unknown(String),
    #[error("wrong arguments, usage: {0}")]
    Usage(&'static str),
    #[error("invalid expiry '{0}'")]
    BadExpiry(String),
}

fn expiry(unit: &str, amount: &str) -> Result<u64, ParseError> {
    let n: u64 = amount.parse().map_err(|_| ParseError::BadExpiry(amount.to_string()))?;
    match unit.to_ascii_uppercase().as_str() {
        "PX" => Ok(n),
        "EX" => Ok(n * 1000),
        _ => Err(ParseError::BadExpiry(unit.to_string())),
    }
}

impl FromStr for Command {
    type Err = ParseError;

    fn from_str(line: &str) -> Result<Command, ParseError> {
        let mut words = line.split_whitespace();
        let name = words.next().ok_or(ParseError::Empty)?.to_ascii_uppercase();
        let args: Vec<&str> = words.collect();

        match (name.as_str(), args.as_slice()) {
            ("GET", [key]) => Ok(Command::Get(key.to_string())),
            ("SET", [key, value]) => Ok(Command::Set { key: key.to_string(), value: value.to_string(), ttl_ms: None }),
            ("SET", [key, value, unit, amount]) => {
                Ok(Command::Set { key: key.to_string(), value: value.to_string(), ttl_ms: Some(expiry(unit, amount)?) })
            }
            ("DEL", [key]) => Ok(Command::Del(key.to_string())),
            ("INCR", [key]) => Ok(Command::Incr(key.to_string())),
            ("KEYS", []) => Ok(Command::Keys),
            ("STATS", []) => Ok(Command::Stats),
            ("GET", _) => Err(ParseError::Usage("GET key")),
            ("SET", _) => Err(ParseError::Usage("SET key value [EX seconds | PX milliseconds]")),
            ("DEL", _) => Err(ParseError::Usage("DEL key")),
            ("INCR", _) => Err(ParseError::Usage("INCR key")),
            ("KEYS", _) | ("STATS", _) => Err(ParseError::Usage("KEYS / STATS take no arguments")),
            (other, _) => Err(ParseError::Unknown(other.to_lowercase())),
        }
    }
}

fn main() {
    for line in ["SET user:1 ana", "get user:1", "SET s xyz PX 50", "SET s xyz EX 2", "SET s", "SET s x EX soon", "FLY away", "   "] {
        match line.parse::<Command>() {
            Ok(cmd) => println!("{line:<18} => {cmd:?}"),
            Err(e) => println!("{line:<18} => error: {e}"),
        }
    }
}
Output
SET user:1 ana     => Set { key: "user:1", value: "ana", ttl_ms: None }
get user:1         => Get("user:1")
SET s xyz PX 50    => Set { key: "s", value: "xyz", ttl_ms: Some(50) }
SET s xyz EX 2     => Set { key: "s", value: "xyz", ttl_ms: Some(2000) }
SET s              => error: wrong arguments, usage: SET key value [EX seconds | PX milliseconds]
SET s x EX soon    => error: invalid expiry 'soon'
FLY away           => error: unknown command 'fly'
                   => error: empty command

4. Testing the parser

Parsing is pure logic, so it's easy to test thoroughly. A table of cases keeps the tests compact (Testing).

Rustcargo test
use std::str::FromStr;
use thiserror::Error;

#[derive(Debug, Clone, PartialEq)]
pub enum Command {
    Get(String),
    Set { key: String, value: String, ttl_ms: Option<u64> },
    Del(String),
    Incr(String),
    Keys,
    Stats,
}

#[derive(Debug, Error, PartialEq)]
pub enum ParseError {
    #[error("empty command")]
    Empty,
    #[error("unknown command '{0}'")]
    Unknown(String),
    #[error("wrong arguments, usage: {0}")]
    Usage(&'static str),
    #[error("invalid expiry '{0}'")]
    BadExpiry(String),
}

fn expiry(unit: &str, amount: &str) -> Result<u64, ParseError> {
    let n: u64 = amount.parse().map_err(|_| ParseError::BadExpiry(amount.to_string()))?;
    match unit.to_ascii_uppercase().as_str() {
        "PX" => Ok(n),
        "EX" => Ok(n * 1000),
        _ => Err(ParseError::BadExpiry(unit.to_string())),
    }
}

impl FromStr for Command {
    type Err = ParseError;

    fn from_str(line: &str) -> Result<Command, ParseError> {
        let mut words = line.split_whitespace();
        let name = words.next().ok_or(ParseError::Empty)?.to_ascii_uppercase();
        let args: Vec<&str> = words.collect();
        match (name.as_str(), args.as_slice()) {
            ("GET", [key]) => Ok(Command::Get(key.to_string())),
            ("SET", [key, value]) => Ok(Command::Set { key: key.to_string(), value: value.to_string(), ttl_ms: None }),
            ("SET", [key, value, unit, amount]) => {
                Ok(Command::Set { key: key.to_string(), value: value.to_string(), ttl_ms: Some(expiry(unit, amount)?) })
            }
            ("DEL", [key]) => Ok(Command::Del(key.to_string())),
            ("INCR", [key]) => Ok(Command::Incr(key.to_string())),
            ("KEYS", []) => Ok(Command::Keys),
            ("STATS", []) => Ok(Command::Stats),
            ("GET", _) => Err(ParseError::Usage("GET key")),
            ("SET", _) => Err(ParseError::Usage("SET key value [EX seconds | PX milliseconds]")),
            ("DEL", _) => Err(ParseError::Usage("DEL key")),
            ("INCR", _) => Err(ParseError::Usage("INCR key")),
            ("KEYS", _) | ("STATS", _) => Err(ParseError::Usage("KEYS / STATS take no arguments")),
            (other, _) => Err(ParseError::Unknown(other.to_lowercase())),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_valid_commands() {
        let cases = [
            ("GET a", Command::Get("a".into())),
            ("get a", Command::Get("a".into())),
            ("SET k v", Command::Set { key: "k".into(), value: "v".into(), ttl_ms: None }),
            ("SET k v px 250", Command::Set { key: "k".into(), value: "v".into(), ttl_ms: Some(250) }),
            ("SET k v EX 3", Command::Set { key: "k".into(), value: "v".into(), ttl_ms: Some(3000) }),
            ("  INCR   hits ", Command::Incr("hits".into())),
            ("KEYS", Command::Keys),
        ];
        for (line, expected) in cases {
            assert_eq!(line.parse::<Command>(), Ok(expected), "line {line:?}");
        }
    }

    #[test]
    fn rejects_invalid_commands() {
        assert_eq!("".parse::<Command>(), Err(ParseError::Empty));
        assert_eq!("PING".parse::<Command>(), Err(ParseError::Unknown("ping".into())));
        assert!(matches!("GET".parse::<Command>(), Err(ParseError::Usage(_))));
        assert!(matches!("KEYS x".parse::<Command>(), Err(ParseError::Usage(_))));
        assert_eq!("SET k v EX never".parse::<Command>(), Err(ParseError::BadExpiry("never".into())));
        assert_eq!("SET k v IN 5".parse::<Command>(), Err(ParseError::BadExpiry("IN".into())));
    }

    #[test]
    fn error_messages_are_helpful() {
        let e = "SET onlykey".parse::<Command>().unwrap_err();
        assert_eq!(e.to_string(), "wrong arguments, usage: SET key value [EX seconds | PX milliseconds]");
    }
}
Output
running 3 tests
test tests::error_messages_are_helpful ... ok
test tests::parses_valid_commands ... ok
test tests::rejects_invalid_commands ... ok

test result: ok. 3 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

5. The store, with an injected clock

The store is a HashMap of entries, each with an optional expiry time. Time comes from a Clock trait so tests can control it (Testing, example 18). Expired entries are treated as missing on read and removed by purge_expired, which a background task will call later.

Rust
use std::collections::HashMap;

pub trait Clock {
    fn now_ms(&self) -> u64;
}

#[derive(Debug, Clone)]
pub struct Entry {
    pub value: String,
    pub expires_at: Option<u64>,
}

pub struct Store<C: Clock> {
    entries: HashMap<String, Entry>,
    clock: C,
}

impl<C: Clock> Store<C> {
    pub fn new(clock: C) -> Self {
        Store { entries: HashMap::new(), clock }
    }

    fn is_live(&self, e: &Entry) -> bool {
        e.expires_at.is_none_or(|t| t > self.clock.now_ms())
    }

    pub fn get(&self, key: &str) -> Option<&str> {
        self.entries.get(key).filter(|e| self.is_live(e)).map(|e| e.value.as_str())
    }

    pub fn set(&mut self, key: &str, value: &str, ttl_ms: Option<u64>) {
        let expires_at = ttl_ms.map(|ttl| self.clock.now_ms() + ttl);
        self.entries.insert(key.to_string(), Entry { value: value.to_string(), expires_at });
    }

    pub fn purge_expired(&mut self) -> usize {
        let now = self.clock.now_ms();
        let before = self.entries.len();
        self.entries.retain(|_, e| e.expires_at.is_none_or(|t| t > now));
        before - self.entries.len()
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }
}

// A clock we can move by hand.
struct ManualClock(std::cell::Cell<u64>);

impl Clock for &ManualClock {
    fn now_ms(&self) -> u64 {
        self.0.get()
    }
}

fn main() {
    let clock = ManualClock(std::cell::Cell::new(1_000));
    let mut store = Store::new(&clock);

    store.set("user:1", "ana", None);
    store.set("session", "xyz", Some(50));
    println!("at 1000: {:?} {:?}", store.get("user:1"), store.get("session"));

    clock.0.set(1_049);
    println!("at 1049: {:?}", store.get("session"));
    clock.0.set(1_050);
    println!("at 1050: {:?} (still stored: {})", store.get("session"), store.len());

    println!("purged {} expired entries, {} left", store.purge_expired(), store.len());
}
Output
at 1000: Some("ana") Some("xyz")
at 1049: Some("xyz")
at 1050: None (still stored: 2)
purged 1 expired entries, 1 left

Note impl Clock for &ManualClock: the store owns a reference to the clock, so the test keeps its own reference and can move time forward while the store exists.

6. The rest of the store: DEL, INCR, KEYS, and errors

INCR can fail when the value isn't a number, so it returns a Result with a small error enum. keys() returns only live keys, sorted, so output is predictable.

Rust
use std::collections::HashMap;
use thiserror::Error;

pub trait Clock {
    fn now_ms(&self) -> u64;
}

pub struct FixedClock(pub u64);
impl Clock for FixedClock {
    fn now_ms(&self) -> u64 {
        self.0
    }
}

#[derive(Debug, Clone)]
pub struct Entry {
    pub value: String,
    pub expires_at: Option<u64>,
}

#[derive(Debug, Error, PartialEq)]
pub enum StoreError {
    #[error("value of '{0}' is not an integer")]
    NotAnInteger(String),
    #[error("increment would overflow")]
    Overflow,
}

pub struct Store<C: Clock> {
    entries: HashMap<String, Entry>,
    clock: C,
}

impl<C: Clock> Store<C> {
    pub fn new(clock: C) -> Self {
        Store { entries: HashMap::new(), clock }
    }

    fn live(&self, key: &str) -> Option<&Entry> {
        let now = self.clock.now_ms();
        self.entries.get(key).filter(|e| e.expires_at.is_none_or(|t| t > now))
    }

    pub fn get(&self, key: &str) -> Option<&str> {
        self.live(key).map(|e| e.value.as_str())
    }

    pub fn set(&mut self, key: &str, value: &str, ttl_ms: Option<u64>) {
        let expires_at = ttl_ms.map(|t| self.clock.now_ms() + t);
        self.entries.insert(key.to_string(), Entry { value: value.to_string(), expires_at });
    }

    pub fn del(&mut self, key: &str) -> bool {
        let existed = self.live(key).is_some();
        self.entries.remove(key);
        existed
    }

    pub fn incr(&mut self, key: &str) -> Result<i64, StoreError> {
        let (current, expires_at) = match self.live(key) {
            Some(e) => (e.value.parse::<i64>().map_err(|_| StoreError::NotAnInteger(key.to_string()))?, e.expires_at),
            None => (0, None),
        };
        let next = current.checked_add(1).ok_or(StoreError::Overflow)?;
        self.entries.insert(key.to_string(), Entry { value: next.to_string(), expires_at });
        Ok(next)
    }

    pub fn keys(&self) -> Vec<&str> {
        let mut keys: Vec<&str> = self.entries.keys().filter(|k| self.live(k).is_some()).map(String::as_str).collect();
        keys.sort_unstable();
        keys
    }
}

fn main() {
    let mut store = Store::new(FixedClock(0));
    store.set("name", "ana", None);
    println!("{:?} {:?} {:?}", store.incr("visits"), store.incr("visits"), store.incr("name"));
    store.set("big", &i64::MAX.to_string(), None);
    println!("{:?}", store.incr("big"));
    println!("{:?}", store.keys());
    println!("del name: {}, del again: {}", store.del("name"), store.del("name"));
    println!("{:?}", store.keys());
}
Output
Ok(1) Ok(2) Err(NotAnInteger("name"))
Err(Overflow)
["big", "name", "visits"]
del name: true, del again: false
["big", "visits"]

7. Responses and the wire format

Each command produces a Response. Its Display implementation is the protocol: one line per response, which clients can read with lines().

Rust
use std::fmt;

#[derive(Debug, Clone, PartialEq)]
pub enum Response {
    Ok,
    Value(String),
    Nil,
    Int(i64),
    List(Vec<String>),
    Error(String),
}

impl fmt::Display for Response {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Response::Ok => write!(f, "OK"),
            Response::Value(v) => write!(f, "{v}"),
            Response::Nil => write!(f, "(nil)"),
            Response::Int(n) => write!(f, "{n}"),
            Response::List(items) if items.is_empty() => write!(f, "(empty)"),
            Response::List(items) => write!(f, "{}", items.join(",")),
            Response::Error(msg) => write!(f, "ERR {msg}"),
        }
    }
}

// Any error type with Display can become an error response.
impl<E: std::error::Error> From<E> for Response {
    fn from(e: E) -> Response {
        Response::Error(e.to_string())
    }
}

fn main() {
    let io_err = std::io::Error::other("disk full");
    let all = [
        Response::Ok,
        Response::Value("ana".into()),
        Response::Nil,
        Response::Int(42),
        Response::List(vec!["a".into(), "b".into()]),
        Response::List(vec![]),
        Response::from(io_err),
    ];
    for r in all {
        println!("{r}");
    }
}
Output
OK
ana
(nil)
42
a,b
(empty)
ERR disk full

The blanket impl<E: Error> From<E> for Response means any error can be turned into a response with .into() or ?-style conversions (Conversion Traits).

8. Executing commands

execute connects the pieces: it takes a parsed Command, applies it to the store, and returns a Response. It's a pure function over the store, so it's easy to test without any networking.

Rust
use std::collections::BTreeMap;

// Minimal stand-ins for the store and response from the previous steps.
#[derive(Debug)]
enum Command {
    Get(String),
    Set { key: String, value: String },
    Del(String),
    Incr(String),
    Keys,
}

#[derive(Debug, PartialEq)]
enum Response {
    Ok,
    Value(String),
    Nil,
    Int(i64),
    List(Vec<String>),
    Error(String),
}

#[derive(Default)]
struct Store(BTreeMap<String, String>);

fn execute(store: &mut Store, cmd: Command) -> Response {
    match cmd {
        Command::Get(k) => store.0.get(&k).map_or(Response::Nil, |v| Response::Value(v.clone())),
        Command::Set { key, value } => {
            store.0.insert(key, value);
            Response::Ok
        }
        Command::Del(k) => Response::Int(store.0.remove(&k).is_some() as i64),
        Command::Incr(k) => {
            let current = store.0.get(&k).map_or(Ok(0), |v| v.parse::<i64>());
            match current {
                Ok(n) => {
                    store.0.insert(k, (n + 1).to_string());
                    Response::Int(n + 1)
                }
                Err(_) => Response::Error(format!("value of '{k}' is not an integer")),
            }
        }
        Command::Keys => Response::List(store.0.keys().cloned().collect()),
    }
}

fn main() {
    let mut store = Store::default();
    let script = vec![
        Command::Set { key: "a".into(), value: "1".into() },
        Command::Incr("a".into()),
        Command::Incr("new".into()),
        Command::Get("a".into()),
        Command::Del("a".into()),
        Command::Get("a".into()),
        Command::Set { key: "t".into(), value: "text".into() },
        Command::Incr("t".into()),
        Command::Keys,
    ];
    for cmd in script {
        let label = format!("{cmd:?}");
        println!("{label:<40} -> {:?}", execute(&mut store, cmd));
    }
}
Output
Set { key: "a", value: "1" }             -> Ok
Incr("a")                                -> Int(2)
Incr("new")                              -> Int(1)
Get("a")                                 -> Value("2")
Del("a")                                 -> Int(1)
Get("a")                                 -> Nil
Set { key: "t", value: "text" }          -> Ok
Incr("t")                                -> Error("value of 't' is not an integer")
Keys                                     -> List(["new", "t"])

9. Sharing the database between tasks

Every connection runs in its own Tokio task, and all of them use one store. The Db handle wraps the store in Arc<Mutex<...>> and keeps statistics in atomics, which don't need the lock (Shared State, Atomics). Cloning a Db just clones the Arcs.

The store's mutex is a std::sync::Mutex, not Tokio's, because the lock is never held across an .await: each command locks, runs, and unlocks immediately (Async Channels & Locks).

Rust
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::sync::{Arc, Mutex};

#[derive(Default)]
pub struct Stats {
    pub commands: AtomicU64,
    pub connections: AtomicU64,
}

#[derive(Clone, Default)]
pub struct Db {
    store: Arc<Mutex<HashMap<String, i64>>>,
    pub stats: Arc<Stats>,
}

impl Db {
    pub fn incr(&self, key: &str) -> i64 {
        self.stats.commands.fetch_add(1, Relaxed);
        let mut store = self.store.lock().unwrap(); // held only for this block, never across .await
        let v = store.entry(key.to_string()).or_insert(0);
        *v += 1;
        *v
    }
}

#[tokio::main]
async fn main() {
    let db = Db::default();

    let mut tasks = Vec::new();
    for _ in 0..8 {
        let db = db.clone(); // a cheap handle clone per task
        tasks.push(tokio::spawn(async move {
            db.stats.connections.fetch_add(1, Relaxed);
            for _ in 0..125 {
                db.incr("hits");
                tokio::task::yield_now().await; // let other tasks interleave
            }
        }));
    }
    for t in tasks {
        t.await.unwrap();
    }

    println!("hits = {}", db.incr("hits") - 1);
    println!("commands = {}, connections = {}", db.stats.commands.load(Relaxed), db.stats.connections.load(Relaxed));
}
Output
hits = 1000
commands = 1001, connections = 8

10. The TCP server: one task per connection

Now the network. The server accepts connections in a loop and spawns a task for each, like one virtual thread per request. Each task reads lines, parses and executes commands, and writes back one response line (Async Fundamentals). A simple echo-style handler shows the shape before we plug in the real logic.

Rust
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};

async fn handle_connection(stream: TcpStream) -> std::io::Result<()> {
    let (reader, mut writer) = stream.into_split();
    let mut lines = BufReader::new(reader).lines();
    while let Some(line) = lines.next_line().await? {
        let reply = match line.trim() {
            "PING" => String::from("PONG"),
            other => format!("ERR unknown command '{other}'"),
        };
        writer.write_all(format!("{reply}\n").as_bytes()).await?;
    }
    Ok(()) // client closed the connection
}

async fn request(stream: &mut BufReader<TcpStream>, line: &str) -> std::io::Result<String> {
    stream.get_mut().write_all(format!("{line}\n").as_bytes()).await?;
    let mut reply = String::new();
    stream.read_line(&mut reply).await?;
    Ok(reply.trim_end().to_string())
}

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:0").await?; // port 0: let the OS pick
    let addr = listener.local_addr()?;
    tokio::spawn(async move {
        loop {
            let (stream, _) = listener.accept().await.unwrap();
            tokio::spawn(handle_connection(stream));
        }
    });

    let mut client = BufReader::new(TcpStream::connect(addr).await?);
    for cmd in ["PING", "HELLO"] {
        println!("{cmd} -> {}", request(&mut client, cmd).await?);
    }
    Ok(())
}
Output
PING -> PONG
HELLO -> ERR unknown command 'HELLO'

11. A background expiry task with graceful shutdown

Expired keys are removed lazily on read, but keys nobody reads again would stay forever. A background task purges them periodically. It runs a select! loop over a timer tick and a watch shutdown signal, so it stops cleanly when the server does (Tokio Patterns).

Rust
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::sync::watch;

type Store = Arc<Mutex<HashMap<String, Instant>>>; // key -> expiry time

async fn expiry_task(store: Store, mut shutdown: watch::Receiver<bool>) -> usize {
    let mut tick = tokio::time::interval(Duration::from_millis(20));
    let mut purged = 0;
    loop {
        tokio::select! {
            _ = tick.tick() => {
                let now = Instant::now();
                let mut map = store.lock().unwrap();
                let before = map.len();
                map.retain(|_, expires| *expires > now);
                purged += before - map.len();
            }
            _ = shutdown.changed() => return purged,
        }
    }
}

#[tokio::main]
async fn main() {
    let store: Store = Arc::default();
    let (shutdown_tx, shutdown_rx) = watch::channel(false);
    let task = tokio::spawn(expiry_task(Arc::clone(&store), shutdown_rx));

    {
        let mut map = store.lock().unwrap();
        let now = Instant::now();
        map.insert("short".into(), now + Duration::from_millis(30));
        map.insert("shorter".into(), now + Duration::from_millis(10));
        map.insert("long".into(), now + Duration::from_secs(60));
    }

    tokio::time::sleep(Duration::from_millis(150)).await;
    let remaining: Vec<String> = store.lock().unwrap().keys().cloned().collect();
    shutdown_tx.send(true).unwrap();

    println!("purged by the background task: {}", task.await.unwrap());
    println!("remaining: {remaining:?}");
}
Output
purged by the background task: 2
remaining: ["long"]

12. Snapshots with serde

To survive restarts, the server saves a snapshot on shutdown and loads it on startup. Deriving Serialize and Deserialize on the entry type is all serde needs; anyhow adds context to I/O and parsing errors (Crates & Workspaces).

Rust
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::Path;

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct Entry {
    value: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    expires_at: Option<u64>,
}

async fn save(path: &Path, entries: &BTreeMap<String, Entry>) -> Result<usize> {
    let json = serde_json::to_string_pretty(entries)?;
    tokio::fs::write(path, &json).await.with_context(|| format!("writing snapshot {}", path.display()))?;
    Ok(entries.len())
}

async fn load(path: &Path) -> Result<BTreeMap<String, Entry>> {
    let text = tokio::fs::read_to_string(path).await.with_context(|| format!("reading snapshot {}", path.display()))?;
    serde_json::from_str(&text).context("snapshot is not valid JSON")
}

#[tokio::main]
async fn main() -> Result<()> {
    let path = std::env::temp_dir().join(format!("kvd-snapshot-{}.json", std::process::id()));

    let mut entries = BTreeMap::new();
    entries.insert("user:1".to_string(), Entry { value: "ana".into(), expires_at: None });
    entries.insert("session".to_string(), Entry { value: "xyz".into(), expires_at: Some(1_700_000_000_000) });

    println!("saved {} entries", save(&path, &entries).await?);
    println!("{}", tokio::fs::read_to_string(&path).await?);
    let restored = load(&path).await?;
    println!("round trip ok: {}", restored == entries);
    tokio::fs::remove_file(&path).await?;

    match load(Path::new("/nonexistent/kvd.json")).await {
        Ok(_) => println!("unexpected"),
        Err(e) => println!("error: {e:#}"),
    }
    Ok(())
}
Output
saved 2 entries
{
  "session": {
    "value": "xyz",
    "expires_at": 1700000000000
  },
  "user:1": {
    "value": "ana"
  }
}
round trip ok: true
error: reading snapshot /nonexistent/kvd.json: No such file or directory (os error 2)

{e:#} prints the whole chain of contexts, from "reading snapshot ..." down to the OS error.

13. The complete server

Everything together in one file. Compared with the steps above, the only new pieces are the glue: Db::execute locks the store and runs a command, handle_connection uses it per line and counts statistics, and run owns the accept loop, the expiry task, and shutdown. main starts the server and then plays the role of several clients.

Rust
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::fmt;
use std::str::FromStr;
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use thiserror::Error;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::watch;
use tokio::task::JoinSet;

// ---------- protocol ----------

#[derive(Debug, Clone, PartialEq)]
pub enum Command {
    Get(String),
    Set { key: String, value: String, ttl_ms: Option<u64> },
    Del(String),
    Incr(String),
    Keys,
    Stats,
}

#[derive(Debug, Error, PartialEq)]
pub enum ParseError {
    #[error("empty command")]
    Empty,
    #[error("unknown command '{0}'")]
    Unknown(String),
    #[error("wrong arguments, usage: {0}")]
    Usage(&'static str),
    #[error("invalid expiry '{0}'")]
    BadExpiry(String),
}

fn expiry(unit: &str, amount: &str) -> Result<u64, ParseError> {
    let n: u64 = amount.parse().map_err(|_| ParseError::BadExpiry(amount.to_string()))?;
    match unit.to_ascii_uppercase().as_str() {
        "PX" => Ok(n),
        "EX" => Ok(n * 1000),
        _ => Err(ParseError::BadExpiry(unit.to_string())),
    }
}

impl FromStr for Command {
    type Err = ParseError;
    fn from_str(line: &str) -> Result<Command, ParseError> {
        let mut words = line.split_whitespace();
        let name = words.next().ok_or(ParseError::Empty)?.to_ascii_uppercase();
        let args: Vec<&str> = words.collect();
        match (name.as_str(), args.as_slice()) {
            ("GET", [key]) => Ok(Command::Get(key.to_string())),
            ("SET", [key, value]) => Ok(Command::Set { key: key.to_string(), value: value.to_string(), ttl_ms: None }),
            ("SET", [key, value, unit, amount]) => {
                Ok(Command::Set { key: key.to_string(), value: value.to_string(), ttl_ms: Some(expiry(unit, amount)?) })
            }
            ("DEL", [key]) => Ok(Command::Del(key.to_string())),
            ("INCR", [key]) => Ok(Command::Incr(key.to_string())),
            ("KEYS", []) => Ok(Command::Keys),
            ("STATS", []) => Ok(Command::Stats),
            ("GET", _) => Err(ParseError::Usage("GET key")),
            ("SET", _) => Err(ParseError::Usage("SET key value [EX seconds | PX milliseconds]")),
            ("DEL", _) => Err(ParseError::Usage("DEL key")),
            ("INCR", _) => Err(ParseError::Usage("INCR key")),
            ("KEYS", _) | ("STATS", _) => Err(ParseError::Usage("KEYS / STATS take no arguments")),
            (other, _) => Err(ParseError::Unknown(other.to_lowercase())),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum Response {
    Ok,
    Value(String),
    Nil,
    Int(i64),
    List(Vec<String>),
    Error(String),
}

impl fmt::Display for Response {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Response::Ok => write!(f, "OK"),
            Response::Value(v) => write!(f, "{v}"),
            Response::Nil => write!(f, "(nil)"),
            Response::Int(n) => write!(f, "{n}"),
            Response::List(items) if items.is_empty() => write!(f, "(empty)"),
            Response::List(items) => write!(f, "{}", items.join(",")),
            Response::Error(msg) => write!(f, "ERR {msg}"),
        }
    }
}

impl<E: std::error::Error> From<E> for Response {
    fn from(e: E) -> Response {
        Response::Error(e.to_string())
    }
}

// ---------- store ----------

pub trait Clock: Send + Sync + 'static {
    fn now_ms(&self) -> u64;
}

pub struct SystemClock;
impl Clock for SystemClock {
    fn now_ms(&self) -> u64 {
        SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as u64
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Entry {
    value: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    expires_at: Option<u64>,
}

#[derive(Debug, Error)]
pub enum StoreError {
    #[error("value of '{0}' is not an integer")]
    NotAnInteger(String),
    #[error("increment would overflow")]
    Overflow,
}

pub struct Store<C: Clock> {
    entries: HashMap<String, Entry>,
    clock: C,
}

impl<C: Clock> Store<C> {
    pub fn new(clock: C) -> Self {
        Store { entries: HashMap::new(), clock }
    }

    fn live(&self, key: &str) -> Option<&Entry> {
        let now = self.clock.now_ms();
        self.entries.get(key).filter(|e| e.expires_at.is_none_or(|t| t > now))
    }

    pub fn execute(&mut self, cmd: Command) -> Response {
        match cmd {
            Command::Get(k) => self.live(&k).map_or(Response::Nil, |e| Response::Value(e.value.clone())),
            Command::Set { key, value, ttl_ms } => {
                let expires_at = ttl_ms.map(|t| self.clock.now_ms() + t);
                self.entries.insert(key, Entry { value, expires_at });
                Response::Ok
            }
            Command::Del(k) => {
                let existed = self.live(&k).is_some();
                self.entries.remove(&k);
                Response::Int(existed as i64)
            }
            Command::Incr(k) => match self.incr(&k) {
                Ok(n) => Response::Int(n),
                Err(e) => e.into(),
            },
            Command::Keys => {
                let mut keys: Vec<String> = self.entries.keys().filter(|k| self.live(k).is_some()).cloned().collect();
                keys.sort();
                Response::List(keys)
            }
            Command::Stats => Response::Error(String::from("STATS is handled by the server")),
        }
    }

    fn incr(&mut self, key: &str) -> Result<i64, StoreError> {
        let (current, expires_at) = match self.live(key) {
            Some(e) => (e.value.parse::<i64>().map_err(|_| StoreError::NotAnInteger(key.to_string()))?, e.expires_at),
            None => (0, None),
        };
        let next = current.checked_add(1).ok_or(StoreError::Overflow)?;
        self.entries.insert(key.to_string(), Entry { value: next.to_string(), expires_at });
        Ok(next)
    }

    pub fn purge_expired(&mut self) -> usize {
        let now = self.clock.now_ms();
        let before = self.entries.len();
        self.entries.retain(|_, e| e.expires_at.is_none_or(|t| t > now));
        before - self.entries.len()
    }

    pub fn snapshot(&self) -> BTreeMap<String, Entry> {
        self.entries.iter().filter(|(k, _)| self.live(k).is_some()).map(|(k, v)| (k.clone(), v.clone())).collect()
    }

    pub fn restore(&mut self, entries: BTreeMap<String, Entry>) {
        self.entries.extend(entries);
    }
}

// ---------- server ----------

#[derive(Default)]
pub struct Stats {
    commands: AtomicU64,
    connections: AtomicU64,
    expired: AtomicU64,
}

#[derive(Clone)]
pub struct Db {
    store: Arc<Mutex<Store<SystemClock>>>,
    stats: Arc<Stats>,
}

impl Db {
    pub fn new() -> Db {
        Db { store: Arc::new(Mutex::new(Store::new(SystemClock))), stats: Arc::default() }
    }

    pub fn execute(&self, line: &str) -> Response {
        self.stats.commands.fetch_add(1, Relaxed);
        let cmd = match line.parse::<Command>() {
            Ok(cmd) => cmd,
            Err(e) => return e.into(),
        };
        if cmd == Command::Stats {
            return Response::Value(format!(
                "commands={} connections={} keys={} expired={}",
                self.stats.commands.load(Relaxed),
                self.stats.connections.load(Relaxed),
                self.store.lock().unwrap().snapshot().len(),
                self.stats.expired.load(Relaxed)
            ));
        }
        self.store.lock().unwrap().execute(cmd) // lock held only for this call
    }
}

async fn handle_connection(db: Db, stream: TcpStream) -> std::io::Result<()> {
    db.stats.connections.fetch_add(1, Relaxed);
    let (reader, mut writer) = stream.into_split();
    let mut lines = BufReader::new(reader).lines();
    while let Some(line) = lines.next_line().await? {
        if line.trim().is_empty() {
            continue;
        }
        let response = db.execute(&line);
        writer.write_all(format!("{response}\n").as_bytes()).await?;
    }
    Ok(())
}

async fn expiry_task(db: Db, mut shutdown: watch::Receiver<bool>) {
    let mut tick = tokio::time::interval(Duration::from_millis(20));
    loop {
        tokio::select! {
            _ = tick.tick() => {
                let purged = db.store.lock().unwrap().purge_expired();
                db.stats.expired.fetch_add(purged as u64, Relaxed);
            }
            _ = shutdown.changed() => return,
        }
    }
}

/// Runs the server until `shutdown` flips to true, then saves a snapshot to `snapshot_path`.
pub async fn run(listener: TcpListener, db: Db, mut shutdown: watch::Receiver<bool>, snapshot_path: std::path::PathBuf) -> Result<usize> {
    let mut connections = JoinSet::new();
    let expiry = tokio::spawn(expiry_task(db.clone(), shutdown.clone()));

    loop {
        tokio::select! {
            accepted = listener.accept() => {
                let (stream, _) = accepted.context("accepting a connection")?;
                connections.spawn(handle_connection(db.clone(), stream));
            }
            _ = shutdown.changed() => break,
        }
    }

    connections.shutdown().await; // stop connection tasks
    expiry.await.context("expiry task panicked")?;
    let snapshot = db.store.lock().unwrap().snapshot();
    tokio::fs::write(&snapshot_path, serde_json::to_string(&snapshot)?).await.context("saving snapshot")?;
    Ok(snapshot.len())
}

// ---------- a tiny client ----------

struct Client {
    stream: BufReader<TcpStream>,
    name: &'static str,
}

impl Client {
    async fn connect(addr: std::net::SocketAddr, name: &'static str) -> Result<Client> {
        Ok(Client { stream: BufReader::new(TcpStream::connect(addr).await?), name })
    }

    async fn send(&mut self, line: &str) -> Result<String> {
        self.stream.get_mut().write_all(format!("{line}\n").as_bytes()).await?;
        let mut reply = String::new();
        self.stream.read_line(&mut reply).await?;
        Ok(reply.trim_end().to_string())
    }

    async fn show(&mut self, line: &str) -> Result<()> {
        let reply = self.send(line).await?;
        println!("{:<5} {line:<24} -> {reply}", self.name);
        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let addr = listener.local_addr()?;
    let snapshot_path = std::env::temp_dir().join(format!("kvd-{}.json", std::process::id()));
    let (shutdown_tx, shutdown_rx) = watch::channel(false);
    let server = tokio::spawn(run(listener, Db::new(), shutdown_rx, snapshot_path.clone()));

    // One interactive client.
    let mut ana = Client::connect(addr, "ana").await?;
    for cmd in ["SET user:1 ana", "GET user:1", "SET session xyz PX 50", "GET session", "INCR user:1", "FLY away", "SET only"] {
        ana.show(cmd).await?;
    }

    // Ten clients incrementing the same counter concurrently.
    let mut clients = JoinSet::new();
    for _ in 0..10 {
        clients.spawn(async move {
            let mut c = Client::connect(addr, "bot").await?;
            for _ in 0..20 {
                c.send("INCR visits").await?;
            }
            anyhow::Ok(())
        });
    }
    while let Some(done) = clients.join_next().await {
        done??;
    }

    tokio::time::sleep(Duration::from_millis(120)).await; // let `session` expire and be purged
    for cmd in ["GET visits", "GET session", "DEL user:1", "KEYS", "STATS"] {
        ana.show(cmd).await?;
    }

    // Graceful shutdown: stop the server, which writes a snapshot.
    shutdown_tx.send(true)?;
    let saved = server.await??;
    let restored: BTreeMap<String, Entry> = serde_json::from_str(&tokio::fs::read_to_string(&snapshot_path).await?)?;
    tokio::fs::remove_file(&snapshot_path).await?;
    println!("snapshot saved {saved} key(s): {:?}", restored.keys().collect::<Vec<_>>());
    Ok(())
}
Output
ana   SET user:1 ana           -> OK
ana   GET user:1               -> ana
ana   SET session xyz PX 50    -> OK
ana   GET session              -> xyz
ana   INCR user:1              -> ERR value of 'user:1' is not an integer
ana   FLY away                 -> ERR unknown command 'fly'
ana   SET only                 -> ERR wrong arguments, usage: SET key value [EX seconds | PX milliseconds]
ana   GET visits               -> 200
ana   GET session              -> (nil)
ana   DEL user:1               -> 1
ana   KEYS                     -> visits
ana   STATS                    -> commands=212 connections=11 keys=1 expired=1
snapshot saved 1 key(s): ["visits"]

done?? unwraps two layers: the outer ? handles a panicked task (JoinError), the inner one the client's own anyhow::Error. And every STATS counter is consistent even though eleven clients ran concurrently: commands counts every line from every connection.

14. An end-to-end test

The same server can be tested through its public API with #[tokio::test]: start it on a random port, talk to it over TCP, and assert on the replies. This compact version of the server keeps the test self-contained.

Rustcargo test
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};

type Db = Arc<Mutex<HashMap<String, String>>>;

fn execute(db: &Db, line: &str) -> String {
    let words: Vec<&str> = line.split_whitespace().collect();
    let mut map = db.lock().unwrap();
    match words.as_slice() {
        ["SET", k, v] => {
            map.insert(k.to_string(), v.to_string());
            String::from("OK")
        }
        ["GET", k] => map.get(*k).cloned().unwrap_or_else(|| String::from("(nil)")),
        ["INCR", k] => {
            let n = map.get(*k).and_then(|v| v.parse::<i64>().ok()).unwrap_or(0) + 1;
            map.insert(k.to_string(), n.to_string());
            n.to_string()
        }
        _ => format!("ERR unknown command '{}'", words.first().unwrap_or(&"")),
    }
}

pub async fn serve(listener: TcpListener, db: Db) {
    loop {
        let (stream, _) = listener.accept().await.unwrap();
        let db = Arc::clone(&db);
        tokio::spawn(async move {
            let (r, mut w) = stream.into_split();
            let mut lines = BufReader::new(r).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                let reply = execute(&db, &line);
                w.write_all(format!("{reply}\n").as_bytes()).await.unwrap();
            }
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    async fn start() -> std::net::SocketAddr {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(serve(listener, Db::default()));
        addr
    }

    async fn send(conn: &mut BufReader<TcpStream>, line: &str) -> String {
        conn.get_mut().write_all(format!("{line}\n").as_bytes()).await.unwrap();
        let mut reply = String::new();
        conn.read_line(&mut reply).await.unwrap();
        reply.trim_end().to_string()
    }

    #[tokio::test]
    async fn set_then_get() {
        let mut c = BufReader::new(TcpStream::connect(start().await).await.unwrap());
        assert_eq!(send(&mut c, "SET lang rust").await, "OK");
        assert_eq!(send(&mut c, "GET lang").await, "rust");
        assert_eq!(send(&mut c, "GET missing").await, "(nil)");
    }

    #[tokio::test]
    async fn unknown_commands_get_an_error() {
        let mut c = BufReader::new(TcpStream::connect(start().await).await.unwrap());
        assert_eq!(send(&mut c, "FLY away").await, "ERR unknown command 'FLY'");
    }

    #[tokio::test]
    async fn concurrent_increments_are_not_lost() {
        let addr = start().await;
        let mut tasks = Vec::new();
        for _ in 0..5 {
            tasks.push(tokio::spawn(async move {
                let mut c = BufReader::new(TcpStream::connect(addr).await.unwrap());
                for _ in 0..10 {
                    send(&mut c, "INCR n").await;
                }
            }));
        }
        for t in tasks {
            t.await.unwrap();
        }
        let mut c = BufReader::new(TcpStream::connect(addr).await.unwrap());
        assert_eq!(send(&mut c, "GET n").await, "50");
    }
}
Output
running 3 tests
test tests::concurrent_increments_are_not_lost ... ok
test tests::set_then_get ... ok
test tests::unknown_commands_get_an_error ... ok

test result: ok. 3 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

15. Where to take it next

kvd is a good base for practice. Each extension exercises a different part of the course:

Extension You'll use
a SUBSCRIBE channel / PUBLISH channel msg pair broadcast channels, select! per connection
an LRU eviction policy with a maximum key count HashMap + VecDeque or a linked structure, maybe unsafe
the real Redis protocol (RESP), so redis-cli can connect byte parsing, bytes::BytesMut, state machines
command-line options (--port, --snapshot) the clap crate with derive macros
structured logs and metrics tracing, tracing-subscriber
benchmarks of execute and of the whole server criterion, profiling with cargo flamegraph
a workspace: kvd-core, kvd-server, kvd-cli Cargo workspaces, pub APIs across crates
Ctrl-C handling in main tokio::signal::ctrl_c() feeding the watch shutdown channel

Where to go after this course

  • Port something you built in another language. A typical web service maps onto axum + sqlx + serde + tokio + tracing.
  • Read real Rust code. ripgrep, tokio, serde, and axum are well commented. The standard library's source (the "source" link on every docs page) shows how Vec, HashMap, and Arc are implemented.
  • Run cargo clippy on everything and read the explanation behind each lint it reports.
  • Official references: The Rust Programming Language, the Rust Reference, the Async Book, and the Rustonomicon for unsafe code.

Summary

  • The project built a working TCP key-value server from small, tested parts: an enum Command parsed with FromStr and thiserror, a Store with TTLs and an injected Clock, a Response whose Display is the protocol, a shared Db with Arc<Mutex<...>> and atomic statistics, a Tokio server with a task per connection, a background expiry task, serde snapshots, and graceful shutdown with a watch channel.
  • Each part was testable on its own because logic (parsing, the store, execute) was kept separate from I/O (the server loop).
  • Beyond the tests, the compiler guarantees there are no data races between the connection tasks.