erigon-pulse/ethdb/kv_badger.go

package ethdb

import (
	"context"
	"runtime"
	"time"

	"github.com/dgraph-io/badger/v2"
	"github.com/ledgerwatch/turbo-geth/log"
)

type badgerOpts struct {
	Badger badger.Options
}

func (opts badgerOpts) Path(path string) badgerOpts {
	opts.Badger = opts.Badger.WithDir(path).WithValueDir(path)
	return opts
}

func (opts badgerOpts) InMem() badgerOpts {
	opts.Badger = opts.Badger.WithInMemory(true)
	return opts
}

func (opts badgerOpts) ReadOnly() badgerOpts {
	opts.Badger = opts.Badger.WithReadOnly(true)
	return opts
}

func (opts badgerOpts) Open(ctx context.Context) (KV, error) {
	logger := log.New("badger_db", opts.Badger.Dir)

	oldMaxProcs := runtime.GOMAXPROCS(0)
	if oldMaxProcs < minGoMaxProcs {
		runtime.GOMAXPROCS(minGoMaxProcs)
		logger.Info("Bumping GOMAXPROCS", "old", oldMaxProcs, "new", minGoMaxProcs)
	}
	opts.Badger = opts.Badger.WithMaxTableSize(512 << 20)

	db, err := badger.Open(opts.Badger)
	if err != nil {
		return nil, err
	}

	ticker := time.NewTicker(gcPeriod)
	// Start GC in backround
	go func() {
		for range ticker.C {
			err := db.RunValueLogGC(0.5)
			if err != badger.ErrNoRewrite {
				logger.Info("Badger GC run", "err", err)
			}
		}
	}()

	return &badgerDB{
		opts:     opts,
		badger:   db,
		log:      logger,
		gcTicker: ticker, // Garbage Collector
	}, nil
}

func (opts badgerOpts) MustOpen(ctx context.Context) KV {
	db, err := opts.Open(ctx)
	if err != nil {
		panic(err)
	}
	return db
}

type badgerDB struct {
	opts     badgerOpts
	badger   *badger.DB
	gcTicker *time.Ticker
	log      log.Logger
}

func NewBadger() badgerOpts {
	return badgerOpts{Badger: badger.DefaultOptions("")}
}

// Close closes BoltKV
// All transactions must be closed before closing the database.
func (db *badgerDB) Close() {
	if db.gcTicker != nil {
		db.gcTicker.Stop()
	}
	if err := db.badger.Close(); err != nil {
		db.log.Warn("failed to close badger DB", "err", err)
	} else {
		db.log.Info("badger database closed")
	}
}

func (db *badgerDB) Begin(ctx context.Context, writable bool) (Tx, error) {
	return &badgerTx{
		db:     db,
		ctx:    ctx,
		badger: db.badger.NewTransaction(writable),
	}, nil
}

type badgerTx struct {
	ctx context.Context
	db  *badgerDB

	badger          *badger.Txn
	badgerIterators []*badger.Iterator
}

type badgerBucket struct {
	tx *badgerTx

	prefix  []byte
	nameLen uint
}

type badgerCursor struct {
	ctx    context.Context
	bucket badgerBucket
	prefix []byte

	badgerOpts badger.IteratorOptions

	badger *badger.Iterator

	k   []byte
	v   []byte
	err error
}

func (db *badgerDB) View(ctx context.Context, f func(tx Tx) error) (err error) {
	t := &badgerTx{db: db, ctx: ctx}
	return db.badger.View(func(tx *badger.Txn) error {
		defer t.cleanup()
		t.badger = tx
		return f(t)
	})
}

func (db *badgerDB) Update(ctx context.Context, f func(tx Tx) error) (err error) {
	t := &badgerTx{db: db, ctx: ctx}
	return db.badger.Update(func(tx *badger.Txn) error {
		defer t.cleanup()
		t.badger = tx
		return f(t)
	})
}

func (tx *badgerTx) Bucket(name []byte) Bucket {
	b := badgerBucket{tx: tx, nameLen: uint(len(name))}
	b.prefix = name
	return b
}

func (tx *badgerTx) Commit(ctx context.Context) error {
	tx.cleanup()
	return tx.badger.Commit()
}

func (tx *badgerTx) Rollback() error {
	tx.cleanup()
	tx.badger.Discard()
	return nil
}

func (tx *badgerTx) cleanup() {
	for _, it := range tx.badgerIterators {
		it.Close()
	}
}

func (c *badgerCursor) Prefix(v []byte) Cursor {
	c.prefix = append(c.bucket.prefix[:c.bucket.nameLen], v...)
	c.badgerOpts.Prefix = c.prefix
	return c
}

func (c *badgerCursor) MatchBits(n uint) Cursor {
	panic("not implemented yet")
}

func (c *badgerCursor) Prefetch(v uint) Cursor {
	c.badgerOpts.PrefetchSize = int(v)
	return c
}

func (c *badgerCursor) NoValues() NoValuesCursor {
	c.badgerOpts.PrefetchValues = false
	return &badgerNoValuesCursor{badgerCursor: *c}
}

func (b badgerBucket) Get(key []byte) (val []byte, err error) {
	select {
	case <-b.tx.ctx.Done():
		return nil, b.tx.ctx.Err()
	default:
	}

	var item *badger.Item
	b.prefix = append(b.prefix[:b.nameLen], key...)
	item, err = b.tx.badger.Get(b.prefix)
	if item != nil {
		val, err = item.ValueCopy(nil) // can improve this by using pool
	}
	return val, err
}

