feat(repository): ensure we don't run parallel fetches for the same blob IDs (#1835)

* feat(repository): ensure we don't run parallel fetches for the same blob IDs

Fixed a bunch of test flakes.

* fixed race condition, limit the size of mutex cache by using LRU

go.mod

@@ -117,6 +117,7 @@ require (
 	github.com/gobwas/httphead v0.1.0 // indirect
 	github.com/gobwas/pool v0.2.1 // indirect
 	github.com/gobwas/ws v1.1.0 // indirect
+	github.com/hashicorp/golang-lru v0.5.4 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 )

go.sum

@@ -349,6 +349,8 @@ github.com/hashicorp/go-version v1.2.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09
 github.com/hashicorp/go.net v0.0.1/go.mod h1:hjKkEWcCURg++eb33jQU7oqQcI9XDCnUzHA0oac0k90=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
+github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
+github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
 github.com/hashicorp/logutils v1.0.0/go.mod h1:QIAnNjmIWmVIIkWDTG1z5v++HQmx9WQRO+LraFDTW64=
 github.com/hashicorp/mdns v1.0.0/go.mod h1:tL+uN++7HEJ6SQLQ2/p+z2pH24WQKWjBPkE0mNTz8vQ=
 github.com/hashicorp/memberlist v0.1.3/go.mod h1:ajVTdAv/9Im8oMAAj5G31PhhMCZJV2pPBoIllUwCN7I=

internal/cache/content_cache_concurrency_test.go

