daemon/internal/extensionstorage/bbolt: Allow concurrent reads

Storage reads queued behind every transaction, including unrelated
writes.
Use RWMutex instead.

Signed-off-by: Paweł Gronowski <pawel.gronowski@docker.com>

daemon/internal/extensionstorage/bbolt/storage.go

@@ -44,10 +44,23 @@ func (definition) Declaration() extensions.Declaration {
 }
 
 type storage struct {
-	mu       sync.Mutex
-	dataRoot string
-	db       *bolt.DB
-	closed   bool
+	lifecycle sync.RWMutex
+	mu        sync.Mutex
+	dataRoot  string
+	db        *bolt.DB
+	closed    bool
+}
+
+type databaseUnavailableError struct{ error }
+
+func (databaseUnavailableError) Unavailable() {}
+
+func (e databaseUnavailableError) GRPCStatus() *status.Status {
+	return status.New(codes.Unavailable, e.Error())
+}
+
+func (e databaseUnavailableError) Unwrap() error {
+	return e.error
 }
 
 func (s *storage) init(ctx context.Context, _ extensions.Config, resolver extensions.Resolver) error {
@@ -65,40 +78,65 @@ func (s *storage) init(ctx context.Context, _ extensions.Config, resolver extens
 	if !filepath.IsAbs(response.RootDir) {
 		return fmt.Errorf("daemon config persistent root %q must be absolute", response.RootDir)
 	}
+	s.mu.Lock()
 	s.dataRoot = response.RootDir
+	s.mu.Unlock()
 	return nil
 }
 
 func (s *storage) shutdown(context.Context) error {
+	s.lifecycle.Lock()
+	defer s.lifecycle.Unlock()
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	if s.closed {
+		return nil
+	}
 	s.closed = true
 	if s.db != nil {
-		return s.db.Close()
+		err := s.db.Close()
+		s.db = nil
+		return err
 	}
 	return nil
 }
 
-// transaction also serializes lazy opening and shutdown with active operations.
+// transaction holds the lifecycle read lock through transaction completion so
+// shutdown cannot close the database between capturing it and beginning work.
 func (s *storage) transaction(ctx context.Context, write bool, fn func(*bolt.Tx) error) error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
 	if err := ctx.Err(); err != nil {
 		return err
 	}
-	if s.closed {
-		return errdefs.Unavailable(status.Error(codes.Unavailable, "extension storage is closed"))
-	}
-	if s.dataRoot == "" {
-		return errors.New("extension storage is not initialized")
-	}
-	if s.db == nil {
-		db, err := openDatabase(s.dataRoot)
-		if err != nil {
+	s.lifecycle.RLock()
+	defer s.lifecycle.RUnlock()
+
+	var db *bolt.DB
+	err := func() error {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		if err := ctx.Err(); err != nil {
 			return err
 		}
-		s.db = db
+		if s.closed {
+			return errdefs.Unavailable(status.Error(codes.Unavailable, "extension storage is closed"))
+		}
+		if s.dataRoot == "" {
+			return errors.New("extension storage is not initialized")
+		}
+		if s.db == nil {
+			var err error
+			s.db, err = openDatabase(s.dataRoot)
+			if err != nil {
+				return databaseUnavailableError{err}
+			}
+		}
+		db = s.db
+		return nil
+	}()
+	if err != nil {
+		return err
 	}
+
 	operation := func(tx *bolt.Tx) error {
 		if err := ctx.Err(); err != nil {
 			return err
@@ -109,9 +147,9 @@ func (s *storage) transaction(ctx context.Context, write bool, fn func(*bolt.Tx)
 		return ctx.Err()
 	}
 	if write {
-		return s.db.Update(operation)
+		return db.Update(operation)
 	}
-	return s.db.View(operation)
+	return db.View(operation)
 }
 
 func (s *storage) Get(ctx context.Context, req *storagekv.RecordRequest) (*storagekv.GetResponse, error) {
@@ -121,11 +159,15 @@ func (s *storage) Get(ctx context.Context, req *storagekv.RecordRequest) (*stora
 	resp := &storagekv.GetResponse{}
 	err := s.transaction(ctx, false, func(tx *bolt.Tx) error {
 		bucket := tx.Bucket([]byte(req.Namespace))
-		if bucket == nil || bucket.Get([]byte(req.Key)) == nil {
+		if bucket == nil {
+			return recordNotFound(req.Key)
+		}
+		value := bucket.Get([]byte(req.Key))
+		if value == nil {
 			return recordNotFound(req.Key)
 		}
 		// Bolt values are valid only for the lifetime of the transaction.
-		resp.Value = bytes.Clone(bucket.Get([]byte(req.Key)))
+		resp.Value = bytes.Clone(value)
 		return nil
 	})
 	if err != nil {

daemon/internal/extensionstorage/bbolt/storage_test.go

@@ -11,12 +11,15 @@ import (
 	"strings"
 	"sync"
 	"testing"
+	"time"
 
 	cerrdefs "github.com/containerd/errdefs"
 	"github.com/moby/extensions"
 	daemonconfigv0 "github.com/moby/moby/v2/extpoints/daemonconfig/v0"
 	storagekv "github.com/moby/moby/v2/extpoints/storage/kv/v0"
 	storagepb "github.com/moby/moby/v2/extpoints/storage/kv/v0/protogen"
+	bolt "go.etcd.io/bbolt"
+	bolterrors "go.etcd.io/bbolt/errors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"gotest.tools/v3/assert"
@@ -275,6 +278,175 @@ func TestStorageConcurrentCreate(t *testing.T) {
 	assert.Equal(t, successes, 1)
 }
 
+func TestStorageConcurrentReads(t *testing.T) {
+	t.Parallel()
+	s := newTestStorage(t, t.TempDir()).(*storage)
+	const readers = 2
+	entered := make(chan struct{}, readers)
+	release := make(chan struct{})
+	var releaseOnce sync.Once
+	releaseReaders := func() { releaseOnce.Do(func() { close(release) }) }
+	var wg sync.WaitGroup
+	results := make(chan error, readers)
+	for range readers {
+		wg.Go(func() {
+			results <- s.transaction(t.Context(), false, func(*bolt.Tx) error {
+				entered <- struct{}{}
+				<-release
+				return nil
+			})
+		})
+	}
+	finished := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(finished)
+	}()
+
+	timer := time.NewTimer(2 * time.Second)
+	defer timer.Stop()
+	for range readers {
+		select {
+		case <-entered:
+		case <-timer.C:
+			releaseReaders()
+			select {
+			case <-finished:
+			case <-time.After(2 * time.Second):
+				t.Fatal("read transactions did not finish after release")
+			}
+			t.Fatal("read transactions did not overlap")
+		}
+	}
+	releaseReaders()
+	select {
+	case <-finished:
+	case <-time.After(2 * time.Second):
+		t.Fatal("read transactions did not finish")
+	}
+	for range readers {
+		assert.NilError(t, <-results)
+	}
+}
+
+func TestStorageShutdownWaitsForTransaction(t *testing.T) {
+	t.Parallel()
+	root := t.TempDir()
+	provider := newTestStorage(t, root)
+	s := provider.(*storage)
+	entered := make(chan struct{})
+	release := make(chan struct{})
+	transactionResult := make(chan error, 1)
+	go func() {
+		transactionResult <- s.transaction(t.Context(), true, func(tx *bolt.Tx) error {
+			bucket, err := tx.CreateBucket([]byte("shutdown"))
+			if err != nil {
+				return err
+			}
+			if err := bucket.Put([]byte("key"), []byte("committed")); err != nil {
+				return err
+			}
+			close(entered)
+			<-release
+			return nil
+		})
+	}()
+
+	select {
+	case <-entered:
+	case <-time.After(2 * time.Second):
+		close(release)
+		t.Fatal("transaction did not reach its synchronization point")
+	}
+
+	shutdownResult := make(chan error, 1)
+	go func() {
+		shutdownResult <- s.shutdown(t.Context())
+	}()
+
+	// A pending shutdown writer prevents new readers from acquiring the lock.
+	// Wait for that state instead of assuming the goroutine has been scheduled.
+	deadline := time.NewTimer(2 * time.Second)
+	defer deadline.Stop()
+	ticker := time.NewTicker(time.Millisecond)
+	defer ticker.Stop()
+	for s.lifecycle.TryRLock() {
+		s.lifecycle.RUnlock()
+		select {
+		case err := <-shutdownResult:
+			close(release)
+			t.Fatalf("shutdown returned before the active transaction completed: %v", err)
+		case <-deadline.C:
+			close(release)
+			t.Fatal("shutdown did not wait on the active transaction")
+		case <-ticker.C:
+		}
+	}
+	close(release)
+	select {
+	case err := <-transactionResult:
+		assert.NilError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("transaction did not finish after release")
+	}
+	select {
+	case err := <-shutdownResult:
+		assert.NilError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("shutdown did not finish after the transaction committed")
+	}
+	assert.Check(t, cerrdefs.IsUnavailable(s.Create(t.Context(), &storagekv.WriteRequest{
+		Namespace: "shutdown",
+		Key:       "after",
+	})))
+
+	reopened := testKV(t, newTestStorage(t, root), "shutdown")
+	value, err := reopened.Get(t.Context(), "key")
+	assert.NilError(t, err)
+	assert.Equal(t, string(value), "committed")
+}
+
+func TestStorageOpenFailureUnavailable(t *testing.T) {
+	t.Parallel()
+	root := filepath.Join(t.TempDir(), "missing")
+	provider := newTestStorage(t, root)
+	_, err := provider.Get(t.Context(), &storagekv.RecordRequest{Namespace: "jobs", Key: "key"})
+	assert.Check(t, cerrdefs.IsUnavailable(err), "%v", err)
+	assert.ErrorIs(t, err, os.ErrNotExist)
+
+	err = testKV(t, provider, "jobs").Create(t.Context(), "key", nil)
+	assert.Check(t, cerrdefs.IsUnavailable(err), "%v", err)
+	assert.ErrorIs(t, err, os.ErrNotExist)
+}
+
+func TestStorageOpenContentionUnavailable(t *testing.T) {
+	t.Parallel()
+	root := t.TempDir()
+	first := newTestStorage(t, root)
+	second := newTestStorage(t, root)
+	firstKV := testKV(t, first, "jobs")
+	assert.NilError(t, firstKV.Create(t.Context(), "held", []byte("value")))
+
+	err := second.Create(t.Context(), &storagekv.WriteRequest{
+		Namespace: "jobs",
+		Key:       "provider",
+	})
+	assert.Check(t, cerrdefs.IsUnavailable(err), "%v", err)
+	assert.ErrorIs(t, err, bolterrors.ErrTimeout)
+
+	secondKV := testKV(t, second, "jobs")
+	err = secondKV.Create(t.Context(), "direct-helper", nil)
+	assert.Check(t, cerrdefs.IsUnavailable(err), "%v", err)
+	assert.ErrorIs(t, err, bolterrors.ErrTimeout)
+
+	rpcKV := testKV(t, rpcStorage(t, second), "jobs")
+	err = rpcKV.Create(t.Context(), "rpc-helper", nil)
+	assert.Check(t, cerrdefs.IsUnavailable(err), "%v", err)
+
+	assert.NilError(t, first.(*storage).shutdown(t.Context()))
+	assert.NilError(t, secondKV.Create(t.Context(), "retry", []byte("available")))
+}
+
 func TestStorageValidationBeforeIO(t *testing.T) {
 	t.Parallel()
 	root := t.TempDir()

extpoints/storage/kv/v0/kv.go

@@ -3,6 +3,7 @@ package storagekvv0
 import (
 	"context"
 
+	cerrdefs "github.com/containerd/errdefs"
 	"github.com/containerd/errdefs/pkg/errgrpc"
 	"github.com/moby/extensions"
 	"github.com/moby/moby/v2/errdefs"
@@ -80,8 +81,11 @@ func (kv *KV) List(ctx context.Context, opts ListOptions) (*ListResponse, error)
 // Preserve native errors, including context cancellation and filesystem causes.
 // RPC clients need their status errors translated back to Go error categories.
 func recordError(err error) error {
+	if err == nil || cerrdefs.IsUnavailable(err) {
+		return err
+	}
 	s, ok := status.FromError(err)
-	if !ok || err == nil {
+	if !ok {
 		return err
 	}
 	native := errgrpc.ToNative(err)