Key-Value Store: Router (Low-Level Design)

With chapter 4 complete, we have a shard that survives the loss of its leader without anybody’s help. Measured against the design of chapter 1, 2 things are still wanting. A client has to discover for itself which member of the shard is the leader, as the program probe did by trying each node in turn and collecting 51 refusals in the process. And all the data is in a single shard, whereas the design calls for the keys to be divided among many.

A single small component supplies both. The router stands between the clients and the shards. Clients send every request to it, and for each request it works out which shard owns the key and which member of that shard is at present the leader, and passes the request on. The clients are thereby relieved of all knowledge of the cluster. They know 1 address, and neither the number of shards nor a change of leader is visible to them.

The router is the least intricate program in this part of the book, and it contains 1 decision that is easy to get wrong and expensive when it is: what to do when the member it has asked does not give a clear answer. A router that retries too readily will perform a write twice, and a router that retries too reluctantly will report failures that it could have concealed. The chapter builds the router, states the rule that settles this question, and then repeats the experiment of chapter 4, killing a leader under load, with the router in front.

What a router does

A router keeps no data. Everything it knows comes from a ring file, read when it starts, and from what the members tell it in the course of its work. This is the statelessness that chapter 1 required of routers, and its practical value is considerable: any number of routers may be run side by side, and any of them may be stopped or restarted at any time, since a router that has lost its memory rebuilds it from the first few answers it receives.

The ring file

$ cat ring.json
{
  "virtual_nodes": 256,
  "shards": [
    { "name": "shard-a", "members": ["http://127.0.0.1:7011", "http://127.0.0.1:7012", "http://127.0.0.1:7013"] },
    { "name": "shard-b", "members": ["http://127.0.0.1:7021", "http://127.0.0.1:7022", "http://127.0.0.1:7023"] }
  ]
}

Each shard is a group of 3 kvraft members, listed in the order of their ids. In the design of chapter 1, the ring is held by a coordinator and changes while the system runs. Here it is a file that the router reads once at startup. That is sufficient for routing requests, and it leaves changes to the ring, which are a substantial subject in themselves, outside this chapter.

Listing 5.1 The imports of kvrouter/src/main.rs. kvrouter/src/main.rs