@@ -0,0 +1,301 @@
+package cache_test
+
+import (
+	"context"
+	"fmt"
+	"runtime/debug"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/kopia/kopia/internal/blobtesting"
+	"github.com/kopia/kopia/internal/cache"
+	"github.com/kopia/kopia/internal/gather"
+	"github.com/kopia/kopia/internal/testlogging"
+	"github.com/kopia/kopia/repo/blob"
+)
+
+type newContentCacheFunc func(ctx context.Context, st blob.Storage, cacheStorage cache.Storage) (cache.ContentCache, error)
+
+func newContentDataCache(ctx context.Context, st blob.Storage, cacheStorage cache.Storage) (cache.ContentCache, error) {
+	return cache.NewContentCacheForData(ctx, st, cacheStorage, cache.SweepSettings{
+		MaxSizeBytes: 100,
+	}, []byte{1, 2, 3, 4})
+}
+
+func newContentMetadataCache(ctx context.Context, st blob.Storage, cacheStorage cache.Storage) (cache.ContentCache, error) {
+	return cache.NewContentCacheForMetadata(ctx, st, cacheStorage, cache.SweepSettings{
+		MaxSizeBytes: 100,
+	})
+}
+
+func TestPrefetchBlocksGetContent_DataCache(t *testing.T) {
+	testContentCachePrefetchBlocksGetContent(t, newContentDataCache)
+}
+
+func TestPrefetchBlocksGetContent_MetadataCache(t *testing.T) {
+	testContentCachePrefetchBlocksGetContent(t, newContentMetadataCache)
+}
+
+func TestGetContentForDifferentContentIDsExecutesInParallel_DataCache(t *testing.T) {
+	testGetContentForDifferentContentIDsExecutesInParallel(t, newContentDataCache, 2)
+}
+
+func TestGetContentForDifferentContentIDsExecutesInParallel_MetadataCache(t *testing.T) {
+	testGetContentForDifferentContentIDsExecutesInParallel(t, newContentMetadataCache, 1)
+}
+
+func TestGetContentForDifferentBlobsExecutesInParallel_DataCache(t *testing.T) {
+	testGetContentForDifferentBlobsExecutesInParallel(t, newContentDataCache)
+}
+
+func TestGetContentForDifferentBlobsExecutesInParallel_MetadataCache(t *testing.T) {
+	testGetContentForDifferentBlobsExecutesInParallel(t, newContentMetadataCache)
+}
+
+func TestGetContentRaceFetchesOnce_DataCache(t *testing.T) {
+	testGetContentRaceFetchesOnce(t, newContentDataCache)
+}
+
+func TestGetContentRaceFetchesOnce_MetadataCache(t *testing.T) {
+	testGetContentRaceFetchesOnce(t, newContentMetadataCache)
+}
+
+// nolint:thelper
+func testContentCachePrefetchBlocksGetContent(t *testing.T, newCache newContentCacheFunc) {
+	ctx := testlogging.Context(t)
+
+	underlyingData := blobtesting.DataMap{}
+	underlying := blobtesting.NewMapStorage(underlyingData, nil, nil)
+	faulty := blobtesting.NewFaultyStorage(underlying)
+
+	cacheData := blobtesting.DataMap{}
+	metadataCacheStorage := blobtesting.NewMapStorage(cacheData, nil, nil).(cache.Storage)
+
+	dataCache, err := newCache(ctx, faulty, metadataCacheStorage)
+	require.NoError(t, err)
+
+	defer dataCache.Close(ctx)
+
+	var tmp gather.WriteBuffer
+	defer tmp.Close()
+
+	require.NoError(t, underlying.PutBlob(ctx, "blob1", gather.FromSlice([]byte{1, 2, 3, 4, 5, 6}), blob.PutOptions{}))
+
+	getBlobStarted := make(chan struct{})
+
+	var (
+		wg sync.WaitGroup
+
+		counter               = new(int32)
+		getBlobFinishedCnt    int32
+		getContentFinishedCnt int32
+	)
+
+	wg.Add(1)
+
+	go func() {
+		var tmp gather.WriteBuffer
+		defer tmp.Close()
+		defer wg.Done()
+
+		<-getBlobStarted // wait until underlying blob starts being fetched
+
+		// start reading content while this is ongoing
+		t.Logf("GetContent started")
+		dataCache.GetContent(ctx, "c1", "blob1", 0, 1, &tmp)
+		t.Logf("GetContent finished")
+
+		getContentFinishedCnt = atomic.AddInt32(counter, 1)
+	}()
+
+	faulty.AddFault(blobtesting.MethodGetBlob).Before(func() {
+		t.Logf("GetBlob started")
+		close(getBlobStarted)
+		time.Sleep(500 * time.Millisecond)
+		t.Logf("GetBlob finished")
+	})
+
+	faulty.AddFault(blobtesting.MethodGetBlob).Before(func() {
+		t.Error("will not be called")
+	})
+
+	dataCache.PrefetchBlob(ctx, "blob1")
+	dataCache.PrefetchBlob(ctx, "blob1")
+
+	wg.Wait()
+
+	// ensure getBlob finishes before getContent finishes despite GetBlob taking non-trivial time
+	require.Less(t, getBlobFinishedCnt, getContentFinishedCnt)
+}
+
+// nolint:thelper
+func testGetContentForDifferentContentIDsExecutesInParallel(t *testing.T, newCache newContentCacheFunc, minGetBlobParallelism int) {
+	ctx := testlogging.Context(t)
+
+	underlyingData := blobtesting.DataMap{}
+	underlying := blobtesting.NewMapStorage(underlyingData, nil, nil)
+	faulty := blobtesting.NewFaultyStorage(underlying)
+
+	cacheData := blobtesting.DataMap{}
+	metadataCacheStorage := blobtesting.NewMapStorage(cacheData, nil, nil).(cache.Storage)
+
+	dataCache, err := newCache(ctx, faulty, metadataCacheStorage)
+	require.NoError(t, err)
+
+	defer dataCache.Close(ctx)
+
+	require.NoError(t, underlying.PutBlob(ctx, "blob1", gather.FromSlice([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10}), blob.PutOptions{}))
+
+	var ct concurrencyTester
+
+	faulty.AddFault(blobtesting.MethodGetBlob).Before(func() {
+		ct.enter()
+		time.Sleep(time.Second)
+		ct.exit()
+	}).Repeat(100)
+
+	var wg sync.WaitGroup
+
+	for i := 0; i < 20; i++ {
+		i := i
+
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+
+			var tmp gather.WriteBuffer
+			defer tmp.Close()
+
+			dataCache.GetContent(ctx, fmt.Sprintf("c%v", i), "blob1", int64(i), 1, &tmp)
+		}()
+	}
+
+	wg.Wait()
+
+	require.GreaterOrEqual(t, ct.maxConcurrencyLevel, minGetBlobParallelism)
+}
+
+// nolint:thelper
+func testGetContentForDifferentBlobsExecutesInParallel(t *testing.T, newCache newContentCacheFunc) {
+	ctx := testlogging.Context(t)
+
+	underlyingData := blobtesting.DataMap{}
+	underlying := blobtesting.NewMapStorage(underlyingData, nil, nil)
+	faulty := blobtesting.NewFaultyStorage(underlying)
+
+	cacheData := blobtesting.DataMap{}
+	metadataCacheStorage := blobtesting.NewMapStorage(cacheData, nil, nil).(cache.Storage)
+
+	dataCache, err := newCache(ctx, faulty, metadataCacheStorage)
+	require.NoError(t, err)
+
+	defer dataCache.Close(ctx)
+
+	for i := 0; i < 100; i++ {
+		require.NoError(t, underlying.PutBlob(ctx, blob.ID(fmt.Sprintf("blob%v", i)), gather.FromSlice([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10}), blob.PutOptions{}))
+	}
+
+	var ct concurrencyTester
+
+	faulty.AddFault(blobtesting.MethodGetBlob).Before(func() {
+		ct.enter()
+		time.Sleep(time.Second)
+		ct.exit()
+	}).Repeat(100)
+
+	var wg sync.WaitGroup
+
+	for i := 0; i < 20; i++ {
+		i := i
+
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+
+			var tmp gather.WriteBuffer
+			defer tmp.Close()
+
+			dataCache.GetContent(ctx, fmt.Sprintf("c%v", i), blob.ID(fmt.Sprintf("blob%v", i)), int64(i), 1, &tmp)
+		}()
+	}
+
+	wg.Wait()
+
+	require.GreaterOrEqual(t, ct.maxConcurrencyLevel, 2)
+}
+
+// nolint:thelper
+func testGetContentRaceFetchesOnce(t *testing.T, newCache newContentCacheFunc) {
+	ctx := testlogging.Context(t)
+
+	underlyingData := blobtesting.DataMap{}
+	underlying := blobtesting.NewMapStorage(underlyingData, nil, nil)
+	faulty := blobtesting.NewFaultyStorage(underlying)
+
+	cacheData := blobtesting.DataMap{}
+	metadataCacheStorage := blobtesting.NewMapStorage(cacheData, nil, nil).(cache.Storage)
+
+	dataCache, err := newCache(ctx, faulty, metadataCacheStorage)
+	require.NoError(t, err)
+
+	defer dataCache.Close(ctx)
+
+	require.NoError(t, underlying.PutBlob(ctx, "blob1", gather.FromSlice([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10}), blob.PutOptions{}))
+
+	faulty.AddFault(blobtesting.MethodGetBlob).Before(func() {
+		t.Logf("stack1: %s", debug.Stack())
+		time.Sleep(time.Second)
+	})
+
+	// this should not execute
+	faulty.AddFault(blobtesting.MethodGetBlob).Before(func() {
+		t.Errorf("GetBlob was called more than once - stack2: %s", debug.Stack())
+	})
+
+	var wg sync.WaitGroup
+
+	for i := 0; i < 20; i++ {
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+
+			var tmp gather.WriteBuffer
+			defer tmp.Close()
+
+			dataCache.GetContent(ctx, "c1", "blob1", 0, 1, &tmp)
+		}()
+	}
+
+	wg.Wait()
+}
+
+type concurrencyTester struct {
+	mu sync.Mutex
+
+	concurrencyLevel    int
+	maxConcurrencyLevel int
+}
+
+func (c *concurrencyTester) enter() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.concurrencyLevel++
+	if c.concurrencyLevel > c.maxConcurrencyLevel {
+		c.maxConcurrencyLevel = c.concurrencyLevel
+	}
+}
+
+func (c *concurrencyTester) exit() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.concurrencyLevel--
+}

