Add fetch call

Author: bajb

background.go

@@ -35,6 +35,20 @@ func (km *keyMutex) Lock(key string) bool {
 	return true
 }
 
+// Wait blocks until the given key is unlocked by a prior Lock call.
+// It does not acquire the lock; callers typically call Get after waiting
+// or attempt to Lock again to become the next owner.
+func (km *keyMutex) Wait(key string) {
+	km.l.Lock()
+	for {
+		if _, ok := km.s[key]; !ok {
+			km.l.Unlock()
+			return
+		}
+		km.c.Wait()
+	}
+}
+
 func (m CacheMap) BackgroundUpdate(key string, updater func() (interface{}, error)) {
 	// Lock the key from writes
 	locked := backgroundMutex.Lock(key)

fetch.go

@@ -0,0 +1,65 @@
+package ttlmap
+
+import (
+	"errors"
+	"time"
+)
+
+// ErrTypeMismatch is returned when the cached value cannot be cast to the requested generic type.
+var ErrTypeMismatch = errors.New("ttlmap: cached value has different type")
+
+// Fetch returns a strictly typed value from the cache, fetching from the provided source function when missing.
+func Fetch[T any](m CacheMap, key string, source func(string) (T, error)) (T, error) {
+	var zero T
+	var returnValue T
+	var okCast bool
+
+	shard := m.GetShard(key)
+	shard.RLock()
+	itm, ok := shard.items[key]
+	if ok {
+		returnValue, okCast = itm.GetValue().(T)
+		shard.RUnlock()
+		if !okCast {
+			return zero, ErrTypeMismatch
+		}
+
+		if !itm.Expired() {
+			return returnValue, nil
+		}
+
+		if !itm.isUpdating && itm.updateMutex.TryLock() {
+			itm.isUpdating = true
+			value, err := source(key)
+			if err == nil {
+				m.Set(key, value, nil)
+			}
+			itm.updateMutex.Unlock()
+			itm.isUpdating = false
+			return value, nil
+		}
+
+		// Item has expired, but another thread is updateMutex
+		return returnValue, nil
+	}
+	shard.RUnlock()
+	shard.Lock()
+	defer shard.Unlock()
+
+	itm, ok = shard.items[key]
+	if ok {
+		// check the value was not already processed when waiting for the lock
+		returnValue, okCast = itm.GetValue().(T)
+		if !okCast {
+			return zero, ErrTypeMismatch
+		}
+		return returnValue, nil
+	}
+
+	value, err := source(key)
+	if err == nil {
+		itm = newItem(value, m.options.defaultCacheDuration, time.Now().Add(m.options.maxLifetime), nil)
+		shard.items[key] = itm
+	}
+	return value, err
+}

fetch_test.go

@@ -0,0 +1,113 @@
+package ttlmap_test
+
+import (
+	"log"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/packaged/ttlmap"
+)
+
+func TestFetch_MissSingleflight(t *testing.T) {
+	cache := ttlmap.New(ttlmap.WithDefaultTTL(100 * time.Millisecond))
+
+	var calls int32
+	source := func(key string) (int, error) {
+		time.Sleep(10 * time.Millisecond)
+		atomic.AddInt32(&calls, 1)
+		return 123, nil
+	}
+
+	// Launch many concurrent callers for a missing key
+	const n = 32
+	var wg sync.WaitGroup
+	wg.Add(n)
+	results := make([]int, n)
+	for i := 0; i < n; i++ {
+		go func(ix int) {
+			defer wg.Done()
+			v, err := ttlmap.Fetch[int](cache, "k1", source)
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+				return
+			}
+			results[ix] = v
+		}(i)
+	}
+	wg.Wait()
+
+	if atomic.LoadInt32(&calls) != 1 {
+		t.Fatalf("expected single source call, got %d", calls)
+	}
+	for i := 0; i < n; i++ {
+		if results[i] != 123 {
+			t.Fatalf("unexpected result at %d: %d", i, results[i])
+		}
+	}
+}
+
+func TestFetch_TypeMismatch(t *testing.T) {
+	cache := ttlmap.New(ttlmap.WithDefaultTTL(100 * time.Millisecond))
+	cache.Set("k", "hello", nil)
+
+	var called int32
+	src := func(key string) (int, error) {
+		atomic.AddInt32(&called, 1)
+		return 42, nil
+	}
+	_, err := ttlmap.Fetch[int](cache, "k", src)
+	if err == nil {
+		t.Fatalf("expected ErrTypeMismatch, got nil")
+	}
+	if err != ttlmap.ErrTypeMismatch {
+		t.Fatalf("expected ErrTypeMismatch, got %v", err)
+	}
+	if atomic.LoadInt32(&called) != 0 {
+		t.Fatalf("source should not be called on type mismatch hit")
+	}
+}
+
+func TestFetch_StaleWhileRevalidate(t *testing.T) {
+	// Short TTL to trigger expiry
+	cache := ttlmap.New(ttlmap.WithDefaultTTL(time.Second), ttlmap.WithMaxLifetime(2*time.Second))
+	ttl := 20 * time.Millisecond
+	cache.Set("sk", 1, &ttl)
+
+	src := func(key string) (int, error) {
+		time.Sleep(15 * time.Millisecond)
+		return 2, nil
+	}
+	// First call should return stale value 1 and trigger background refresh
+	v1, err := ttlmap.Fetch[int](cache, "sk", src)
+	if err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+	if v1 != 1 {
+		t.Fatalf("expected stale value 1, got %d", v1)
+	}
+
+	wg := sync.WaitGroup{}
+	for i := 0; i < 100000; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_, fetchErr := ttlmap.Fetch[int](cache, "sk", src)
+			if fetchErr != nil {
+				log.Print(fetchErr.Error())
+			}
+		}()
+	}
+	wg.Wait()
+
+	time.Sleep(10 * time.Millisecond)
+	// Next call should observe refreshed value 2
+	v2, err := ttlmap.Fetch[int](cache, "sk", src)
+	if err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+	if v2 != 2 {
+		t.Fatalf("expected refreshed value 2, got %d", v2)
+	}
+}

go.mod

@@ -1,3 +1,3 @@
 module github.com/packaged/ttlmap
 
-go 1.15
+go 1.20

item.go

@@ -8,11 +8,13 @@ import (
 // Item represents a record in the map
 type Item struct {
 	sync.RWMutex
-	data     interface{}
-	deadline time.Time
-	ttl      time.Duration
-	expires  *time.Time
-	onDelete func(*Item)
+	updateMutex sync.RWMutex
+	isUpdating  bool
+	data        interface{}
+	deadline    time.Time
+	ttl         time.Duration
+	expires     *time.Time
+	onDelete    func(*Item)
 }
 
 func newItem(value interface{}, duration time.Duration, deadline time.Time, onDelete func(*Item)) *Item {