use std::collections::HashMap;
use std::error::Error;
use std::fs;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use axum::body::Bytes;
use axum::extract::{Path, State};
use axum::http::header::{CONTENT_TYPE, ETAG, IF_MATCH, IF_NONE_MATCH};
use axum::http::{HeaderMap, Method, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::{Json, Router};
use clap::Parser;
use kvd::api::{ApiError, valid_key};
use ring::Ring;
use serde::{Deserialize, Serialize};
use tokio::net::TcpListener;
use tokio::time::sleep;

Listing 5.2 The ring file, and what the router keeps for each shard. kvrouter/src/main.rs

#[derive(Parser)]
struct Options {
    #[arg(long, default_value = "127.0.0.1:7000")]
    listen: SocketAddr,

    /// A JSON file that lists the shards and the addresses of their members.
    #[arg(long)]
    ring: PathBuf,
}

/// The contents of the ring file.
#[derive(Deserialize)]
struct Config {
    virtual_nodes: u32,
    shards: Vec<ShardConfig>,
}

#[derive(Deserialize)]
struct ShardConfig {
    name: String,
    /// The HTTP address of each member, in the order of their ids: member 1 first.
    members: Vec<String>,
}

/// A shard, as the router knows it.
struct Shard {
    members: Vec<String>,
    /// The position in `members` of the member that led when the router last asked.
    leader: AtomicUsize,
}

struct Routes {
    ring: Ring,
    shards: HashMap<String, Shard>,
    client: reqwest::Client,
}

For every shard the router remembers a single number, the position in the list of the member that was the leader when it last asked. This is a guess and nothing more. The sections that follow show what happens when the guess is wrong, and the important property of the design is that a wrong guess costs an extra request and never a wrong result.

Finding the shard

Each shard is placed on the ring of chapter 1 with 256 tokens. To look up a key then takes 2 steps: the ring is asked for the name of the owner, and the shard of that name is fetched.

Listing 5.3 Building the ring, and finding the shard of a key. kvrouter/src/main.rs

impl Routes {
    /// Builds the ring from the file: every shard receives its tokens.
    fn load(config: Config) -> Routes {
        let mut ring = Ring::new(config.virtual_nodes);
        let mut shards = HashMap::new();

        for shard in config.shards {
            ring.add(&shard.name);
            shards.insert(
                shard.name,
                Shard {
                    members: shard.members,
                    leader: AtomicUsize::new(0),
                },
            );
        }

        Routes {
            ring,
            shards,
            client: reqwest::Client::new(),
        }
    }

    /// The shard that owns a key.
    fn shard_of(&self, key: &str) -> Option<&Shard> {
        self.shards
            .get(self.ring.owner(key)?)
    }
}

Finding the leader

The router sends the request to the member that it believes to be the leader. Usually the belief is correct. When it is not, the answer of the member tells the router what to do next.

All the cases are decided by a single rule: the router asks another member only if the first attempt has certainly changed nothing. A follower that answers 421 did not touch its log. Nor did a member that reports that the shard is electing a leader, nor one whose process is not running and whose connection was refused. In all 3 cases, to ask another member is safe, since whatever that member does will be the only effect of the request.

A timeout is a different matter altogether. The leader may have appended the record, and even committed it, before the connection failed, and the router cannot know. If it were to send the request to another member, a write might be performed twice, and a conditional write that had succeeded might be reported as failed. The router therefore answers 503 and leaves the decision to the client, which is the only party that knows whether its request is safe to repeat. This is the lost response of chapter 1 once again, and the router’s rule is a refusal to guess about it on somebody else’s behalf.

In the code, the possible outcomes of asking a member are the variants of an enum.

Listing 5.4 The 3 outcomes of asking a member. kvrouter/src/main.rs

/// What a member said about a request it did not serve.
#[derive(Deserialize)]
struct Problem {
    #[serde(rename = "type")]
    kind: String,
    #[serde(default)]
    detail: String,
}

impl Problem {
    /// The position of the leader in the shard's list of members, when the
    /// member named it, as in "the leader is member 3".
    fn named_leader(&self) -> Option<usize> {
        let id: usize = self
            .detail
            .rsplit(' ')
            .next()?
            .parse()
            .ok()?;

        id.checked_sub(1)
    }
}

/// The result of asking one member.
enum Attempt {
    /// The member served the request. This is its response.
    Served(Response),
    /// The member changed nothing. It may have named the leader.
    Refused(Option<usize>),
    /// The request may or may not have taken effect.
    Unknown(String),
}

ask sends the request to 1 member and classifies what comes back as 1 of the outcomes.

Listing 5.5 Asking 1 member. kvrouter/src/main.rs

const NO_LEADER: &str = "/problems/no-leader";

impl Routes {
    /// Sends the request to one member and classifies what comes back.
    async fn ask(
        &self,
        member: &str,
        method: Method,
        key: &str,
        headers: &HeaderMap,
        body: Bytes,
    ) -> Attempt {
        let conditions = [IF_MATCH, IF_NONE_MATCH]
            .into_iter()
            .filter_map(|name| {
                headers
                    .get(&name)
                    .map(|value| (name, value))
            });
        let request = self
            .client
            .request(method, format!("{member}/v1/keys/{key}"))
            .body(body);
        let request = conditions.fold(request, |request, (name, value)| {
            request.header(name, value)
        });

        let response = match request.send().await {
            Ok(response) => response,
            Err(error) if error.is_connect() => return Attempt::Refused(None),
            Err(error) => return Attempt::Unknown(error.to_string()),
        };

        let status = response.status();
        let redirected = status == StatusCode::MISDIRECTED_REQUEST;

        if !redirected && status != StatusCode::SERVICE_UNAVAILABLE {
            return Attempt::Served(relay(response).await);
        }

        match response.json::<Problem>().await {
            Ok(problem) if redirected || problem.kind == NO_LEADER => {
                Attempt::Refused(problem.named_leader())
            }
            Ok(problem) => Attempt::Unknown(problem.detail),
            Err(error) => Attempt::Unknown(error.to_string()),
        }
    }
}

/// Passes a member's response on to the client: status, version and body.
async fn relay(response: reqwest::Response) -> Response {
    let status = response.status();
    let headers: HeaderMap = [ETAG, CONTENT_TYPE]
        .into_iter()
        .filter_map(|name| {
            response
                .headers()
                .get(&name)
                .map(|value| (name, value.clone()))
        })
        .collect();

    match response.bytes().await {
        Ok(body) => (status, headers, body).into_response(),
        Err(error) => ApiError::Unavailable(error.to_string()).into_response(),
    }
}

forward is the loop around it. A refusal moves the guess to the member that the refusal named, or failing that to the next in the list. When every member has refused once, an election must be in progress, and the router then waits 150 ms between further attempts, so that it does not flood a shard which is busy choosing a leader.

Listing 5.6 Trying members until one serves the request. kvrouter/src/main.rs

/// How many times the router asks a member before it gives up.
const ATTEMPTS: usize = 12;

/// How long the router waits before it asks again, once every member of a
/// shard has refused.
const PAUSE: Duration = Duration::from_millis(150);

impl Routes {
    /// Sends a request to the leader of the key's shard. A member that
    /// refuses has changed nothing, so the router asks another.
    async fn forward(
        &self,
        method: Method,
        key: &str,
        headers: &HeaderMap,
        body: Bytes,
    ) -> Result<Response, ApiError> {
        valid_key(key)?;

        let shard = self
            .shard_of(key)
            .ok_or(ApiError::NoLeader)?;

        for attempt in 1..=ATTEMPTS {
            let tried = shard.leader.load(Ordering::Relaxed);
            let member = &shard.members[tried];

            let named = match self
                .ask(member, method.clone(), key, headers, body.clone())
                .await
            {
                Attempt::Served(response) => return Ok(response),
                Attempt::Unknown(detail) => return Err(ApiError::Unavailable(detail)),
                Attempt::Refused(named) => named,
            };

            let next = named
                .filter(|&position| position < shard.members.len())
                .unwrap_or((tried + 1) % shard.members.len());

            shard
                .leader
                .store(next, Ordering::Relaxed);

            // Every member has refused once: the shard is electing a leader.
            if attempt >= shard.members.len() {
                sleep(PAUSE).await;
            }
        }

        Err(ApiError::NoLeader)
    }
}

Listing 5.7 The HTTP handlers, and the start of the program. kvrouter/src/main.rs

async fn key(
    State(routes): State<Arc<Routes>>,
    method: Method,
    Path(key): Path<String>,
    headers: HeaderMap,
    body: Bytes,
) -> Result<Response, ApiError> {
    routes
        .forward(method, &key, &headers, body)
        .await
}

#[derive(Serialize)]
struct ShardHealth<'a> {
    name: &'a str,
    /// The address the router currently sends this shard's requests to.
    leader: &'a str,
}