internal/cache/content_cache_data.go

@@ -31,6 +31,11 @@ func BlobIDCacheKey(id blob.ID) string {
 }
 
 func (c *contentCacheForData) GetContent(ctx context.Context, contentID string, blobID blob.ID, offset, length int64, output *gather.WriteBuffer) error {
+	// acquire shared lock
+	mut := c.pc.GetFetchingMutex(string(blobID))
+	mut.RLock()
+	defer mut.RUnlock()
+
 	if c.pc.GetPartial(ctx, BlobIDCacheKey(blobID), offset, length, output) {
 		return nil
 	}
@@ -48,7 +53,7 @@ func (c *contentCacheForData) Close(ctx context.Context) {
 	c.pc.Close(ctx)
 }
 
-func (c contentCacheForData) PrefetchBlob(ctx context.Context, blobID blob.ID) error {
+func (c *contentCacheForData) PrefetchBlob(ctx context.Context, blobID blob.ID) error {
 	var blobData gather.WriteBuffer
 	defer blobData.Close()
 
@@ -56,6 +61,15 @@ func (c contentCacheForData) PrefetchBlob(ctx context.Context, blobID blob.ID) e
 		return nil
 	}
 
+	// acquire exclusive lock
+	mut := c.pc.GetFetchingMutex(string(blobID))
+	mut.Lock()
+	defer mut.Unlock()
+
+	if c.pc.GetPartial(ctx, BlobIDCacheKey(blobID), 0, 1, &blobData) {
+		return nil
+	}
+
 	// read the entire blob
 	if err := c.st.GetBlob(ctx, blobID, 0, -1, &blobData); err != nil {
 		reportMissError()
@@ -71,7 +85,7 @@ func (c contentCacheForData) PrefetchBlob(ctx context.Context, blobID blob.ID) e
 	return nil
 }
 
-func (c contentCacheForData) CacheStorage() Storage {
+func (c *contentCacheForData) CacheStorage() Storage {
 	return c.pc.cacheStorage
 }
 

internal/cache/content_cache_metadata.go

@@ -2,9 +2,6 @@ package cache
 
 import (
 	"context"
-	"hash/fnv"
-	"io"
-	"sync"
 
 	"github.com/pkg/errors"
 
@@ -13,24 +10,10 @@ import (
 	"github.com/kopia/kopia/repo/blob"
 )
 
-const (
-	metadataCacheMutexShards = 256
-)
-
 type contentCacheForMetadata struct {
 	pc *PersistentCache
 
-	st             blob.Storage
-	shardedMutexes [metadataCacheMutexShards]sync.Mutex
-}
-
-func (c *contentCacheForMetadata) mutexForBlob(blobID blob.ID) *sync.Mutex {
-	// hash the blob ID to pick one of the sharded mutexes.
-	h := fnv.New32()
-	io.WriteString(h, string(blobID)) //nolint:errcheck
-	mutexID := h.Sum32() % metadataCacheMutexShards
-
-	return &c.shardedMutexes[mutexID]
+	st blob.Storage
 }
 
 func (c *contentCacheForMetadata) GetContent(ctx context.Context, contentID string, blobID blob.ID, offset, length int64, output *gather.WriteBuffer) error {
@@ -39,10 +22,10 @@ func (c *contentCacheForMetadata) GetContent(ctx context.Context, contentID stri
 		return nil
 	}
 
-	// lock the mutex
-	m := c.mutexForBlob(blobID)
-	m.Lock()
-	defer m.Unlock()
+	// acquire exclusive lock
+	mut := c.pc.GetFetchingMutex(string(blobID))
+	mut.Lock()
+	defer mut.Unlock()
 
 	// check again to see if we perhaps lost the race and the data is now in cache.
 	if c.pc.GetPartial(ctx, string(blobID), offset, length, output) {
@@ -77,10 +60,10 @@ func (c *contentCacheForMetadata) PrefetchBlob(ctx context.Context, blobID blob.
 	var blobData gather.WriteBuffer
 	defer blobData.Close()
 
-	// lock the mutex
-	m := c.mutexForBlob(blobID)
-	m.Lock()
-	defer m.Unlock()
+	// acquire exclusive lock
+	mut := c.pc.GetFetchingMutex(string(blobID))
+	mut.Lock()
+	defer mut.Unlock()
 
 	// check to see if the data is now in cache.
 	if c.pc.GetPartial(ctx, string(blobID), 0, 1, &blobData) {

internal/cache/persistent_lru_cache.go

@@ -8,6 +8,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	lru "github.com/hashicorp/golang-lru"
 	"github.com/pkg/errors"
 
 	"github.com/kopia/kopia/internal/clock"
@@ -27,6 +28,11 @@ const (
 	// DefaultTouchThreshold specifies the resolution of timestamps used to determine which cache items
 	// to expire. This helps cache storage writes on frequently accessed items.
 	DefaultTouchThreshold = 10 * time.Minute
+
+	// Size of the mutex cache LRU.
+	// In case a mutex is evicted of the cache, the impact will be some redundant read,
+	// which given the size should be extremely rare.
+	mutexCacheSize = 10000
 )
 
 // PersistentCache provides persistent on-disk cache.
@@ -41,6 +47,8 @@ type PersistentCache struct {
 
 	periodicSweepRunning sync.WaitGroup
 	periodicSweepClosed  chan struct{}
+
+	mutexCache *lru.Cache
 }
 
 // CacheStorage returns cache storage.
@@ -48,6 +56,23 @@ func (c *PersistentCache) CacheStorage() Storage {
 	return c.cacheStorage
 }
 
+// GetFetchingMutex returns a RWMutex used to lock a blob or content during loading.
+func (c *PersistentCache) GetFetchingMutex(key string) *sync.RWMutex {
+	if v, ok := c.mutexCache.Get(key); ok {
+		// nolint:forcetypeassert
+		return v.(*sync.RWMutex)
+	}
+
+	newVal := &sync.RWMutex{}
+
+	if prevVal, ok, _ := c.mutexCache.PeekOrAdd(key, newVal); ok {
+		// nolint:forcetypeassert
+		return prevVal.(*sync.RWMutex)
+	}
+
+	return newVal
+}
+
 // GetOrLoad is utility function gets the provided item from the cache or invokes the provided fetch function.
 // The function also appends and verifies HMAC checksums using provided secret on all cached items to ensure data integrity.
 func (c *PersistentCache) GetOrLoad(ctx context.Context, key string, fetch func(output *gather.WriteBuffer) error, output *gather.WriteBuffer) error {
@@ -62,6 +87,15 @@ func (c *PersistentCache) GetOrLoad(ctx context.Context, key string, fetch func(
 
 	output.Reset()
 
+	mut := c.GetFetchingMutex(key)
+	mut.Lock()
+	defer mut.Unlock()
+
+	// check again while holding the mutex
+	if c.GetFull(ctx, key, output) {
+		return nil
+	}
+
 	if err := fetch(output); err != nil {
 		reportMissError()
 
@@ -245,6 +279,14 @@ func (c *PersistentCache) sweepDirectory(ctx context.Context) (err error) {
 
 	dur := timer.Elapsed()
 
+	const hundredPercent = 100
+
+	inUsePercent := int64(hundredPercent)
+
+	if c.sweep.MaxSizeBytes != 0 {
+		inUsePercent = hundredPercent * totalRetainedSize / c.sweep.MaxSizeBytes
+	}
+
 	log(ctx).Debugw(
 		"finished sweeping",
 		"cache", c.description,
@@ -253,7 +295,7 @@ func (c *PersistentCache) sweepDirectory(ctx context.Context) (err error) {
 		"tooRecentBytes", tooRecentBytes,
 		"tooRecentCount", tooRecentCount,
 		"maxSizeBytes", c.sweep.MaxSizeBytes,
-		"inUsePercent", 100*totalRetainedSize/c.sweep.MaxSizeBytes,
+		"inUsePercent", inUsePercent,
 	)
 
 	return nil
@@ -299,6 +341,8 @@ func NewPersistentCache(ctx context.Context, description string, cacheStorage St
 		storageProtection:   storageProtection,
 	}
 
+	c.mutexCache, _ = lru.New(mutexCacheSize)
+
 	// verify that cache storage is functional by listing from it
 	if _, err := c.cacheStorage.GetMetadata(ctx, "test-blob"); err != nil && !errors.Is(err, blob.ErrBlobNotFound) {
 		return nil, errors.Wrapf(err, "unable to open %v", c.description)

internal/cache/persistent_lru_cache_test.go

@@ -325,6 +325,8 @@ func TestPersistentLRUCache_Defaults(t *testing.T) {
 	})
 	require.NoError(t, err)
 
+	defer pc.Close(ctx)
+
 	pc.Put(ctx, "key1", gather.FromSlice([]byte{1, 2, 3}))
 
 	var tmp gather.WriteBuffer