feat: simplify pebble datastore creation with single set of options

datastore.go

@@ -20,53 +20,53 @@
 // It supports batching. It does not support TTL or transactions, because pebble
 // doesn't have those features.
 type Datastore struct {
-	db      *pebble.DB
-	status  int32
-	closing chan struct{}
-	wg      sync.WaitGroup
+	db *pebble.DB
 
-	opts *pebble.Options
+	cache      *pebble.Cache
+	closing    chan struct{}
+	disableWAL bool
+	status     int32
+	wg         sync.WaitGroup
 }
 
 var _ ds.Datastore = (*Datastore)(nil)
 var _ ds.Batching = (*Datastore)(nil)
 
-type DatastoreOption func(*Datastore)
-
-// WithPebbleDB is used to configure the Datastore with a custom DB.
-func WithPebbleDB(db *pebble.DB) DatastoreOption {
-	return func(ds *Datastore) {
-		ds.db = db
-	}
-}
-
 // NewDatastore creates a pebble-backed datastore.
-// Users can provide pebble options or rely on Pebble's defaults.
-func NewDatastore(path string, opts *pebble.Options, options ...DatastoreOption) (*Datastore, error) {
-	if opts == nil {
-		opts = &pebble.Options{}
-		opts.EnsureDefaults()
-	}
-	opts.Logger = logger
-
-	store := &Datastore{
-		opts:    opts,
-		closing: make(chan struct{}),
-	}
-
-	for _, opt := range options {
-		opt(store)
-	}
-
-	if store.db == nil {
-		db, err := pebble.Open(path, opts)
+//
+// Users can provide pebble options using WithPebbleOpts or rely on Pebble's
+// defaults. Any pebble options that are not assigned a value are assigned
+// pebble's default value for the option.
+func NewDatastore(path string, options ...Option) (*Datastore, error) {
+	opts := getOpts(options)
+
+	// Use the provided database or create a new one.
+	db := opts.db
+	var disableWAL bool
+	var cache *pebble.Cache
+	if db == nil {
+		pebbleOpts := opts.pebbleOpts.EnsureDefaults()
+		pebbleOpts.Logger = logger
+		disableWAL = pebbleOpts.DisableWAL
+		// Use the provided cache, create a custom-sized cache, or use default.
+		if pebbleOpts.Cache == nil && opts.cacheSize != 0 {
+			cache = pebble.NewCache(opts.cacheSize)
+			// Keep ref to cache if it is created here.
+			pebbleOpts.Cache = cache
+		}
+		var err error
+		db, err = pebble.Open(path, pebbleOpts)
 		if err != nil {
 			return nil, fmt.Errorf("failed to open pebble database: %w", err)
 		}
-		store.db = db
 	}
 
-	return store, nil
+	return &Datastore{
+		db:         db,
+		disableWAL: disableWAL,
+		cache:      cache,
+		closing:    make(chan struct{}),
+	}, nil
 }
 
 // get performs a get on the database, copying the value to a new slice and
@@ -337,7 +337,7 @@
 	// crash. In pebble this is done by fsyncing the WAL, which can be requested when
 	// performing write operations. But there is no separate operation to fsync
 	// only. The closest is LogData, which actually writes a log entry on the WAL.
-	if d.opts.DisableWAL { // otherwise this errors
+	if d.disableWAL { // otherwise this errors
 		return nil
 	}
 	err := d.db.LogData(nil, pebble.Sync)
@@ -357,6 +357,9 @@
 		d.wg.Wait()
 		return nil
 	}
+	if d.cache != nil {
+		defer d.cache.Unref()
+	}
 	close(d.closing)
 	d.wg.Wait()
 	_ = d.db.Flush()

datastore_test.go

@@ -26,7 +26,7 @@ func newDatastore(t *testing.T) (*Datastore, func()) {
 		t.Fatal(err)
 	}
 
-	d, err := NewDatastore(path, nil)
+	d, err := NewDatastore(path)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -50,7 +50,7 @@ func newDatastoreWithPebbleDB(t *testing.T) (*Datastore, func()) {
 		t.Fatal(err)
 	}
 
-	d, err := NewDatastore(path, nil, WithPebbleDB(db))
+	d, err := NewDatastore(path, WithPebbleDB(db))
 	if err != nil {
 		t.Fatal(err)
 	}

option.go

@@ -0,0 +1,48 @@
+package pebbleds
+
+import (
+	"github.com/cockroachdb/pebble"
+)
+
+type config struct {
+	cacheSize  int64
+	db         *pebble.DB
+	pebbleOpts *pebble.Options
+}
+
+type Option func(*config)
+
+func getOpts(options []Option) config {
+	var cfg config
+	for _, opt := range options {
+		if opt == nil {
+			continue
+		}
+		opt(&cfg)
+	}
+	return cfg
+}
+
+// WithCacheSize configures the size of pebble's shared block cache. A value of
+// 0 (the default) uses the default cache size.
+func WithCacheSize(size int64) Option {
+	return func(c *config) {
+		c.cacheSize = size
+	}
+}
+
+// WithPebbleDB is used to configure the Datastore with a custom DB.
+func WithPebbleDB(db *pebble.DB) Option {
+	return func(c *config) {
+		c.db = db
+	}
+}
+
+// WithPebbleOpts sets any/all configurable values for pebble. If not set, the
+// default configuration values are used. Any unspecified value in opts is
+// replaced by the default value.
+func WithPebbleOpts(opts *pebble.Options) Option {
+	return func(c *config) {
+		c.pebbleOpts = opts
+	}
+}