[Storage Refactor] Init pebble DB in scaffold.

cmd/node_builder.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/cockroachdb/pebble"
 	"github.com/dgraph-io/badger/v2"
 	madns "github.com/multiformats/go-multiaddr-dns"
 	"github.com/onflow/crypto"
@@ -151,6 +152,7 @@ type BaseConfig struct {
 	DynamicStartupEpoch         string
 	DynamicStartupSleepInterval time.Duration
 	datadir                     string
+	pebbleDir                   string
 	secretsdir                  string
 	secretsDBEnabled            bool
 	InsecureSecretsDB           bool
@@ -164,7 +166,6 @@ type BaseConfig struct {
 	MetricsEnabled              bool
 	guaranteesCacheSize         uint
 	receiptsCacheSize           uint
-	db                          *badger.DB
 	HeroCacheMetricsEnable      bool
 	SyncCoreConfig              chainsync.Config
 	CodecFactory                func() network.Codec
@@ -198,6 +199,7 @@ type NodeConfig struct {
 	MetricsRegisterer prometheus.Registerer
 	Metrics           Metrics
 	DB                *badger.DB
+	PebbleDB          *pebble.DB
 	SecretsDB         *badger.DB
 	Storage           Storage
 	ProtocolEvents    *events.Distributor
@@ -253,6 +255,7 @@ type StateExcerptAtBoot struct {
 
 func DefaultBaseConfig() *BaseConfig {
 	datadir := "/data/protocol"
+	pebbleDir := "/data/protocol-pebble"
 
 	// NOTE: if the codec used in the network component is ever changed any code relying on
 	// the message format specific to the codec must be updated. i.e: the AuthorizedSenderValidator.
@@ -269,6 +272,7 @@ func DefaultBaseConfig() *BaseConfig {
 		ObserverMode:     false,
 		BootstrapDir:     "bootstrap",
 		datadir:          datadir,
+		pebbleDir:        pebbleDir,
 		secretsdir:       NotSet,
 		secretsDBEnabled: true,
 		level:            "info",

cmd/scaffold.go

@@ -31,6 +31,7 @@ import (
 	"github.com/onflow/flow-go/admin/commands/common"
 	storageCommands "github.com/onflow/flow-go/admin/commands/storage"
 	"github.com/onflow/flow-go/cmd/build"
+	"github.com/onflow/flow-go/cmd/scaffold"
 	"github.com/onflow/flow-go/config"
 	"github.com/onflow/flow-go/consensus/hotstuff/persister"
 	"github.com/onflow/flow-go/fvm/initialize"
@@ -164,6 +165,7 @@ func (fnb *FlowNodeBuilder) BaseFlags() {
 	fnb.flags.StringVar(&fnb.BaseConfig.BindAddr, "bind", defaultConfig.BindAddr, "address to bind on")
 	fnb.flags.StringVarP(&fnb.BaseConfig.BootstrapDir, "bootstrapdir", "b", defaultConfig.BootstrapDir, "path to the bootstrap directory")
 	fnb.flags.StringVarP(&fnb.BaseConfig.datadir, "datadir", "d", defaultConfig.datadir, "directory to store the public database (protocol state)")
+	fnb.flags.StringVar(&fnb.BaseConfig.pebbleDir, "pebble-dir", defaultConfig.pebbleDir, "directory to store the public pebble database (protocol state)")
 	fnb.flags.StringVar(&fnb.BaseConfig.secretsdir, "secretsdir", defaultConfig.secretsdir, "directory to store private database (secrets)")
 	fnb.flags.StringVarP(&fnb.BaseConfig.level, "loglevel", "l", defaultConfig.level, "level for logging output")
 	fnb.flags.Uint32Var(&fnb.BaseConfig.debugLogLimit, "debug-log-limit", defaultConfig.debugLogLimit, "max number of debug/trace log events per second")
@@ -1056,13 +1058,7 @@ func (fnb *FlowNodeBuilder) initProfiler() error {
 	return nil
 }
 
-func (fnb *FlowNodeBuilder) initDB() error {
-
-	// if a db has been passed in, use that instead of creating one
-	if fnb.BaseConfig.db != nil {
-		fnb.DB = fnb.BaseConfig.db
-		return nil
-	}
+func (fnb *FlowNodeBuilder) initBadgerDB() error {
 
 	// Pre-create DB path (Badger creates only one-level dirs)
 	err := os.MkdirAll(fnb.BaseConfig.datadir, 0700)
@@ -1110,6 +1106,17 @@ func (fnb *FlowNodeBuilder) initDB() error {
 	return nil
 }
 
+func (fnb *FlowNodeBuilder) initPebbleDB() error {
+	db, closer, err := scaffold.InitPebbleDB(fnb.BaseConfig.pebbleDir)
+	if err != nil {
+		return err
+	}
+
+	fnb.PebbleDB = db
+	fnb.ShutdownFunc(closer.Close)
+	return nil
+}
+
 func (fnb *FlowNodeBuilder) initSecretsDB() error {
 
 	// if the secrets DB is disabled (only applicable for Consensus Follower,
@@ -1889,9 +1896,7 @@ func WithBindAddress(bindAddress string) Option {
 
 func WithDataDir(dataDir string) Option {
 	return func(config *BaseConfig) {
-		if config.db == nil {
-			config.datadir = dataDir
-		}
+		config.datadir = dataDir
 	}
 }
 
@@ -1925,14 +1930,6 @@ func WithLogLevel(level string) Option {
 	}
 }
 
-// WithDB takes precedence over WithDataDir and datadir will be set to empty if DB is set using this option
-func WithDB(db *badger.DB) Option {
-	return func(config *BaseConfig) {
-		config.db = db
-		config.datadir = ""
-	}
-}
-
 // FlowNode creates a new Flow node builder with the given name.
 func FlowNode(role string, opts ...Option) *FlowNodeBuilder {
 	config := DefaultBaseConfig()
@@ -2038,7 +2035,11 @@ func (fnb *FlowNodeBuilder) onStart() error {
 		return err
 	}
 
-	if err := fnb.initDB(); err != nil {
+	if err := fnb.initBadgerDB(); err != nil {
+		return err
+	}
+
+	if err := fnb.initPebbleDB(); err != nil {
 		return err
 	}
 

cmd/scaffold/pebble_db.go

@@ -0,0 +1,26 @@
+package scaffold
+
+import (
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/cockroachdb/pebble"
+
+	pebblestorage "github.com/onflow/flow-go/storage/pebble"
+)
+
+func InitPebbleDB(dir string) (*pebble.DB, io.Closer, error) {
+	// Pre-create DB path
+	err := os.MkdirAll(dir, 0700)
+	if err != nil {
+		return nil, nil, fmt.Errorf("could not create pebble db (path: %s): %w", dir, err)
+	}
+
+	db, err := pebblestorage.OpenDefaultPebbleDB(dir)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return db, db, nil
+}

cmd/scaffold/pebble_db_test.go

@@ -0,0 +1,17 @@
+package scaffold
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/onflow/flow-go/utils/unittest"
+)
+
+func TestInitPebbleDB(t *testing.T) {
+	unittest.RunWithTempDir(t, func(dir string) {
+		_, closer, err := InitPebbleDB(dir)
+		require.NoError(t, err)
+		require.NoError(t, closer.Close())
+	})
+}

follower/consensus_follower.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/dgraph-io/badger/v2"
 	"github.com/onflow/crypto"
 	"github.com/rs/zerolog"
 
@@ -36,7 +35,6 @@ type Config struct {
 	networkPrivKey   crypto.PrivateKey   // the network private key of this node
 	bootstrapNodes   []BootstrapNodeInfo // the bootstrap nodes to use
 	bindAddr         string              // address to bind on
-	db               *badger.DB          // the badger DB storage to use for the protocol state
 	dataDir          string              // directory to store the protocol state (if the badger storage is not provided)
 	bootstrapDir     string              // path to the bootstrap directory
 	logLevel         string              // log level
@@ -51,9 +49,7 @@ type Option func(c *Config)
 // If a database is supplied, then data directory will be set to empty string
 func WithDataDir(dataDir string) Option {
 	return func(cf *Config) {
-		if cf.db == nil {
-			cf.dataDir = dataDir
-		}
+		cf.dataDir = dataDir
 	}
 }
 
@@ -69,15 +65,6 @@ func WithLogLevel(level string) Option {
 	}
 }
 
-// WithDB sets the underlying database that will be used to store the chain state
-// WithDB takes precedence over WithDataDir and datadir will be set to empty if DB is set using this option
-func WithDB(db *badger.DB) Option {
-	return func(cf *Config) {
-		cf.db = db
-		cf.dataDir = ""
-	}
-}
-
 func WithExposeMetrics(expose bool) Option {
 	return func(c *Config) {
 		c.exposeMetrics = expose
@@ -142,9 +129,6 @@ func getBaseOptions(config *Config) []cmd.Option {
 	if config.logLevel != "" {
 		options = append(options, cmd.WithLogLevel(config.logLevel))
 	}
-	if config.db != nil {
-		options = append(options, cmd.WithDB(config.db))
-	}
 	if config.exposeMetrics {
 		options = append(options, cmd.WithMetricsEnabled(config.exposeMetrics))
 	}

storage/pebble/open.go

@@ -82,6 +82,8 @@ func MustOpenDefaultPebbleDB(dir string) (*pebble.DB, error) {
 }
 
 // IsPebbleInitialized checks if the given folder contains a valid Pebble DB.
+// return error if the folder does not exist, is not a directory, or is missing required files
+// return nil if the folder contains a valid Pebble DB
 func IsPebbleInitialized(folderPath string) error {
 	// Check if the folder exists
 	info, err := os.Stat(folderPath)