Skip to Content
Course content

228: Building a Distributed Key-Value Store

Click on the "Edit" button in the top corner of the screen to edit your slide content.

How do I keep the local data store thread-safe without killing performance?

When you're building a KV store, the most tempting thing to do is wrap a HashMap in a Mutex. I've seen a lot of people do this early on, but in a distributed system where you're hammering the store with thousands of concurrent reads, a Mutex becomes a massive bottleneck because it blocks every single reader while another reader is just... reading.

You want an RwLock (Read-Write Lock). This allows any number of concurrent readers, but only one writer. In Rust, std::sync::RwLock is okay, but if you're using tokio for your networking—which you should be—use tokio::sync::RwLock to avoid blocking the entire async runtime thread.

use tokio::sync::RwLock;
use std::collections::HashMap;
use std::sync::Arc;

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

async fn get_value(db: Db, key: String) -> Option<String> {
    let read_guard = db.read().await;
    read_guard.get(&key).cloned()
}

async fn set_value(db: Db, key: String, value: String) {
    let mut write_guard = db.write().await;
    write_guard.insert(key, value);
}

Pro tip: If you find the RwLock is still too slow, look into DashMap. It shards the map internally so you aren't locking the entire database for every single write.

How do we actually decide which node owns which key?

You can't just use key % node_count. Why? Because the second you add a new node to your cluster to scale up, your node_count changes, and suddenly almost every key in your system maps to a different node. You'd have to move nearly all your data across the network, which is a nightmare.

This is where Consistent Hashing comes in. Imagine a circle (a ring) from 0 to 232-1. You hash your node IDs and place them on that ring. When a request comes in for a key, you hash the key and move clockwise around the ring until you hit the first node. If a node drops out, only the keys that were mapped to that specific node move to the next neighbor. Everyone else stays put.

I usually implement this using a BTreeMap in Rust, where the keys are the hash values of the nodes and the values are the node addresses. It makes finding the "next" node in the ring a simple range query.

What's the most "Rust-y" way to handle the communication between nodes?

Don't waste your time writing a custom TCP protocol with manual byte parsing unless you're doing this for a PhD thesis. Use gRPC via the tonic crate. It gives you a strictly typed contract via Protocol Buffers, which is essential when you have multiple nodes that need to agree on exactly what a "PutRequest" looks like.

Here is a simplified version of how your service definition would look in .proto:

syntax = "proto3";
package kvstore;

service KvService {
  rpc Get (GetRequest) returns (GetResponse);
  rpc Put (PutRequest) returns (PutResponse);
}

message GetRequest { string key = 1; }
message GetResponse { string value = 1; bool found = 2; }
message PutRequest { string key = 1; string value = 2; }
message PutResponse { bool success = 1; }

Tonic generates the traits and structs for you. You just implement the trait for your store. It handles the async plumbing and HTTP/2 multiplexing under the hood, letting you focus on the distribution logic rather than worrying about packet fragmentation.

How do I prevent data loss when a node restarts?

If you keep everything in a HashMap, a power flicker wipes your entire dataset. You could serialize the whole map to disk every few seconds, but that's horribly inefficient as the DB grows. Instead, use a Write-Ahead Log (WAL).

Every time a Put request comes in, you append that operation to a simple, append-only file on disk before updating the in-memory map. Appending is incredibly fast because it's a sequential write. If the node crashes, you simply replay the log from the beginning to rebuild your HashMap in memory.

use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;

async fn log_operation(key: &str, value: &str) -> tokio::io::Result<()> {
    let mut file = OpenOptions::new()
        .append(true)
        .create(true)
        .open("wal.log")
        .await?;

    let entry = format!("SET {} {}\n", key, value);
    file.write_all(entry.as_bytes()).await?;
    file.sync_all().await?; // Ensure it actually hits the platter
    Ok(())
}

Eventually, your log will get too huge. That's when you'll want to implement "compaction" (or snapshotting), where you save the current state of the map and delete the old log entries, but that's a topic for a much deeper dive.




📋 Practical Task

Exercise: Implement a Consistent Hash Ring

Your goal is to build a basic Consistent Hashing mechanism to determine which node should handle a specific key. This is the foundation of your distributed store's routing logic.

  • Create a HashRing struct that holds a BTreeMap<u64, String>, where the key is a hash and the value is the node's address (e.g., "127.0.0.1:8080").
  • Implement an add_node method that hashes the node's address and inserts it into the map.
  • Implement a get_node method that takes a key (string), hashes it, and finds the first node in the map with a hash greater than or equal to the key's hash.
  • Ensure that if the key's hash is greater than the highest node hash, it "wraps around" to the first node in the map.
  • Use the DefaultHasher from std::collections::hash_map for simplicity.

Test your implementation by adding three nodes and verifying that a set of 10 different keys are distributed across them predictably.

Rating
0 0

There are no comments for now.

to be the first to leave a comment.