async fn health(State(routes): State<Arc<Routes>>) -> Json<serde_json::Value> {
    let mut shards: Vec<ShardHealth> = routes
        .shards
        .iter()
        .map(|(name, shard)| ShardHealth {
            name,
            leader: &shard.members[shard.leader.load(Ordering::Relaxed)],
        })
        .collect();
    shards.sort_by_key(|shard| shard.name);

    Json(serde_json::json!({ "status": "ok", "shards": shards }))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let options = Options::parse();
    let config: Config = serde_json::from_str(&fs::read_to_string(&options.ring)?)?;
    let routes = Arc::new(Routes::load(config));

    println!("routing to {} shards", routes.shards.len());

    let app = Router::new()
        .route("/v1/keys/{key}", get(key).put(key).delete(key))
        .route("/healthz", get(health))
        .with_state(routes);

    let listener = TcpListener::bind(options.listen).await?;
    println!("listening on {}", listener.local_addr()?);

    axum::serve(listener, app).await?;

    Ok(())
}

2 shards behind 1 router

We start 6 kvraft members in 2 groups of 3, then the router, and write 400 keys through it.

$ ./target/release/kvrouter --listen 127.0.0.1:7000 --ring ring.json &
routing to 2 shards
listening on 127.0.0.1:7000
$ ./target/release/kvctl load --server http://127.0.0.1:7000 --count 400 --acked acked.txt
400 writes acknowledged
$ for port in 7013 7023; do curl http://127.0.0.1:$port/healthz; echo; done
{"id":3,"role":"leader","term":1,"leader":3,"last_sequence":206,"committed":206}
{"id":3,"role":"leader","term":1,"leader":3,"last_sequence":196,"committed":196}
$ ./target/release/kvctl verify --server http://127.0.0.1:7000 --acked acked.txt
400 acknowledged, 400 found, 0 missing

The client used a single address throughout. It never learned how many shards exist, or which members lead them.

A leader is lost behind the router

probe, from chapter 4, is now given the router’s address and no other. While it is writing, the leader of shard-a is killed.

$ ./target/release/probe --seconds 12 --acked acked.txt --server http://127.0.0.1:7000 &
$ sleep 4; kill -9 <pid of member 3 of shard-a, its leader>
780 writes acknowledged in 12 seconds, 0 attempts refused
the longest pause between 2 acknowledged writes was 1.08 s, ending 5.2 s after the start
$ curl http://127.0.0.1:7000/healthz
{"shards":[{"leader":"http://127.0.0.1:7011","name":"shard-a"},{"leader":"http://127.0.0.1:7023","name":"shard-b"}],"status":"ok"}
$ ./target/release/kvctl verify --server http://127.0.0.1:7000 --acked acked.txt
780 acknowledged, 780 found, 0 missing

