Concurrent caches, singleflight and pub/sub

Go · sf/stampede/main.go

// stampede/main.go: 100 requests for one product arrive together while its cache entry is cold. Without help, each
// one misses the cache and loads the product from the database. With singleflight.Group.Do, the first request loads
// it and the others wait for that load and get its result.
package main

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/sync/singleflight"
)

var dbLoads atomic.Int64

// loadProduct stands for a database query of 50 ms.
func loadProduct(id string) (string, error) {
	dbLoads.Add(1)
	time.Sleep(50 * time.Millisecond)
	return id + ": steel bottle, 1 litre", nil
}

// cache is a map behind a mutex, enough for this program.
type cache struct {
	mu sync.Mutex
	m  map[string]string
}

func (c *cache) get(id string) (string, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	v, ok := c.m[id]
	return v, ok
}

func (c *cache) set(id, v string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.m[id] = v
}

// burst sends 100 requests for the same product at once and returns how many database loads they caused.
func burst(get func(id string) (string, error)) int64 {
	dbLoads.Store(0)
	var wg sync.WaitGroup
	for range 100 {
		wg.Go(func() { get("P-42") })
	}
	wg.Wait()
	return dbLoads.Load()
}

func main() {
	plain := &cache{m: map[string]string{}}
	n := burst(func(id string) (string, error) {
		if v, ok := plain.get(id); ok {
			return v, nil
		}
		v, err := loadProduct(id)
		if err == nil {
			plain.set(id, v)
		}
		return v, err
	})
	fmt.Printf("without singleflight: 100 requests, %d database loads\n", n)

	var g singleflight.Group
	collapsed := &cache{m: map[string]string{}}
	n = burst(func(id string) (string, error) {
		if v, ok := collapsed.get(id); ok {
			return v, nil
		}
		v, err, _ := g.Do(id, func() (any, error) {
			if v, ok := collapsed.get(id); ok { // a late caller may come after the load has finished
				return v, nil
			}
			v, err := loadProduct(id)
			if err == nil {
				collapsed.set(id, v)
			}
			return v, err
		})
		if err != nil {
			return "", err
		}
		return v.(string), nil
	})
	fmt.Printf("with singleflight:    100 requests, %d database load\n", n)
}

MOD · sf/go.mod

module example.com/sf

go 1.27

require golang.org/x/sync v0.23.0

Go (broker.go) · pubsub/broker.go

// Package pubsub is a small broker. Each subscriber has a buffered channel of its own, and Publish never waits for a
// slow subscriber: when a subscriber's buffer is full, the oldest event in it is dropped to make room for the new one.
package pubsub

import "sync"

// Broker delivers every published event to every subscriber.
type Broker struct {
	mu   sync.Mutex
	subs []chan string // in the order they subscribed
}

// Subscribe returns a channel of the events published from now on, with room for buffer of them, and a function that
// ends the subscription and closes the channel.
func (b *Broker) Subscribe(buffer int) (<-chan string, func()) {
	ch := make(chan string, buffer)
	b.mu.Lock()
	b.subs = append(b.subs, ch)
	b.mu.Unlock()
	var once sync.Once
	return ch, func() {
		once.Do(func() {
			b.mu.Lock()
			defer b.mu.Unlock()
			for i, s := range b.subs {
				if s == ch {
					b.subs = append(b.subs[:i], b.subs[i+1:]...)
					break
				}
			}
			close(ch) // under the lock, so no Publish can send on it afterwards
		})
	}
}

// Publish delivers ev to every subscriber without waiting, dropping the oldest event of a subscriber whose buffer is
// full.
func (b *Broker) Publish(ev string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for _, ch := range b.subs {
		dropOldest(ch, ev)
	}
}

// dropOldest sends ev on ch, first taking the oldest event out of ch for as long as its buffer is full.
func dropOldest(ch chan string, ev string) {
	for {
		select {
		case ch <- ev:
			return
		default:
		}
		select {
		case <-ch: // the buffer is full: drop its oldest event
		default:
		}
	}
}

Go (broker_test.go) · pubsub/broker_test.go

package pubsub

import (
	"fmt"
	"strings"
	"sync"
	"testing"
	"testing/synctest"
	"time"
)

// TestTwoSubscribers publishes six events, one every 10 ms of the bubble's clock, to a subscriber that reads each at
// once and to one that needs 35 ms for each, both with room for 2 events.
func TestTwoSubscribers(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		var b Broker
		fast, stopFast := b.Subscribe(2)
		slow, stopSlow := b.Subscribe(2)
		var gotFast, gotSlow []string
		var wg sync.WaitGroup
		wg.Go(func() {
			for ev := range fast {
				gotFast = append(gotFast, ev)
			}
		})
		wg.Go(func() {
			for ev := range slow {
				gotSlow = append(gotSlow, ev)
				time.Sleep(35 * time.Millisecond)
			}
		})
		var waited time.Duration
		for i := 1; i <= 6; i++ {
			time.Sleep(10 * time.Millisecond)
			t0 := time.Now()
			b.Publish(fmt.Sprintf("e%d", i))
			waited += time.Since(t0)
		}
		time.Sleep(200 * time.Millisecond) // the slow subscriber finishes what it holds
		stopFast()
		stopSlow()
		wg.Wait()
		fmt.Println("fast subscriber got:", strings.Join(gotFast, " "))
		fmt.Println("slow subscriber got:", strings.Join(gotSlow, " "))
		fmt.Println("Publish waited for subscribers:", waited)
	})
}