func (b badgerBucket) Put(key []byte, value []byte) error {
	select {
	case <-b.tx.ctx.Done():
		return b.tx.ctx.Err()
	default:
	}

	b.prefix = append(b.prefix[:b.nameLen], key...)
	return b.tx.badger.Set(b.prefix, value)
}

func (b badgerBucket) Delete(key []byte) error {
	select {
	case <-b.tx.ctx.Done():
		return b.tx.ctx.Err()
	default:
	}

	b.prefix = append(b.prefix[:b.nameLen], key...)
	return b.tx.badger.Delete(b.prefix)
}

func (b badgerBucket) Cursor() Cursor {
	c := &badgerCursor{bucket: b, ctx: b.tx.ctx, badgerOpts: badger.DefaultIteratorOptions}
	c.prefix = append(c.prefix, b.prefix[:b.nameLen]...) // set bucket
	c.badgerOpts.Prefix = c.prefix
	return c
}

func (c *badgerCursor) initCursor() {
	if c.badger != nil {
		return
	}

	_ = c.bucket
	_ = c.bucket.tx
	_ = c.bucket.tx.badger

	c.badger = c.bucket.tx.badger.NewIterator(c.badgerOpts)
	// add to auto-cleanup on end of transactions
	if c.bucket.tx.badgerIterators == nil {
		c.bucket.tx.badgerIterators = make([]*badger.Iterator, 0, 1)
	}
	c.bucket.tx.badgerIterators = append(c.bucket.tx.badgerIterators, c.badger)
}

func (c *badgerCursor) First() ([]byte, []byte, error) {
	c.initCursor()

	c.badger.Rewind()
	if !c.badger.Valid() {
		c.k = nil
		return c.k, c.v, c.err
	}
	item := c.badger.Item()
	c.k = item.Key()[c.bucket.nameLen:]
	if c.badgerOpts.PrefetchValues {
		c.v, c.err = item.ValueCopy(c.v) // bech show: using .ValueCopy on same buffer has same speed as item.Value()
	}
	return c.k, c.v, c.err
}

func (c *badgerCursor) Seek(seek []byte) ([]byte, []byte, error) {
	select {
	case <-c.ctx.Done():
		return nil, nil, c.ctx.Err()
	default:
	}

	c.initCursor()

	c.badger.Seek(append(c.bucket.prefix[:c.bucket.nameLen], seek...))
	if !c.badger.Valid() {
		c.k = nil
		return c.k, c.v, c.err
	}
	item := c.badger.Item()
	c.k = item.Key()[c.bucket.nameLen:]
	if c.badgerOpts.PrefetchValues {
		c.v, c.err = item.ValueCopy(c.v)
	}
	return c.k, c.v, c.err
}

func (c *badgerCursor) Next() ([]byte, []byte, error) {
	select {
	case <-c.ctx.Done():
		return nil, nil, c.ctx.Err()
	default:
	}

	c.badger.Next()
	if !c.badger.Valid() {
		c.k = nil
		return c.k, c.v, c.err
	}
	item := c.badger.Item()
	c.k = item.Key()[c.bucket.nameLen:]
	if c.badgerOpts.PrefetchValues {
		c.v, c.err = item.ValueCopy(c.v)
	}

	return c.k, c.v, c.err
}

func (c *badgerCursor) Walk(walker func(k, v []byte) (bool, error)) error {
	for k, v, err := c.First(); k != nil || err != nil; k, v, err = c.Next() {
		if err != nil {
			return err
		}
		ok, err := walker(k, v)
		if err != nil {
			return err
		}
		if !ok {
			return nil
		}
	}
	return nil
}

type badgerNoValuesCursor struct {
	badgerCursor
}

func (c *badgerNoValuesCursor) Walk(walker func(k []byte, vSize uint32) (bool, error)) error {
	for k, vSize, err := c.First(); k != nil || err != nil; k, vSize, err = c.Next() {
		if err != nil {
			return err
		}
		ok, err := walker(k, vSize)
		if err != nil {
			return err
		}
		if !ok {
			return nil
		}
	}
	return nil
}

func (c *badgerNoValuesCursor) First() ([]byte, uint32, error) {
	c.initCursor()
	c.badger.Rewind()
	if !c.badger.Valid() {
		c.k = nil
		return c.k, 0, c.err
	}
	item := c.badger.Item()
	c.k = item.Key()[c.bucket.nameLen:]
	return c.k, uint32(item.ValueSize()), c.err
}

func (c *badgerNoValuesCursor) Seek(seek []byte) ([]byte, uint32, error) {
	select {
	case <-c.ctx.Done():
		return nil, 0, c.ctx.Err()
	default:
	}

	c.initCursor()

	c.badger.Seek(append(c.bucket.prefix[:c.bucket.nameLen], seek...))
	if !c.badger.Valid() {
		c.k = nil
		return c.k, 0, c.err
	}
	item := c.badger.Item()
	c.k = item.Key()[c.bucket.nameLen:]

	return c.k, uint32(item.ValueSize()), c.err
}

func (c *badgerNoValuesCursor) Next() ([]byte, uint32, error) {
	select {
	case <-c.ctx.Done():
		return nil, 0, c.ctx.Err()
	default:
	}

	c.badger.Next()
	if !c.badger.Valid() {
		c.k = nil
		return c.k, 0, c.err
	}
	item := c.badger.Item()
	c.k = item.Key()[c.bucket.nameLen:]
	return c.k, uint32(item.ValueSize()), c.err
}