Design a distributed cache

package main
 
import (
	"fmt"
	"hash/fnv"
	"slices"
	"sort"
	"strconv"
)
 
func hash(s string) uint32 {
	h := fnv.New32a()
	h.Write([]byte(s))
	x := h.Sum32() // mix the bits so similar names land far apart
	x ^= x >> 16
	x *= 0x85ebca6b
	x ^= x >> 13
	x *= 0xc2b2ae35
	x ^= x >> 16
	return x
}
 
type ring struct {
	points []uint32
	owner  map[uint32]string
}
 
func newRing(nodes []string, vnodes int) *ring {
	r := &ring{owner: map[uint32]string{}}
	for _, n := range nodes {
		for v := range vnodes {
			p := hash(n + "#" + strconv.Itoa(v))
			r.points = append(r.points, p)
			r.owner[p] = n
		}
	}
	slices.Sort(r.points)
	return r
}
 
func (r *ring) get(key string) string {
	h := hash(key)
	i := sort.Search(len(r.points), func(i int) bool { return r.points[i] >= h })
	if i == len(r.points) {
		i = 0 // wrap around the ring
	}
	return r.owner[r.points[i]]
}
 
func spread(r *ring, keys int) (lo, hi int) {
	count := map[string]int{}
	for k := range keys {
		count[r.get("key"+strconv.Itoa(k))]++
	}
	lo, hi = keys, 0
	for _, c := range count {
		lo, hi = min(lo, c), max(hi, c)
	}
	return lo, hi
}
 
func main() {
	nodes := []string{"a", "b", "c", "d", "e"}
	for _, v := range []int{1, 10, 100, 500} {
		lo, hi := spread(newRing(nodes, v), 100_000)
		fmt.Printf("%3d vnodes: busiest node has %.1fx the keys of the quietest\n", v, float64(hi)/float64(lo))
	}
 
	before := newRing(nodes, 200)
	after := newRing(append(slices.Clone(nodes), "f"), 200)
	moved := 0
	for k := range 100_000 {
		key := "key" + strconv.Itoa(k)
		if before.get(key) != after.get(key) {
			moved++
		}
	}
	fmt.Printf("adding a 6th node moved %.1f%% of keys (ideal: 1/6 = 16.7%%; modulo hashing: ~83%%)\n", float64(moved)/1000)
}

  1 vnodes: busiest node has 12.2x the keys of the quietest
 10 vnodes: busiest node has 2.7x the keys of the quietest
100 vnodes: busiest node has 1.1x the keys of the quietest
500 vnodes: busiest node has 1.1x the keys of the quietest
adding a 6th node moved 17.1% of keys (ideal: 1/6 = 16.7%; modulo hashing: ~83%)

kvcache/
├── go.mod
├── main.go                       HTTP API + page
├── main_test.go
├── internal/lru/lru.go           sharded LRU with TTL
├── internal/lru/lru_test.go
├── internal/ring/ring.go         consistent hashing with virtual nodes
├── internal/ring/ring_test.go
├── internal/cluster/cluster.go   routing + migration
├── internal/cluster/cluster_test.go
└── web/index.html

mkdir kvcache && cd kvcache && go mod init kvcache
mkdir -p internal/lru internal/ring internal/cluster web

module kvcache
 
go 1.24

// Package lru is a concurrency-safe, sharded LRU cache with per-entry TTLs.
package lru
 
import (
	"container/list"
	"hash/maphash"
	"sync"
	"time"
)
 
type entry struct {
	key     string
	val     string
	expires time.Time // zero: never
}
 
type shard struct {
	mu    sync.Mutex
	cap   int
	ll    *list.List // front = most recently used
	items map[string]*list.Element
}
 
// Cache splits keys over independent shards to reduce lock contention.
type Cache struct {
	shards []*shard
	seed   maphash.Seed
	now    func() time.Time
}
 