Go (cache.go) · practice/cache.go

// Package practice is the practice of the lesson on caches: a generic cache that loads a missing key once, however
// many goroutines ask for it at the same time.
package practice

import "time"

// Cache keeps each value for ttl after it was loaded. Get loads a key that is missing or has expired with load, once
// for all the callers that ask for that key while it loads, and never caches an error.
type Cache[K comparable, V any] struct {
	ttl  time.Duration
	load func(key K) (V, error)
}

// New returns a cache that loads values with load and keeps them for ttl.
func New[K comparable, V any](ttl time.Duration, load func(key K) (V, error)) *Cache[K, V] {
	return &Cache[K, V]{ttl: ttl, load: load}
}

// Get returns the value of key, loading it if it is missing or has expired. Callers that ask for a key while it loads
// wait for that load and get its result; loads of different keys run at the same time.
func (c *Cache[K, V]) Get(key K) (V, error) {
	var zero V
	return zero, nil // replace this with a map behind a mutex and one load per key
}

// Close stops any goroutine the cache started and waits for it.
func (c *Cache[K, V]) Close() {}

Go (cache_test.go) · practice/cache_test.go

package practice

import (
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"testing/synctest"
	"time"
)

// TestOneLoadFor200Callers: 200 goroutines ask for the same cold key at once, and the load takes 50 ms. Exactly one
// load may happen, and every caller gets its value. It runs on the real clock: a goroutine waiting for a sync.Mutex
// would stop a synctest bubble's clock, and a cache that loads under its lock would then hang instead of failing.
func TestOneLoadFor200Callers(t *testing.T) {
	var loads atomic.Int64
	c := New(time.Minute, func(key string) (string, error) {
		loads.Add(1)
		time.Sleep(50 * time.Millisecond)
		return "value of " + key, nil
	})
	defer c.Close()
	var wg sync.WaitGroup
	var right atomic.Int64
	for range 200 {
		wg.Go(func() {
			if v, err := c.Get("P-42"); err == nil && v == "value of P-42" {
				right.Add(1)
			}
		})
	}
	wg.Wait()
	if loads.Load() != 1 || right.Load() != 200 {
		t.Fatalf("%d loads, %d callers got the value; want 1 load and 200", loads.Load(), right.Load())
	}
}

// TestReloadAfterTTL: a value is kept for its ttl of 100 ms, then loaded again.
func TestReloadAfterTTL(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		var loads atomic.Int64
		c := New(100*time.Millisecond, func(key string) (int, error) {
			return int(loads.Add(1)), nil
		})
		defer c.Close()
		a, _ := c.Get("k")
		time.Sleep(50 * time.Millisecond)
		b, _ := c.Get("k")
		time.Sleep(100 * time.Millisecond)
		d, _ := c.Get("k")
		if a != 1 || b != 1 || d != 2 {
			t.Fatalf("Get gave %d, %d, %d at 0, 50 and 150 ms; want 1, 1, 2", a, b, d)
		}
	})
}

// TestKeysLoadInParallel: three keys whose loads take 50 ms each must all be ready well before 150 ms, so a load
// must not run while the cache's lock is held. It runs on the real clock, for the reason given above.
func TestKeysLoadInParallel(t *testing.T) {
	c := New(time.Minute, func(key string) (string, error) {
		time.Sleep(50 * time.Millisecond)
		return key, nil
	})
	defer c.Close()
	keys := []string{"a", "b", "c"}
	got := make([]string, len(keys))
	start := time.Now()
	var wg sync.WaitGroup
	for i, k := range keys {
		wg.Go(func() { got[i], _ = c.Get(k) })
	}
	wg.Wait()
	took := time.Since(start)
	if got[0] != "a" || got[1] != "b" || got[2] != "c" {
		t.Fatalf("got %q; want [a b c], loaded at the same time", got)
	}
	if took >= 120*time.Millisecond {
		t.Fatalf("three keys loaded in %v; want well under 150ms, about the 50ms of one load", took.Round(time.Millisecond))
	}
}

// TestErrorNotCached: a failed load returns its error, and the next Get loads again.
func TestErrorNotCached(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		broken := errors.New("database: connection refused")
		var loads atomic.Int64
		c := New(time.Minute, func(key string) (string, error) {
			if loads.Add(1) == 1 {
				return "", broken
			}
			return "ok", nil
		})
		defer c.Close()
		_, err1 := c.Get("k")
		v, err2 := c.Get("k")
		if !errors.Is(err1, broken) || err2 != nil || v != "ok" || loads.Load() != 2 {
			t.Fatalf("first Get: %v; second: %q, %v; %d loads; want the error, then \"ok\" after 2 loads", err1, v, err2, loads.Load())
		}
	})
}