File locking module data store

.golangci.yml

@@ -281,4 +281,7 @@ issues:
       path: private/buf/cmd/buf/buf_test.go
       # The LsModules tests call chdir and cannot be parallelized.
       text: "LsModules"
-
+    - linters:
+        - forbidigo
+      # We are using errgroup in a test for concurrent reads/writes to the cache.
+      path: private/bufpkg/bufmodule/bufmodulecache/bufmodulecache_test.go

private/buf/bufcli/cache.go

@@ -29,6 +29,7 @@ import (
 	"github.com/bufbuild/buf/private/bufpkg/bufmodule/bufmodulestore"
 	"github.com/bufbuild/buf/private/pkg/app/appext"
 	"github.com/bufbuild/buf/private/pkg/command"
+	"github.com/bufbuild/buf/private/pkg/filelock"
 	"github.com/bufbuild/buf/private/pkg/normalpath"
 	"github.com/bufbuild/buf/private/pkg/storage/storageos"
 )
@@ -166,12 +167,17 @@ func newModuleDataProvider(
 	if err != nil {
 		return nil, err
 	}
+	filelocker, err := filelock.NewLocker(fullCacheDirPath)
+	if err != nil {
+		return nil, err
+	}
 	return bufmodulecache.NewModuleDataProvider(
 		container.Logger(),
 		delegateModuleDataProvider,
 		bufmodulestore.NewModuleDataStore(
 			container.Logger(),
 			cacheBucket,
+			bufmodulestore.ModuleDataStoreWithFileLocker(filelocker),
 		),
 	), nil
 }

private/bufpkg/bufmodule/bufmodulecache/bufmodulecache_test.go

@@ -16,15 +16,21 @@ package bufmodulecache
 
 import (
 	"context"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/bufbuild/buf/private/bufpkg/bufmodule"
 	"github.com/bufbuild/buf/private/bufpkg/bufmodule/bufmodulestore"
 	"github.com/bufbuild/buf/private/bufpkg/bufmodule/bufmoduletesting"
+	"github.com/bufbuild/buf/private/pkg/filelock"
 	"github.com/bufbuild/buf/private/pkg/slicesext"
 	"github.com/bufbuild/buf/private/pkg/storage/storagemem"
+	"github.com/bufbuild/buf/private/pkg/storage/storageos"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
+	"golang.org/x/sync/errgroup"
 )
 
 func TestCommitProviderForModuleKeyBasic(t *testing.T) {
@@ -214,6 +220,57 @@ func TestModuleDataProviderBasic(t *testing.T) {
 	)
 }
 