// New creates a cache holding about capacity entries in total, split over n shards.
func New(capacity, n int, now func() time.Time) *Cache {
	if now == nil {
		now = time.Now
	}
	c := &Cache{shards: make([]*shard, n), seed: maphash.MakeSeed(), now: now}
	for i := range c.shards {
		c.shards[i] = &shard{cap: max(1, capacity/n), ll: list.New(), items: map[string]*list.Element{}}
	}
	return c
}
 
func (c *Cache) shard(key string) *shard {
	return c.shards[maphash.String(c.seed, key)%uint64(len(c.shards))]
}
 
// Set stores key with an optional TTL (0 means no expiry). It reports whether an entry was evicted.
func (c *Cache) Set(key, val string, ttl time.Duration) (evicted bool) {
	s := c.shard(key)
	s.mu.Lock()
	defer s.mu.Unlock()
	var exp time.Time
	if ttl > 0 {
		exp = c.now().Add(ttl)
	}
	if el, ok := s.items[key]; ok {
		e := el.Value.(*entry)
		e.val, e.expires = val, exp
		s.ll.MoveToFront(el)
		return false
	}
	s.items[key] = s.ll.PushFront(&entry{key, val, exp})
	if s.ll.Len() > s.cap {
		oldest := s.ll.Back()
		s.ll.Remove(oldest)
		delete(s.items, oldest.Value.(*entry).key)
		return true
	}
	return false
}
 
// Get returns the value if present and not expired.
func (c *Cache) Get(key string) (string, bool) {
	s := c.shard(key)
	s.mu.Lock()
	defer s.mu.Unlock()
	el, ok := s.items[key]
	if !ok {
		return "", false
	}
	e := el.Value.(*entry)
	if !e.expires.IsZero() && !c.now().Before(e.expires) {
		s.ll.Remove(el) // lazy expiry
		delete(s.items, key)
		return "", false
	}
	s.ll.MoveToFront(el)
	return e.val, true
}
 
// Delete removes key.
func (c *Cache) Delete(key string) {
	s := c.shard(key)
	s.mu.Lock()
	defer s.mu.Unlock()
	if el, ok := s.items[key]; ok {
		s.ll.Remove(el)
		delete(s.items, key)
	}
}
 
// Len counts live and not-yet-swept entries.
func (c *Cache) Len() int {
	n := 0
	for _, s := range c.shards {
		s.mu.Lock()
		n += len(s.items)
		s.mu.Unlock()
	}
	return n
}
 
// Each calls fn for every unexpired entry with its remaining TTL (0 = none). fn must not call back into c.
func (c *Cache) Each(fn func(key, val string, ttl time.Duration)) {
	now := c.now()
	for _, s := range c.shards {
		s.mu.Lock()
		for el := s.ll.Front(); el != nil; el = el.Next() {
			e := el.Value.(*entry)
			var ttl time.Duration
			if !e.expires.IsZero() {
				if ttl = e.expires.Sub(now); ttl <= 0 {
					continue
				}
			}
			fn(e.key, e.val, ttl)
		}
		s.mu.Unlock()
	}
}

package lru
 
import (
	"strconv"
	"sync"
	"testing"
	"time"
)
 
func TestEvictsLeastRecentlyUsed(t *testing.T) {
	c := New(2, 1, nil) // one shard, so eviction order is exact
	c.Set("a", "1", 0)
	c.Set("b", "2", 0)
	c.Get("a")               // a is now the most recent
	if !c.Set("c", "3", 0) { // evicts b
		t.Fatal("expected an eviction")
	}
	if _, ok := c.Get("b"); ok {
		t.Error("b should have been evicted")
	}
	if v, ok := c.Get("a"); !ok || v != "1" {
		t.Error("a should still be there")
	}
}
 
