Merge pull request moby#5190 from tonistiigi/db-transactions

db: move DB dependencies to transactor interface

cache/metadata/metadata.go

@@ -8,6 +8,8 @@ import (
 	"sync"
 
 	"github.com/moby/buildkit/util/bklog"
+	"github.com/moby/buildkit/util/db"
+	"github.com/moby/buildkit/util/db/boltutil"
 	"github.com/pkg/errors"
 	bolt "go.etcd.io/bbolt"
 )
@@ -21,18 +23,18 @@ const (
 var errNotFound = errors.Errorf("not found")
 
 type Store struct {
-	db *bolt.DB
+	db db.DB
 }
 
 func NewStore(dbPath string) (*Store, error) {
-	db, err := bolt.Open(dbPath, 0600, nil)
+	db, err := boltutil.Open(dbPath, 0600, nil)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to open database file %s", dbPath)
 	}
 	return &Store{db: db}, nil
 }
 
-func (s *Store) DB() *bolt.DB {
+func (s *Store) DB() db.Transactor {
 	return s.db
 }
 
@@ -183,21 +185,28 @@ func (s *Store) Get(id string) (*StorageItem, bool) {
 		si, _ := newStorageItem(id, nil, s)
 		return si
 	}
-	tx, err := s.db.Begin(false)
-	if err != nil {
-		return empty(), false
-	}
-	defer tx.Rollback()
-	b := tx.Bucket([]byte(mainBucket))
-	if b == nil {
+
+	var si *StorageItem
+	if err := s.db.Update(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(mainBucket))
+		if b == nil {
+			return nil
+		}
+		b = b.Bucket([]byte(id))
+		if b == nil {
+			return nil
+		}
+		si, _ = newStorageItem(id, b, s)
+		return nil
+	}); err != nil {
 		return empty(), false
 	}