+func TestConcurrentCacheReadWrite(t *testing.T) {
+	t.Skip("skipping expensive test for concurrent cache reads/writes")
+	t.Parallel()
+
+	bsrProvider, moduleKeys := testGetBSRProviderAndModuleKeys(t, context.Background())
+	tempDir := t.TempDir()
+	cacheDir := filepath.Join(tempDir, "cache")
+
+	for i := 0; i < 20; i++ {
+		require.NoError(t, os.MkdirAll(cacheDir, 0755))
+		errs, ctx := errgroup.WithContext(context.Background())
+
+		for j := 0; j < 5; j++ {
+			bucket, err := storageos.NewProvider().NewReadWriteBucket(cacheDir)
+			require.NoError(t, err)
+			filelocker, err := filelock.NewLocker(cacheDir)
+			require.NoError(t, err)
+
+			cacheProvider := newModuleDataProvider(
+				zap.NewNop(),
+				bsrProvider,
+				bufmodulestore.NewModuleDataStore(
+					zap.NewNop(),
+					bucket,
+					bufmodulestore.ModuleDataStoreWithFileLocker(filelocker),
+				),
+			)
+
+			errs.Go(func() error {
+				moduleDatas, err := cacheProvider.GetModuleDatasForModuleKeys(
+					ctx,
+					moduleKeys,
+				)
+				if err != nil {
+					return err
+				}
+				for _, moduleData := range moduleDatas {
+					// Calling moduleData.Bucket() checks the digest
+					if _, err := moduleData.Bucket(); err != nil {
+						return err
+					}
+				}
+				return nil
+			})
+		}
+
+		assert.NoError(t, errs.Wait()) // Waits for all go routines to finish and returns the first error, if any
+		require.NoError(t, os.RemoveAll(cacheDir))
+	}
+}
+
 func testGetBSRProviderAndModuleKeys(t *testing.T, ctx context.Context) (bufmoduletesting.OmniProvider, []bufmodule.ModuleKey) {
 	bsrProvider, err := bufmoduletesting.NewOmniProvider(
 		bufmoduletesting.ModuleData{
@@ -235,7 +292,10 @@ func testGetBSRProviderAndModuleKeys(t *testing.T, ctx context.Context) (bufmodu
 		bufmoduletesting.ModuleData{
 			Name: "buf.build/foo/mod3",
 			PathToData: map[string][]byte{
-				"mod3.proto": []byte(
+				"mod3a.proto": []byte(
+					`syntax = proto3; package mod3;`,
+				),
+				"mod3b.proto": []byte(
 					`syntax = proto3; package mod3;`,
 				),
 			},

private/bufpkg/bufmodule/bufmodulecache/module_data_provider.go

@@ -62,6 +62,6 @@ func newModuleDataProvider(
 func (p *moduleDataProvider) GetModuleDatasForModuleKeys(
 	ctx context.Context,
 	moduleKeys []bufmodule.ModuleKey,
-) ([]bufmodule.ModuleData, error) {
+) (_ []bufmodule.ModuleData, retErr error) {
 	return p.baseProvider.getValuesForKeys(ctx, moduleKeys)
 }

private/bufpkg/bufmodule/bufmodulestore/module_data_store.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/bufbuild/buf/private/bufpkg/bufmodule"
 	"github.com/bufbuild/buf/private/pkg/encoding"
+	"github.com/bufbuild/buf/private/pkg/filelock"
 	"github.com/bufbuild/buf/private/pkg/normalpath"
 	"github.com/bufbuild/buf/private/pkg/slicesext"
 	"github.com/bufbuild/buf/private/pkg/storage"
@@ -96,13 +97,27 @@ func ModuleDataStoreWithTar() ModuleDataStoreOption {
 	}
 }
 
+// TODO(doria): clean up GoDoc.
+// ModuleDataStoreWithFileLocker returns a new ModuleDataStoreOption that sets a filelock.Locker.
+//
+// If a filelocker is set, then the module data store will use this to grab a file lock on
+// module.yaml to synchronize reading and writing to the cache.
+//
+// If ModuleDataStoreWithTar is set, then this is ignored.
+func ModuleDataStoreWithFileLocker(filelocker filelock.Locker) ModuleDataStoreOption {
+	return func(moduleDataStore *moduleDataStore) {
+		moduleDataStore.filelocker = filelocker
+	}
+}
+
 /// *** PRIVATE ***
 
 type moduleDataStore struct {
 	logger *zap.Logger
 	bucket storage.ReadWriteBucket
 
-	tar bool
+	tar        bool
+	filelocker filelock.Locker
 }
 
 func newModuleDataStore(
@@ -167,11 +182,26 @@ func (p *moduleDataStore) getModuleDataForModuleKey(
 			return nil, err
 		}
 	} else {
-		moduleCacheBucket, err = p.getReadWriteBucketForDir(moduleKey)
-		// Not checking for fs.ErrNotExist. Function only returns error on actual system error.
+		dirPath, err := getModuleDataStoreDirPath(moduleKey)
 		if err != nil {
 			return nil, err
 		}
+		p.logDebugModuleKey(
+			moduleKey,
+			"module data store dir read write bucket",
+			zap.String("dirPath", dirPath),
+		)
+		moduleCacheBucket = storage.MapReadWriteBucket(p.bucket, storage.MapOnPrefix(dirPath))
+		// Only attempt to get a file lock when storing individual files
+		if p.filelocker != nil {
+			unlocker, err := p.filelocker.RLock(ctx, dirPath+"/"+externalModuleDataFileName)
+			if err != nil {
+				return nil, err
+			}
+			defer func() {
+				retErr = multierr.Append(retErr, unlocker.Unlock())
+			}()
+		}
 	}
 	defer func() {
 		if retErr != nil {
@@ -297,6 +327,15 @@ func (p *moduleDataStore) deleteInvalidModuleData(
 	if err != nil {
 		return err
 	}
+	if p.filelocker != nil {
+		unlocker, err := p.filelocker.Lock(ctx, dirPath+"/"+externalModuleDataFileName)
+		if err != nil {
+			return err
+		}
+		defer func() {
+			retErr = multierr.Append(retErr, unlocker.Unlock())
+		}()
+	}
 	return p.bucket.DeleteAll(ctx, dirPath)
 }
 
@@ -317,10 +356,26 @@ func (p *moduleDataStore) putModuleData(
 			}
 		}()
 	} else {
-		moduleCacheBucket, err = p.getReadWriteBucketForDir(moduleKey)
+		dirPath, err := getModuleDataStoreDirPath(moduleKey)
 		if err != nil {
 			return err
 		}
+		p.logDebugModuleKey(
+			moduleKey,
+			"module data store dir read write bucket",
+			zap.String("dirPath", dirPath),
+		)
+		moduleCacheBucket = storage.MapReadWriteBucket(p.bucket, storage.MapOnPrefix(dirPath))
+		// Only attempt to get a file lock when storing individual files
+		if p.filelocker != nil {
+			unlocker, err := p.filelocker.Lock(ctx, dirPath+"/"+externalModuleDataFileName)
+			if err != nil {
+				return err
+			}
+			defer func() {
+				retErr = multierr.Append(retErr, unlocker.Unlock())
+			}()
+		}
 	}
 	depModuleKeys, err := moduleData.DeclaredDepModuleKeys()
 	if err != nil {
@@ -391,22 +446,6 @@ func (p *moduleDataStore) putModuleData(
 	return storage.PutPath(ctx, moduleCacheBucket, externalModuleDataFileName, data)
 }
 
-// Only returns error on actual system error.
-func (p *moduleDataStore) getReadWriteBucketForDir(
-	moduleKey bufmodule.ModuleKey,
-) (storage.ReadWriteBucket, error) {
-	dirPath, err := getModuleDataStoreDirPath(moduleKey)
-	if err != nil {
-		return nil, err
-	}
-	p.logDebugModuleKey(
-		moduleKey,
-		"module data store dir read write bucket",
-		zap.String("dirPath", dirPath),
-	)
-	return storage.MapReadWriteBucket(p.bucket, storage.MapOnPrefix(dirPath)), nil
-}
-
 // May return fs.ErrNotExist error if tar not found.
 func (p *moduleDataStore) getReadBucketForTar(
 	ctx context.Context,