func TestTTL(t *testing.T) {
	now := time.Unix(1000, 0)
	c := New(10, 2, func() time.Time { return now })
	c.Set("session", "xyz", 30*time.Second)
	c.Set("config", "on", 0)
	now = now.Add(29 * time.Second)
	if _, ok := c.Get("session"); !ok {
		t.Fatal("not expired yet")
	}
	now = now.Add(time.Second)
	if _, ok := c.Get("session"); ok {
		t.Fatal("expired at exactly 30s")
	}
	if _, ok := c.Get("config"); !ok {
		t.Fatal("no-TTL entries never expire")
	}
}
 
func TestConcurrentAccess(t *testing.T) {
	c := New(1000, 16, nil)
	var wg sync.WaitGroup
	for g := range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := range 1000 {
				k := strconv.Itoa(g*1000 + i)
				c.Set(k, k, 0)
				c.Get(k)
			}
		}()
	}
	wg.Wait()
	if n := c.Len(); n > 1000+16 { // each shard holds at most cap/n
		t.Errorf("Len = %d, too many", n)
	}
}

// Package ring implements consistent hashing with virtual nodes.
package ring
 
import (
	"hash/fnv"
	"slices"
	"sort"
	"strconv"
	"sync"
)
 
func hash(s string) uint32 {
	h := fnv.New32a()
	h.Write([]byte(s))
	x := h.Sum32()
	x ^= x >> 16 // finalizer: spread similar names
	x *= 0x85ebca6b
	x ^= x >> 13
	x *= 0xc2b2ae35
	x ^= x >> 16
	return x
}
 
// Ring maps keys to nodes. It is safe for concurrent use.
type Ring struct {
	mu     sync.RWMutex
	vnodes int
	points []uint32
	owner  map[uint32]string
	nodes  map[string]bool
}
 
func New(vnodes int) *Ring {
	return &Ring{vnodes: vnodes, owner: map[uint32]string{}, nodes: map[string]bool{}}
}
 
func (r *Ring) Add(node string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.nodes[node] {
		return
	}
	r.nodes[node] = true
	for v := range r.vnodes {
		p := hash(node + "#" + strconv.Itoa(v))
		if _, taken := r.owner[p]; taken {
			continue // an astronomically rare collision: skip this point
		}
		r.owner[p] = node
		r.points = append(r.points, p)
	}
	slices.Sort(r.points)
}
 
func (r *Ring) Remove(node string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if !r.nodes[node] {
		return
	}
	delete(r.nodes, node)
	r.points = slices.DeleteFunc(r.points, func(p uint32) bool {
		if r.owner[p] == node {
			delete(r.owner, p)
			return true
		}
		return false
	})
}
 
// Get returns the node owning key, or "" if the ring is empty.
func (r *Ring) Get(key string) string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	if len(r.points) == 0 {
		return ""
	}
	h := hash(key)
	i := sort.Search(len(r.points), func(i int) bool { return r.points[i] >= h })
	if i == len(r.points) {
		i = 0
	}
	return r.owner[r.points[i]]
}
 
// Nodes lists members in sorted order.
func (r *Ring) Nodes() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]string, 0, len(r.nodes))
	for n := range r.nodes {
		out = append(out, n)
	}
	slices.Sort(out)
	return out
}

package ring
 
import (
	"strconv"
	"testing"
)
 
func TestMinimalMovement(t *testing.T) {
	r := New(200)
	for _, n := range []string{"a", "b", "c", "d"} {
		r.Add(n)
	}
	before := map[string]string{}
	for i := range 20_000 {
		k := strconv.Itoa(i)
		before[k] = r.Get(k)
	}
	r.Add("e")
	moved := 0
	for k, old := range before {
		now := r.Get(k)
		if now != old {
			moved++
			if now != "e" {
				t.Fatalf("key %s moved from %s to %s, but only moves to the new node are allowed", k, old, now)
			}
		}
	}
	if frac := float64(moved) / 20_000; frac < 0.12 || frac > 0.28 {
		t.Errorf("moved %.1f%% of keys, want about 20%%", frac*100)
	}
}
 