The comparison with chapter 4 shows what the router contributes. The election took as long as it did before, about a second, since the router does nothing to hasten it. What has changed is what the client experiences. A second of refusals, which every client would have had to be written to handle, has become a second of waiting, which requires nothing of the client at all.

Writes to shard-b were at no time affected. probe sends its keys in turn, and so it waited behind the 1 request for shard-a that was being held. Independent clients writing to shard-b would have noticed nothing.

What is still missing

The complete program

Listing 5.8 The complete kvrouter/Cargo.toml. kvrouter/Cargo.toml

[package]
name = "kvrouter"
version = "0.1.0"
edition.workspace = true
rust-version.workspace = true
license.workspace = true
publish.workspace = true

[dependencies]
kvd.workspace = true
ring.workspace = true
axum.workspace = true
clap.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true

Listing 5.9 The complete kvrouter/src/main.rs. kvrouter/src/main.rs

//! Command kvrouter sends each request to the shard that owns its key, and
//! within that shard to the member that leads. Clients talk to the router
//! and need to know nothing about shards or leaders.
//!
//! Run it from code/ with:
//!
//!     cargo run -p kvrouter -- --listen 127.0.0.1:7000 --ring ring.json

use std::collections::HashMap;
use std::error::Error;
use std::fs;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use axum::body::Bytes;
use axum::extract::{Path, State};
use axum::http::header::{CONTENT_TYPE, ETAG, IF_MATCH, IF_NONE_MATCH};
use axum::http::{HeaderMap, Method, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::{Json, Router};
use clap::Parser;
use kvd::api::{ApiError, valid_key};
use ring::Ring;
use serde::{Deserialize, Serialize};
use tokio::net::TcpListener;
use tokio::time::sleep;

#[derive(Parser)]
struct Options {
    #[arg(long, default_value = "127.0.0.1:7000")]
    listen: SocketAddr,

    /// A JSON file that lists the shards and the addresses of their members.
    #[arg(long)]
    ring: PathBuf,
}

/// The contents of the ring file.
#[derive(Deserialize)]
struct Config {
    virtual_nodes: u32,
    shards: Vec<ShardConfig>,
}

#[derive(Deserialize)]
struct ShardConfig {
    name: String,
    /// The HTTP address of each member, in the order of their ids: member 1 first.
    members: Vec<String>,
}

/// A shard, as the router knows it.
struct Shard {
    members: Vec<String>,
    /// The position in `members` of the member that led when the router last asked.
    leader: AtomicUsize,
}

struct Routes {
    ring: Ring,
    shards: HashMap<String, Shard>,
    client: reqwest::Client,
}

impl Routes {
    /// Builds the ring from the file: every shard receives its tokens.
    fn load(config: Config) -> Routes {
        let mut ring = Ring::new(config.virtual_nodes);
        let mut shards = HashMap::new();

        for shard in config.shards {
            ring.add(&shard.name);
            shards.insert(
                shard.name,
                Shard {
                    members: shard.members,
                    leader: AtomicUsize::new(0),
                },
            );
        }

        Routes {
            ring,
            shards,
            client: reqwest::Client::new(),
        }
    }

    /// The shard that owns a key.
    fn shard_of(&self, key: &str) -> Option<&Shard> {
        self.shards
            .get(self.ring.owner(key)?)
    }
}

/// What a member said about a request it did not serve.
#[derive(Deserialize)]
struct Problem {
    #[serde(rename = "type")]
    kind: String,
    #[serde(default)]
    detail: String,
}

impl Problem {
    /// The position of the leader in the shard's list of members, when the
    /// member named it, as in "the leader is member 3".
    fn named_leader(&self) -> Option<usize> {
        let id: usize = self
            .detail
            .rsplit(' ')
            .next()?
            .parse()
            .ok()?;

        id.checked_sub(1)
    }
}

/// The result of asking one member.
enum Attempt {
    /// The member served the request. This is its response.
    Served(Response),
    /// The member changed nothing. It may have named the leader.
    Refused(Option<usize>),
    /// The request may or may not have taken effect.
    Unknown(String),
}

const NO_LEADER: &str = "/problems/no-leader";

impl Routes {
    /// Sends the request to one member and classifies what comes back.
    async fn ask(
        &self,
        member: &str,
        method: Method,
        key: &str,
        headers: &HeaderMap,
        body: Bytes,
    ) -> Attempt {
        let conditions = [IF_MATCH, IF_NONE_MATCH]
            .into_iter()
            .filter_map(|name| {
                headers
                    .get(&name)
                    .map(|value| (name, value))
            });
        let request = self
            .client
            .request(method, format!("{member}/v1/keys/{key}"))
            .body(body);
        let request = conditions.fold(request, |request, (name, value)| {
            request.header(name, value)
        });

        let response = match request.send().await {
            Ok(response) => response,
            Err(error) if error.is_connect() => return Attempt::Refused(None),
            Err(error) => return Attempt::Unknown(error.to_string()),
        };

        let status = response.status();
        let redirected = status == StatusCode::MISDIRECTED_REQUEST;

        if !redirected && status != StatusCode::SERVICE_UNAVAILABLE {
            return Attempt::Served(relay(response).await);
        }

        match response.json::<Problem>().await {
            Ok(problem) if redirected || problem.kind == NO_LEADER => {
                Attempt::Refused(problem.named_leader())
            }
            Ok(problem) => Attempt::Unknown(problem.detail),
            Err(error) => Attempt::Unknown(error.to_string()),
        }
    }
}

/// Passes a member's response on to the client: status, version and body.
async fn relay(response: reqwest::Response) -> Response {
    let status = response.status();
    let headers: HeaderMap = [ETAG, CONTENT_TYPE]
        .into_iter()
        .filter_map(|name| {
            response
                .headers()
                .get(&name)
                .map(|value| (name, value.clone()))
        })
        .collect();

    match response.bytes().await {
        Ok(body) => (status, headers, body).into_response(),
        Err(error) => ApiError::Unavailable(error.to_string()).into_response(),
    }
}

/// How many times the router asks a member before it gives up.
const ATTEMPTS: usize = 12;

/// How long the router waits before it asks again, once every member of a
/// shard has refused.
const PAUSE: Duration = Duration::from_millis(150);

impl Routes {
    /// Sends a request to the leader of the key's shard. A member that
    /// refuses has changed nothing, so the router asks another.
    async fn forward(
        &self,
        method: Method,
        key: &str,
        headers: &HeaderMap,
        body: Bytes,
    ) -> Result<Response, ApiError> {
        valid_key(key)?;

        let shard = self
            .shard_of(key)
            .ok_or(ApiError::NoLeader)?;

        for attempt in 1..=ATTEMPTS {
            let tried = shard.leader.load(Ordering::Relaxed);
            let member = &shard.members[tried];

            let named = match self
                .ask(member, method.clone(), key, headers, body.clone())
                .await
            {
                Attempt::Served(response) => return Ok(response),
                Attempt::Unknown(detail) => return Err(ApiError::Unavailable(detail)),
                Attempt::Refused(named) => named,
            };

            let next = named
                .filter(|&position| position < shard.members.len())
                .unwrap_or((tried + 1) % shard.members.len());

            shard
                .leader
                .store(next, Ordering::Relaxed);

            // Every member has refused once: the shard is electing a leader.
            if attempt >= shard.members.len() {
                sleep(PAUSE).await;
            }
        }

        Err(ApiError::NoLeader)
    }
}

async fn key(
    State(routes): State<Arc<Routes>>,
    method: Method,
    Path(key): Path<String>,
    headers: HeaderMap,
    body: Bytes,
) -> Result<Response, ApiError> {
    routes
        .forward(method, &key, &headers, body)
        .await
}

#[derive(Serialize)]
struct ShardHealth<'a> {
    name: &'a str,
    /// The address the router currently sends this shard's requests to.
    leader: &'a str,
}

async fn health(State(routes): State<Arc<Routes>>) -> Json<serde_json::Value> {
    let mut shards: Vec<ShardHealth> = routes
        .shards
        .iter()
        .map(|(name, shard)| ShardHealth {
            name,
            leader: &shard.members[shard.leader.load(Ordering::Relaxed)],
        })
        .collect();
    shards.sort_by_key(|shard| shard.name);

    Json(serde_json::json!({ "status": "ok", "shards": shards }))
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let options = Options::parse();
    let config: Config = serde_json::from_str(&fs::read_to_string(&options.ring)?)?;
    let routes = Arc::new(Routes::load(config));

    println!("routing to {} shards", routes.shards.len());

    let app = Router::new()
        .route("/v1/keys/{key}", get(key).put(key).delete(key))
        .route("/healthz", get(health))
        .with_state(routes);

    let listener = TcpListener::bind(options.listen).await?;
    println!("listening on {}", listener.local_addr()?);

    axum::serve(listener, app).await?;

    Ok(())
}