Require only single threaded map, move locking to backend

Author: mgnsk

backend.go

@@ -14,7 +14,7 @@ type recordList[V any] struct {
 type backend[K comparable, V any] struct {
 	timer            *time.Timer
 	done             chan struct{}
-	xmap             builtinMap[K, *record[V]]
+	xmap             backendMap[K, *record[V]]
 	list             recordList[V]
 	earliestExpireAt int64
 	cap              int
@@ -39,9 +39,48 @@ func (b *backend[K, V]) close() error {
 	return nil
 }
 
+func (b *backend[K, V]) LoadOrStore(key K, new *record[V]) (old *record[V], loaded bool) {
+	b.mu.RLock()
+	if r, ok := b.xmap.Load(key); ok {
+		b.mu.RUnlock()
+		return r, true
+	}
+	b.mu.RUnlock()
+
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	if r, ok := b.xmap.Load(key); ok {
+		return r, true
+	}
+
+	b.xmap.Store(key, new)
+
+	return new, false
+}
+
+func (b *backend[K, V]) Range(f func(key K, r *record[V]) bool) {
+	b.mu.RLock()
+	keys := make([]K, 0, b.xmap.Len())
+	b.xmap.Range(func(key K, _ *record[V]) bool {
+		keys = append(keys, key)
+		return true
+	})
+	b.mu.RUnlock()
+
+	for _, key := range keys {
+		b.mu.RLock()
+		r, ok := b.xmap.Load(key)
+		b.mu.RUnlock()
+		if ok && !f(key, r) {
+			return
+		}
+	}
+}
+
 func (b *backend[K, V]) evict(key K) (*record[V], bool) {
-	if r, ok := b.xmap[key]; ok && r.initialized.Load() {
-		delete(b.xmap, key)
+	if r, ok := b.xmap.Load(key); ok && r.initialized.Load() {
+		b.xmap.Delete(key)
 		b.list.Remove(r)
 		return r, true
 	}
@@ -101,16 +140,17 @@ func (b *backend[K, V]) runGC(now int64) {
 
 	var earliest int64
 
-	for key, r := range b.xmap {
+	b.xmap.Range(func(key K, r *record[V]) bool {
 		if r.initialized.Load() {
 			if len(overflowed) > 0 && overflowed[r] || r.deadline > 0 && r.deadline < now {
-				delete(b.xmap, key)
+				b.xmap.Delete(key)
 				b.list.Remove(r)
 			} else if r.deadline > 0 && (earliest == 0 || r.deadline < earliest) {
 				earliest = r.deadline
 			}
 		}
-	}
+		return true
+	})
 
 	b.earliestExpireAt = earliest
 	if earliest > 0 {

cache.go

@@ -32,7 +32,7 @@ func (c *Cache[K, V]) Exists(key K) bool {
 	c.backend.mu.RLock()
 	defer c.backend.mu.RUnlock()
 
-	r, ok := c.backend.xmap[key]
+	r, ok := c.backend.xmap.Load(key)
 	return ok && r.initialized.Load()
 }
 
@@ -41,7 +41,7 @@ func (c *Cache[K, V]) Get(key K) (value V, exists bool) {
 	c.backend.mu.RLock()
 	defer c.backend.mu.RUnlock()
 
-	if r, ok := c.backend.xmap[key]; ok && r.initialized.Load() {
+	if r, ok := c.backend.xmap.Load(key); ok && r.initialized.Load() {
 		return r.value, true
 	}
 
@@ -54,7 +54,7 @@ func (c *Cache[K, V]) Get(key K) (value V, exists bool) {
 //
 // Range is allowed to modify the cache.
 func (c *Cache[K, V]) Range(f func(key K, value V) bool) {
-	c.backend.xmap.Range(&c.backend.mu, func(key K, r *record[V]) bool {
+	c.backend.Range(func(key K, r *record[V]) bool {
 		if r.initialized.Load() {
 			return f(key, r.value)
 		}
@@ -126,7 +126,7 @@ func (c *Cache[K, V]) TryFetch(key K, f func() (V, time.Duration, error)) (value
 	defer new.wg.Done()
 
 loadOrStore:
-	if r, loaded := c.backend.xmap.LoadOrStore(&c.backend.mu, key, new); loaded {
+	if r, loaded := c.backend.LoadOrStore(key, new); loaded {
 		if r.initialized.Load() {
 			c.pool.Put(new)
 			return r.value, nil
@@ -147,7 +147,7 @@ loadOrStore:
 			c.backend.mu.Lock()
 			defer c.backend.mu.Unlock()
 
-			delete(c.backend.xmap, key)
+			c.backend.xmap.Delete(key)
 
 			panic(r)
 		}
@@ -158,7 +158,7 @@ loadOrStore:
 		c.backend.mu.Lock()
 		defer c.backend.mu.Unlock()
 
-		delete(c.backend.xmap, key)
+		c.backend.xmap.Delete(key)
 
 		var zero V
 		return zero, err

map.go

@@ -1,42 +1,35 @@
 package evcache
 
-import "sync"
+type backendMap[K comparable, V any] interface {
+	Load(K) (V, bool)
+	Store(K, V)
+	Delete(K)
+	Len() int
+	Range(func(K, V) bool)
+}
 
 type builtinMap[K comparable, V any] map[K]V
 
-func (m builtinMap[K, V]) LoadOrStore(mu *sync.RWMutex, key K, value V) (actual V, loaded bool) {
-	mu.RLock()
-	if v, ok := m[key]; ok {
-		mu.RUnlock()
-		return v, true
-	}
-	mu.RUnlock()
-
-	mu.Lock()
-	defer mu.Unlock()
-
-	if v, ok := m[key]; ok {
-		return v, true
-	}
+func (m builtinMap[K, V]) Load(key K) (value V, ok bool) {
+	value, ok = m[key]
+	return
+}
 
+func (m builtinMap[K, V]) Store(key K, value V) {
 	m[key] = value
+}
 
-	return value, false
+func (m builtinMap[K, V]) Delete(key K) {
+	delete(m, key)
 }
 
-func (m builtinMap[K, V]) Range(mu *sync.RWMutex, f func(key K, value V) bool) {
-	mu.RLock()
-	keys := make([]K, 0, len(m))
-	for key := range m {
-		keys = append(keys, key)
-	}
-	mu.RUnlock()
+func (m builtinMap[K, V]) Len() int {
+	return len(m)
+}
 
-	for _, key := range keys {
-		mu.RLock()
-		v, ok := m[key]
-		mu.RUnlock()
-		if ok && !f(key, v) {
+func (m builtinMap[K, V]) Range(f func(key K, value V) bool) {
+	for k, v := range m {
+		if !f(k, v) {
 			return
 		}
 	}