func TestBalance(t *testing.T) {
	r := New(200)
	for _, n := range []string{"a", "b", "c"} {
		r.Add(n)
	}
	count := map[string]int{}
	for i := range 30_000 {
		count[r.Get("k"+strconv.Itoa(i))]++
	}
	for n, c := range count {
		if c < 7_000 || c > 13_000 {
			t.Errorf("node %s got %d of 30000 keys", n, c)
		}
	}
}
 
func TestRemoveAndEmpty(t *testing.T) {
	r := New(10)
	r.Add("only")
	if r.Get("x") != "only" {
		t.Fatal("single node owns everything")
	}
	r.Remove("only")
	if r.Get("x") != "" || len(r.Nodes()) != 0 {
		t.Fatal("empty ring returns no owner")
	}
}

// Package cluster routes keys to nodes with a consistent-hash ring and migrates keys on membership changes.
package cluster
 
import (
	"errors"
	"sync"
	"time"
 
	"kvcache/internal/lru"
	"kvcache/internal/ring"
)
 
var ErrNoNodes = errors.New("cluster has no nodes")
 
type Cluster struct {
	mu       sync.RWMutex // guards nodes; held exclusively while migrating
	ring     *ring.Ring
	nodes    map[string]*lru.Cache
	capacity int
	now      func() time.Time
}
 
func New(capacityPerNode, vnodes int, now func() time.Time) *Cluster {
	return &Cluster{ring: ring.New(vnodes), nodes: map[string]*lru.Cache{}, capacity: capacityPerNode, now: now}
}
 
func (c *Cluster) owner(key string) (*lru.Cache, string, error) {
	name := c.ring.Get(key)
	if name == "" {
		return nil, "", ErrNoNodes
	}
	return c.nodes[name], name, nil
}
 
func (c *Cluster) Set(key, val string, ttl time.Duration) (string, error) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	n, name, err := c.owner(key)
	if err != nil {
		return "", err
	}
	n.Set(key, val, ttl)
	return name, nil
}
 
func (c *Cluster) Get(key string) (val, node string, ok bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	n, name, err := c.owner(key)
	if err != nil {
		return "", "", false
	}
	v, ok := n.Get(key)
	return v, name, ok
}
 
func (c *Cluster) Delete(key string) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if n, _, err := c.owner(key); err == nil {
		n.Delete(key)
	}
}
 
type move struct {
	key, val string
	ttl      time.Duration
}
 
// rebalance moves every key whose owner changed. Caller holds c.mu exclusively.
func (c *Cluster) rebalance(from map[string]*lru.Cache) int {
	var moves []move
	for name, cache := range from {
		cache.Each(func(k, v string, ttl time.Duration) {
			if c.ring.Get(k) != name {
				moves = append(moves, move{k, v, ttl})
			}
		})
	}
	for _, m := range moves {
		for _, cache := range from {
			cache.Delete(m.key)
		}
		if n, _, err := c.owner(m.key); err == nil {
			n.Set(m.key, m.val, m.ttl)
		}
	}
	return len(moves)
}
 
// AddNode joins a node and returns how many keys moved to it.
func (c *Cluster) AddNode(name string) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, ok := c.nodes[name]; ok {
		return 0
	}
	c.nodes[name] = lru.New(c.capacity, 8, c.now)
	c.ring.Add(name)
	return c.rebalance(c.nodes)
}
 
// RemoveNode drains a node into the others and returns how many keys moved.
func (c *Cluster) RemoveNode(name string) int {
	c.mu.Lock()
	defer c.mu.Unlock()
	leaving, ok := c.nodes[name]
	if !ok {
		return 0
	}
	c.ring.Remove(name)
	delete(c.nodes, name)
	return c.rebalance(map[string]*lru.Cache{name: leaving})
}
 
// Stats returns the number of keys per node.
func (c *Cluster) Stats() map[string]int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := map[string]int{}
	for name, n := range c.nodes {
		out[name] = n.Len()
	}
	return out
}

package cluster
 
import (
	"strconv"
	"testing"
	"time"
)
 
