diff --git a/datastore.go b/datastore.go index db2f407..7647c84 100644 --- a/datastore.go +++ b/datastore.go @@ -20,53 +20,53 @@ var logger = log.Logger("pebble") // 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 @@ func (d *Datastore) Sync(ctx context.Context, _ ds.Key) error { // 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 @@ func (d *Datastore) Close() error { d.wg.Wait() return nil } + if d.cache != nil { + defer d.cache.Unref() + } close(d.closing) d.wg.Wait() _ = d.db.Flush() diff --git a/datastore_test.go b/datastore_test.go index 71cdf6f..3d3c9d7 100644 --- a/datastore_test.go +++ b/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) } diff --git a/option.go b/option.go new file mode 100644 index 0000000..8afa512 --- /dev/null +++ b/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 + } +}