learning-go-for-sre-interview/09-safe-counter/main_test.go

package main

import (
	"fmt"
	"sync"
	"testing"
)

const (
	goroutines = 16
	perG       = 1000
)

func TestMutexCounter(t *testing.T) {
	c := NewMutexCounter()
	var wg sync.WaitGroup
	for g := 0; g < goroutines; g++ {
		wg.Add(1)
		go func(g int) {
			defer wg.Done()
			for i := 0; i < perG; i++ {
				c.Inc(fmt.Sprintf("k%d", i%4))
			}
		}(g)
	}
	wg.Wait()
	var total int64
	for i := 0; i < 4; i++ {
		total += c.Get(fmt.Sprintf("k%d", i))
	}
	if want := int64(goroutines * perG); total != want {
		t.Fatalf("lost updates: want %d got %d", want, total)
	}
}

func TestSnapshotIsACopy(t *testing.T) {
	c := NewMutexCounter()
	c.Inc("a")
	snap := c.Snapshot()
	if snap == nil || snap["a"] != 1 {
		t.Fatalf("bad snapshot: %v", snap)
	}
	snap["a"] = 999 // mutating the snapshot must not affect the counter
	if c.Get("a") != 1 {
		t.Fatal("Snapshot returned the internal map, not a copy")
	}
	// And writing to the snapshot while counting must not race (-race checks).
	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); for i := 0; i < 100; i++ { c.Inc("a") } }()
	for i := 0; i < 100; i++ {
		snap["b"] = int64(i)
	}
	wg.Wait()
}

func TestAtomicCounter(t *testing.T) {
	c := NewAtomicCounter()
	var wg sync.WaitGroup
	for g := 0; g < goroutines; g++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < perG; i++ {
				c.Inc()
			}
		}()
	}
	wg.Wait()
	if want := int64(goroutines * perG); c.Get() != want {
		t.Fatalf("lost updates: want %d got %d", want, c.Get())
	}
}

func TestShardedCounter(t *testing.T) {
	c := NewShardedCounter()
	var wg sync.WaitGroup
	for g := 0; g < goroutines; g++ {
		wg.Add(1)
		go func(g int) {
			defer wg.Done()
			for i := 0; i < perG; i++ {
				c.Inc(fmt.Sprintf("key-%d", i%50))
			}
		}(g)
	}
	wg.Wait()
	var total int64
	for k, v := range c.Snapshot() {
		_ = k
		total += v
	}
	if want := int64(goroutines * perG); total != want {
		t.Fatalf("lost updates across shards: want %d got %d", want, total)
	}
	if c.Get("key-0") == 0 {
		t.Fatal("Get by key broken")
	}
}

func BenchmarkMutexCounter(b *testing.B) {
	c := NewMutexCounter()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			c.Inc("hot")
		}
	})
}

func BenchmarkShardedCounter(b *testing.B) {
	c := NewShardedCounter()
	keys := [8]string{"a", "b", "c", "d", "e", "f", "g", "h"}
	b.RunParallel(func(pb *testing.PB) {
		i := 0
		for pb.Next() {
			c.Inc(keys[i&7])
			i++
		}
	})
}

func BenchmarkAtomicCounter(b *testing.B) {
	c := NewAtomicCounter()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			c.Inc()
		}
	})
}