func TestDataSurvivesMembershipChanges(t *testing.T) {
	c := New(10_000, 100, nil)
	if _, err := c.Set("x", "1", 0); err == nil {
		t.Fatal("an empty cluster must refuse writes")
	}
	for _, n := range []string{"n1", "n2", "n3"} {
		c.AddNode(n)
	}
	for i := range 3000 {
		k := strconv.Itoa(i)
		c.Set(k, "v"+k, time.Hour)
	}
	if moved := c.AddNode("n4"); moved < 400 || moved > 1200 {
		t.Errorf("adding n4 moved %d keys, want about 750", moved)
	}
	c.RemoveNode("n2")
	total := 0
	for _, n := range c.Stats() {
		total += n
	}
	if total != 3000 {
		t.Fatalf("keys after changes = %d, want 3000", total)
	}
	for i := range 3000 {
		k := strconv.Itoa(i)
		if v, _, ok := c.Get(k); !ok || v != "v"+k {
			t.Fatalf("lost key %s", k)
		}
	}
}
 
func TestTTLPreservedOnMigration(t *testing.T) {
	now := time.Unix(0, 0)
	clock := func() time.Time { return now }
	c := New(100, 50, clock)
	c.AddNode("a")
	c.Set("k", "v", 10*time.Second)
	now = now.Add(6 * time.Second)
	c.AddNode("b")
	c.AddNode("c")
	now = now.Add(5 * time.Second) // 11s total: expired, wherever it lives now
	if _, _, ok := c.Get("k"); ok {
		t.Fatal("migration must keep the remaining TTL, not reset it")
	}
}

package main
 
import (
	_ "embed"
	"encoding/json"
	"io"
	"log/slog"
	"net/http"
	"os"
	"strconv"
	"time"
 
	"kvcache/internal/cluster"
)
 
//go:embed web/index.html
var indexHTML []byte
 
func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}
 
func newServer(c *cluster.Cluster) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /{$}", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write(indexHTML)
	})
	mux.HandleFunc("PUT /api/kv/{key}", func(w http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 1<<20))
		if err != nil {
			writeJSON(w, http.StatusRequestEntityTooLarge, map[string]string{"error": "value too large"})
			return
		}
		ttl, _ := strconv.Atoi(r.URL.Query().Get("ttl"))
		node, err := c.Set(r.PathValue("key"), string(body), time.Duration(ttl)*time.Second)
		if err != nil {
			writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": err.Error()})
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"node": node})
	})
	mux.HandleFunc("GET /api/kv/{key}", func(w http.ResponseWriter, r *http.Request) {
		v, node, ok := c.Get(r.PathValue("key"))
		if !ok {
			writeJSON(w, http.StatusNotFound, map[string]string{"error": "miss", "node": node})
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"value": v, "node": node})
	})
	mux.HandleFunc("DELETE /api/kv/{key}", func(w http.ResponseWriter, r *http.Request) {
		c.Delete(r.PathValue("key"))
		w.WriteHeader(http.StatusNoContent)
	})
	mux.HandleFunc("GET /api/cluster", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, c.Stats())
	})
	mux.HandleFunc("POST /api/nodes/{name}", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, map[string]int{"moved": c.AddNode(r.PathValue("name"))})
	})
	mux.HandleFunc("DELETE /api/nodes/{name}", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, map[string]int{"moved": c.RemoveNode(r.PathValue("name"))})
	})
	return mux
}
 
func main() {
	c := cluster.New(100_000, 150, nil)
	for _, n := range []string{"node-a", "node-b", "node-c"} {
		c.AddNode(n)
	}
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}
	slog.Info("listening", "port", port)
	srv := &http.Server{Addr: ":" + port, Handler: newServer(c), ReadHeaderTimeout: 5 * time.Second}
	if err := srv.ListenAndServe(); err != nil {
		slog.Error("server", "err", err)
		os.Exit(1)
	}
}

package main
 
