diff --git a/cache.go b/cache.go index 12f1e88..47799ef 100644 --- a/cache.go +++ b/cache.go @@ -23,7 +23,7 @@ type UpdateFunc[T any] func(ctx context.Context, key string) (EntryUpdate[T], er // entry is missing or expired, it will be updated within the Get method from the UpdateFunc. type UpdateCache[T any] struct { keyLocks *keyedMutex - entriesLock *sync.Mutex + entriesLock sync.RWMutex ttl time.Duration entries map[string]*Entry[T] updateFunc UpdateFunc[T] @@ -35,11 +35,10 @@ type UpdateCache[T any] struct { // The UpdateCache is thread-safe and can be called from multiple goroutines. func NewUpdateCache[T any](ttl time.Duration, updateFunc UpdateFunc[T]) *UpdateCache[T] { return &UpdateCache[T]{ - keyLocks: newKeyedMutex(), - entriesLock: &sync.Mutex{}, - ttl: ttl, - entries: make(map[string]*Entry[T]), - updateFunc: updateFunc, + keyLocks: newKeyedMutex(), + ttl: ttl, + entries: make(map[string]*Entry[T]), + updateFunc: updateFunc, } } @@ -54,32 +53,36 @@ func (c *UpdateCache[T]) Get(ctx context.Context, key string) (res T, hit bool, unlock := c.keyLocks.Lock(key) defer unlock() - if entry, exists := c.entries[key]; exists && entry.ExpiresAt.After(time.Now()) { + c.entriesLock.RLock() + entry, exists := c.entries[key] + c.entriesLock.RUnlock() + + if exists && entry.ExpiresAt.After(time.Now()) { return entry.Value, true, entry.Error - } else { - entry, err := c.updateFunc(ctx, key) + } - // In case we receive an *UpdateFunc* error here, we won't store the result and just return the error here. - // In this case, we want to retry loading the entry again on the next Get call. - if err != nil { - return res, false, err - } + update, err := c.updateFunc(ctx, key) - // acquire a write lock on the cache to update the entry - c.entriesLock.Lock() - defer c.entriesLock.Unlock() - - // In case the error is embedded within the EntryUpdate, we still return it as the error below, but also update - // the cache. This allows us to cache persistent errors and reduce the load on any underlying datasource by not - // calling the UpdateFunc again until the cache entry expires. - c.entries[key] = &Entry[T]{ - Value: entry.Value, - Error: entry.Error, - ExpiresAt: time.Now().Add(c.ttl), - } + // In case we receive an *UpdateFunc* error here, we won't store the result and just return the error here. + // In this case, we want to retry loading the entry again on the next Get call. + if err != nil { + return res, false, err + } - return entry.Value, false, entry.Error + // acquire a write lock on the cache to update the entry + c.entriesLock.Lock() + defer c.entriesLock.Unlock() + + // In case the error is embedded within the EntryUpdate, we still return it as the error below, but also update + // the cache. This allows us to cache persistent errors and reduce the load on any underlying datasource by not + // calling the UpdateFunc again until the cache entry expires. + c.entries[key] = &Entry[T]{ + Value: update.Value, + Error: update.Error, + ExpiresAt: time.Now().Add(c.ttl), } + + return update.Value, false, update.Error } // Prune removes all expired entries from the cache. diff --git a/cache_test.go b/cache_test.go index a92d2ed..5c23b84 100644 --- a/cache_test.go +++ b/cache_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "github.com/stretchr/testify/assert" + "strconv" "sync" "testing" "time" @@ -65,6 +66,31 @@ func TestCache_GetParallel(t *testing.T) { assert.Equal(t, 1, updateCnt) } +func TestCache_GetParallelDistinctKeys(t *testing.T) { + + // Each key has its own keyed mutex, so concurrent Get calls on different keys read and write the shared entries + // map at the same time. This reproduces the "concurrent map read and map write" fatal error when run with -race. + testCache := NewUpdateCache(5*time.Minute, func(ctx context.Context, key string) (EntryUpdate[string], error) { + return EntryUpdate[string]{Value: key}, nil + }) + + wg := &sync.WaitGroup{} + for i := 0; i < 100; i++ { + key := strconv.Itoa(i) + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + res, _, err := testCache.Get(context.Background(), key) + assert.NoError(t, err) + assert.Equal(t, key, res) + } + }() + } + + wg.Wait() +} + func TestCache_GetExpires(t *testing.T) { updateCnt := 0