fix panic on concurrent map read

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.

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