import (
	"encoding/json"
	"net/http/httptest"
	"strings"
	"testing"
 
	"kvcache/internal/cluster"
)
 
func TestAPI(t *testing.T) {
	c := cluster.New(1000, 50, nil)
	c.AddNode("a")
	c.AddNode("b")
	h := newServer(c)
	call := func(method, path, body string) (int, map[string]any) {
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, httptest.NewRequest(method, path, strings.NewReader(body)))
		var m map[string]any
		json.NewDecoder(rec.Body).Decode(&m)
		return rec.Code, m
	}
	if code, m := call("PUT", "/api/kv/user:1", `{"name":"Ana"}`); code != 200 || m["node"] == "" {
		t.Fatalf("put: %d %v", code, m)
	}
	if code, m := call("GET", "/api/kv/user:1", ""); code != 200 || m["value"] != `{"name":"Ana"}` {
		t.Fatalf("get: %d %v", code, m)
	}
	call("DELETE", "/api/kv/user:1", "")
	if code, _ := call("GET", "/api/kv/user:1", ""); code != 404 {
		t.Fatalf("after delete: %d", code)
	}
	if code, m := call("POST", "/api/nodes/c", ""); code != 200 || m["moved"] == nil {
		t.Fatalf("add node: %d %v", code, m)
	}
	if _, m := call("GET", "/api/cluster", ""); len(m) != 3 {
		t.Fatalf("cluster: %v", m)
	}
}

<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>Cache cluster</title>
  <style>
    body { font-family: system-ui, sans-serif; max-width: 720px; margin: 3rem auto; padding: 0 1rem; }
    .node { display: flex; align-items: center; gap: .6rem; margin: .4rem 0; font: 14px monospace; }
    .bar { height: 18px; background: #076678; border-radius: 4px; transition: width .3s; }
    button, input { padding: .45rem .7rem; border-radius: 8px; border: 1px solid #999; }
    #msg { margin: .8rem 0; color: #555; }
  </style>
</head>
<body>
  <h1>Consistent-hash cache cluster</h1>
  <p><button id="load">Load 300 keys</button>
     <input id="name" value="node-d" size="10"> <button id="add">Add node</button> <button id="rm">Remove node</button></p>
  <div id="msg"></div>
  <div id="nodes"></div>
  <script>
    const $ = (id) => document.getElementById(id);
    async function refresh() {
      const stats = await (await fetch("/api/cluster")).json();
      const max = Math.max(1, ...Object.values(stats));
      $("nodes").innerHTML = "";
      for (const name of Object.keys(stats).sort()) {
        const row = document.createElement("div");
        row.className = "node";
        row.innerHTML = `<span style="width:7rem">${name}</span><span class="bar" style="width:${(stats[name] / max) * 60}%"></span><span>${stats[name]}</span>`;
        $("nodes").append(row);
      }
    }
    $("load").onclick = async () => {
      await Promise.all(Array.from({ length: 300 }, (_, i) =>
        fetch(`/api/kv/key-${i}?ttl=3600`, { method: "PUT", body: "value " + i })));
      $("msg").textContent = "Stored 300 keys.";
      refresh();
    };
    $("add").onclick = async () => {
      const r = await (await fetch("/api/nodes/" + encodeURIComponent($("name").value), { method: "POST" })).json();
      $("msg").textContent = `Added. ${r.moved} keys moved to the new node.`;
      refresh();
    };
    $("rm").onclick = async () => {
      const r = await (await fetch("/api/nodes/" + encodeURIComponent($("name").value), { method: "DELETE" })).json();
      $("msg").textContent = `Removed. ${r.moved} keys moved to the remaining nodes.`;
      refresh();
    };
    refresh();
  </script>
</body>
</html>

go test -race ./...
go run .
curl -s -X PUT 'localhost:8080/api/kv/greeting?ttl=60' -d 'halo'
# {"node":"node-b"}
curl -s localhost:8080/api/kv/greeting
# {"node":"node-b","value":"halo"}