Golang sharded map

shardedmap.go

package cache

import (
	"sync"
)

type cacheItem interface{}

type Cache interface {
	// Get returns the value associated with the key parameter.
	Get(uint64) (interface{}, bool)
	// Set adds the key-value pair to the Map or updates the value if it's
	// already present.
	Set(uint64, interface{})
	// Del deletes the key-value pair from the Map.
	Del(uint64) (interface{}, bool)

	Clear()
}

const numShards uint64 = 256

type shardedMap struct {
	shards []*lockedMap
}

func NewShardedMap() Cache {
	sm := &shardedMap{
		shards: make([]*lockedMap, int(numShards)),
	}
	for i := range sm.shards {
		sm.shards[i] = newLockedMap()
	}
	return sm
}

func (sm *shardedMap) Get(key uint64) (interface{}, bool) {
	return sm.shards[key%numShards].Get(key)
}

func (sm *shardedMap) Set(key uint64, value interface{}) {
	sm.shards[key%numShards].Set(key, value)
}

func (sm *shardedMap) Del(key uint64) (interface{}, bool) {
	return sm.shards[key%numShards].Del(key)
}

func (sm *shardedMap) Clear() {
	for i := uint64(0); i < numShards; i++ {
		sm.shards[i].Clear()
	}
}

type lockedMap struct {
	sync.RWMutex
	data map[uint64]cacheItem
}

func newLockedMap() *lockedMap {
	return &lockedMap{
		data: make(map[uint64]cacheItem),
	}
}

func (m *lockedMap) Get(key uint64) (interface{}, bool) {
	m.RLock()
	item, ok := m.data[key]
	m.RUnlock()
	if !ok {
		return nil, false
	}
	return item, true
}

func (m *lockedMap) Set(key uint64, value interface{}) {
	m.Lock()
	m.data[key] = value
	m.Unlock()
}

func (m *lockedMap) Del(key uint64) (interface{}, bool) {
	m.Lock()
	item, ok := m.data[key]
	if !ok {
		m.Unlock()
		return nil, false
	}
	delete(m.data, key)
	m.Unlock()
	return item, true
}

func (m *lockedMap) Clear() {
	m.Lock()
	m.data = make(map[uint64]cacheItem)
	m.Unlock()
}