eth-tracker/db/bolt.go

package db

import (
	"bytes"
	"encoding/binary"
	"fmt"

	"github.com/bits-and-blooms/bitset"
	bolt "go.etcd.io/bbolt"
)

type boltDB struct {
	db *bolt.DB
}

const (
	dbFolderName = "tracker_db"

	upperBoundKey = "upper"
	lowerBoundKey = "lower"
)

var sortableOrder = binary.BigEndian

func NewBoltDB() (DB, error) {
	db, err := bolt.Open(dbFolderName, 0600, nil)
	if err != nil {
		return nil, err
	}

	db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("blocks"))
		if err != nil {
			return fmt.Errorf("create bucket: %s", err)
		}
		return nil
	})

	return &boltDB{
		db: db,
	}, nil
}

func (d *boltDB) Close() error {
	return d.db.Close()
}

func (d *boltDB) get(k string) ([]byte, error) {
	var v []byte
	err := d.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("blocks"))
		v = b.Get([]byte(k))
		return nil
	})

	if err != nil {
		return nil, err
	}

	return v, nil
}

func (d *boltDB) setUint64(k string, v uint64) error {
	err := d.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("blocks"))
		return b.Put([]byte(k), marshalUint64(v))
	})
	if err != nil {
		return err
	}
	return nil
}

func (d *boltDB) setUint64AsKey(v uint64) error {
	err := d.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("blocks"))
		return b.Put(marshalUint64(v), nil)
	})
	if err != nil {
		return err
	}
	return nil
}

func unmarshalUint64(b []byte) uint64 {
	return sortableOrder.Uint64(b)
}

func marshalUint64(v uint64) []byte {
	b := make([]byte, 8)
	sortableOrder.PutUint64(b, v)
	return b
}

func (d *boltDB) SetLowerBound(v uint64) error {
	return d.setUint64(lowerBoundKey, v)
}

func (d *boltDB) GetLowerBound() (uint64, error) {
	v, err := d.get(lowerBoundKey)
	if err != nil {
		return 0, err
	}

	if v == nil {
		return 0, nil
	}

	return unmarshalUint64(v), nil
}

func (d *boltDB) SetUpperBound(v uint64) error {
	return d.setUint64(upperBoundKey, v)
}

func (d *boltDB) GetUpperBound() (uint64, error) {
	v, err := d.get(upperBoundKey)
	if err != nil {
		return 0, err
	}
	return unmarshalUint64(v), nil
}

func (d *boltDB) SetValue(v uint64) error {
	return d.setUint64AsKey(v)
}

func (d *boltDB) GetMissingValuesBitSet(lowerBound uint64, upperBound uint64) (*bitset.BitSet, error) {
	var b bitset.BitSet

	err := d.db.View(func(tx *bolt.Tx) error {
		var (
			lowerRaw = marshalUint64(lowerBound)
			upperRaw = marshalUint64(upperBound)
		)

		for i := lowerBound; i <= upperBound; i++ {
			b.Set(uint(i))
		}

		c := tx.Bucket([]byte("blocks")).Cursor()

		for k, _ := c.Seek(lowerRaw); k != nil && bytes.Compare(k, upperRaw) <= 0; k, _ = c.Next() {
			b.Clear(uint(unmarshalUint64(k)))
		}

		return nil
	})
	if err != nil {
		return nil, err
	}

	return &b, nil
}

func (d *boltDB) Cleanup() error {
	lowerBound, err := d.GetLowerBound()
	if err != nil {
		return err
	}
	target := marshalUint64(lowerBound - 1)

	err = d.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("blocks"))
		c := b.Cursor()

		for k, _ := c.First(); k != nil && bytes.Compare(k, target) <= 0; k, _ = c.Next() {
			if err := b.Delete(k); err != nil {
				return err
			}
		}

		return nil
	})
	if err != nil {
		return err
	}

	return nil
}