-	b = b.Bucket([]byte(id))
-	if b == nil {
-		return empty(), false
+
+	if si != nil {
+		return si, true
 	}
-	si, _ := newStorageItem(id, b, s)
-	return si, true
+
+	return empty(), false
 }
 
 func (s *Store) Close() error {

cmd/buildkitd/main.go

@@ -46,6 +46,7 @@ import (
 	"github.com/moby/buildkit/util/appdefaults"
 	"github.com/moby/buildkit/util/archutil"
 	"github.com/moby/buildkit/util/bklog"
+	"github.com/moby/buildkit/util/db/boltutil"
 	"github.com/moby/buildkit/util/grpcerrors"
 	"github.com/moby/buildkit/util/profiler"
 	"github.com/moby/buildkit/util/resolver"
@@ -62,7 +63,6 @@ import (
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli"
-	"go.etcd.io/bbolt"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"go.opentelemetry.io/otel/exporters/prometheus"
 	"go.opentelemetry.io/otel/propagation"
@@ -804,7 +804,7 @@ func newController(c *cli.Context, cfg *config.Config) (*control.Controller, err
 		return nil, err
 	}
 
-	historyDB, err := bbolt.Open(filepath.Join(cfg.Root, "history.db"), 0600, nil)
+	historyDB, err := boltutil.Open(filepath.Join(cfg.Root, "history.db"), 0600, nil)
 	if err != nil {
 		return nil, err
 	}

control/control.go

@@ -35,6 +35,7 @@ import (
 	"github.com/moby/buildkit/solver/llbsolver/proc"
 	"github.com/moby/buildkit/solver/pb"
 	"github.com/moby/buildkit/util/bklog"
+	"github.com/moby/buildkit/util/db"
 	"github.com/moby/buildkit/util/imageutil"
 	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/throttle"
@@ -43,7 +44,6 @@ import (
 	"github.com/moby/buildkit/worker"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
-	"go.etcd.io/bbolt"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	tracev1 "go.opentelemetry.io/proto/otlp/collector/trace/v1"
 	"golang.org/x/sync/errgroup"
@@ -62,7 +62,7 @@ type Opt struct {
 	ResolveCacheImporterFuncs map[string]remotecache.ResolveCacheImporterFunc
 	Entitlements              []string
 	TraceCollector            sdktrace.SpanExporter
-	HistoryDB                 *bbolt.DB
+	HistoryDB                 db.DB
 	CacheStore                *bboltcachestorage.Store
 	LeaseManager              *leaseutil.Manager
 	ContentStore              *containerdsnapshot.Store

solver/bboltcachestorage/storage.go

@@ -9,6 +9,8 @@ import (
 	"github.com/moby/buildkit/identity"
 	"github.com/moby/buildkit/solver"
 	"github.com/moby/buildkit/util/bklog"
+	"github.com/moby/buildkit/util/db"
+	"github.com/moby/buildkit/util/db/boltutil"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
 	bolt "go.etcd.io/bbolt"
@@ -22,11 +24,13 @@ const (
 )
 
 type Store struct {
-	db *bolt.DB
+	db db.DB
 }
 
 func NewStore(dbPath string) (*Store, error) {
-	db, err := safeOpenDB(dbPath)
+	db, err := safeOpenDB(dbPath, &bolt.Options{
+		NoSync: true,
+	})
 	if err != nil {
 		return nil, err
 	}
@@ -42,7 +46,6 @@ func NewStore(dbPath string) (*Store, error) {
 	}); err != nil {
 		return nil, err
 	}
-	db.NoSync = true
 	return &Store{db: db}, nil
 }
 
@@ -465,7 +468,7 @@ func isEmptyBucket(b *bolt.Bucket) bool {
 
 // safeOpenDB opens a bolt database and recovers from panic that
 // can be caused by a corrupted database file.
-func safeOpenDB(dbPath string) (db *bolt.DB, err error) {
+func safeOpenDB(dbPath string, opts *bolt.Options) (db db.DB, err error) {
 	defer func() {
 		if r := recover(); r != nil {
 			err = errors.Errorf("%v", r)
@@ -476,16 +479,16 @@ func safeOpenDB(dbPath string) (db *bolt.DB, err error) {
 		// then fallback to resetting the database since the database
 		// may be corrupt.
 		if err != nil && fileHasContent(dbPath) {
-			db, err = fallbackOpenDB(dbPath, err)
+			db, err = fallbackOpenDB(dbPath, opts, err)
 		}
 	}()
-	return openDB(dbPath)
+	return openDB(dbPath, opts)
 }
 
 // fallbackOpenDB performs database recovery and opens the new database
 // file when the database fails to open. Called after the first database
 // open fails.
-func fallbackOpenDB(dbPath string, openErr error) (*bolt.DB, error) {
+func fallbackOpenDB(dbPath string, opts *bolt.Options, openErr error) (db.DB, error) {
 	backupPath := dbPath + "." + identity.NewID() + ".bak"
 	bklog.L.Errorf("failed to open database file %s, resetting to empty. Old database is backed up to %s. "+
 		"This error signifies that buildkitd likely crashed or was sigkilled abrubtly, leaving the database corrupted. "+
@@ -496,12 +499,12 @@ func fallbackOpenDB(dbPath string, openErr error) (*bolt.DB, error) {
 
 	// Attempt to open the database again. This should be a new database.
 	// If this fails, it is a permanent error.
-	return openDB(dbPath)
+	return openDB(dbPath, opts)
 }
 
 // openDB opens a bolt database in user-only read/write mode.
-func openDB(dbPath string) (*bolt.DB, error) {
-	return bolt.Open(dbPath, 0600, nil)
+func openDB(dbPath string, opts *bolt.Options) (db.DB, error) {
+	return boltutil.Open(dbPath, 0600, opts)
 }
 
 // fileHasContent checks if we have access to the file with appropriate

solver/bboltcachestorage/storage_test.go

@@ -16,7 +16,7 @@ func TestBoltCacheStorage(t *testing.T) {
 		st, err := NewStore(filepath.Join(tmpDir, "cache.db"))
 		require.NoError(t, err)
 		t.Cleanup(func() {
-			require.NoError(t, st.db.Close())
+			require.NoError(t, st.Close())
 		})
 
 		return st

solver/llbsolver/history.go

@@ -22,6 +22,7 @@ import (
 	"github.com/moby/buildkit/cmd/buildkitd/config"
 	"github.com/moby/buildkit/identity"
 	containerdsnapshot "github.com/moby/buildkit/snapshot/containerd"
+	"github.com/moby/buildkit/util/db"
 	"github.com/moby/buildkit/util/grpcerrors"
 	"github.com/moby/buildkit/util/iohelper"
 	"github.com/moby/buildkit/util/leaseutil"
@@ -39,7 +40,7 @@ const (
 )
 
 type HistoryQueueOpt struct {
-	DB           *bolt.DB
+	DB           db.Transactor
 	LeaseManager *leaseutil.Manager
 	ContentStore *containerdsnapshot.Store
 	CleanConfig  *config.HistoryConfig

util/db/boltutil/db.go

@@ -0,0 +1,16 @@
+package boltutil
+
+import (
+	"io/fs"
+
+	"github.com/moby/buildkit/util/db"
+	bolt "go.etcd.io/bbolt"
+)
+
+func Open(p string, mode fs.FileMode, options *bolt.Options) (db.DB, error) {
+	bdb, err := bolt.Open(p, mode, options)
+	if err != nil {
+		return nil, err
+	}
+	return bdb, nil
+}

util/db/db.go

@@ -0,0 +1,8 @@
+package db
+
+import "io"
+
+type DB interface {
+	io.Closer
+	Transactor
+}

util/db/transactor.go

@@ -0,0 +1,11 @@
+package db
+
+import (
+	bolt "go.etcd.io/bbolt"
+)
+
+// Transactor is the database interface for running transactions
+type Transactor interface {
+	View(fn func(*bolt.Tx) error) error
+	Update(fn func(*bolt.Tx) error) error
+}