Sharded Concurrent Map

package main

import (
	"fmt"
	"hash/fnv"
	"sync"
)

type shard[K comparable, V any] struct {
	mu   sync.RWMutex
	data map[K]V
}

type ShardedMap[K comparable, V any] struct {
	shards []*shard[K, V]
}

func NewShardedMap[K comparable, V any](count int) *ShardedMap[K, V] {
	if count <= 0 {
		panic("shard count must be positive")
	}
	m := &ShardedMap[K, V]{shards: make([]*shard[K, V], count)}
	for i := range m.shards {
		m.shards[i] = &shard[K, V]{data: make(map[K]V)}
	}
	return m
}

func (m *ShardedMap[K, V]) shardFor(key K) *shard[K, V] {
	hasher := fnv.New32a()
	fmt.Fprintf(hasher, "%v", key)
	return m.shards[int(hasher.Sum32())%len(m.shards)]
}

func (m *ShardedMap[K, V]) Set(key K, value V) {
	s := m.shardFor(key)
	s.mu.Lock()
	s.data[key] = value
	s.mu.Unlock()
}

func (m *ShardedMap[K, V]) Get(key K) (V, bool) {
	s := m.shardFor(key)
	s.mu.RLock()
	value, ok := s.data[key]
	s.mu.RUnlock()
	return value, ok
}

func (m *ShardedMap[K, V]) Delete(key K) {
	s := m.shardFor(key)
	s.mu.Lock()
	delete(s.data, key)
	s.mu.Unlock()
}

func (m *ShardedMap[K, V]) Len() int {
	total := 0
	for _, s := range m.shards {
		s.mu.RLock()
		total += len(s.data)
		s.mu.RUnlock()
	}
	return total
}

func main() {
	cache := NewShardedMap[string, int](8)
	var wg sync.WaitGroup
	for worker := 0; worker < 16; worker++ {
		wg.Add(1)
		go func(worker int) {
			defer wg.Done()
			for i := 0; i < 100; i++ {
				cache.Set(fmt.Sprintf("key-%d-%d", worker, i), i)
			}
		}(worker)
	}
	wg.Wait()

	fmt.Println("entries:", cache.Len())
	value, ok := cache.Get("key-3-42")
	fmt.Println("lookup:", value, ok)
	cache.Delete("key-3-42")
	_, ok = cache.Get("key-3-42")
	fmt.Println("after delete:", ok)
}