Golang操作BoltDB数据库

package boltdb

import (
	"bytes"
	"math"

	"go.etcd.io/bbolt"
	bolt "go.etcd.io/bbolt"
	"go.uber.org/zap"
)

type client struct {
	db  *bbolt.DB
	log *zap.Logger
}

type Client interface {
	CreateBucketIfNotExists(bucket Bucket)
	Put(bucketType Bucket, key, value []byte)
	Delete(bucketType Bucket, key []byte)
	Range(bucketType Bucket, key, endKey []byte, limit int64) (keys [][]byte, vs [][]byte)
	ForEach(bucketType Bucket, visitor func(k, v []byte) error)
	Close() error
}

func NewClient(path string, log *zap.Logger) Client {
	db, err := bbolt.Open(path, 0600, nil)
	if err != nil {
		log.Panic("failed to open database", zap.String("path", path), zap.Error(err))
	}

	return &client{db: db, log: log}
}

func (c *client) CreateBucketIfNotExists(bucketType Bucket) {
	_ = c.db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(bucketType.Name())
		if err != nil && err != bolt.ErrBucketExists {
			c.log.Error(
				"failed to create a bucket",
				zap.Stringer("bucket-name", bucketType),
				zap.Error(err),
			)
		}
		return err
	})
}

func (c *client) Put(bucketType Bucket, key, value []byte) {
	_ = c.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(bucketType.Name())
		if bucket == nil {
			c.log.Error("failed to find a bucket",
				zap.String("bucket-name", bucketType.String()))
		}

		err := bucket.Put(key, value)
		if err != nil {
			c.log.Error(
				"failed to write to a bucket",
				zap.Stringer("bucket-name", bucketType),
				zap.Error(err),
			)
		}
		return err
	})
}

func (c *client) Delete(bucketType Bucket, key []byte) {
	_ = c.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketType.Name())
		if b == nil {
			return bolt.ErrBucketNotFound
		}

		err := b.Delete(key)
		if err != nil {
			c.log.Error(
				"failed to delete a key",
				zap.Stringer("bucket-name", bucketType),
				zap.Error(err),
			)
		}

		return err
	})
}

func (c *client) Range(bucketType Bucket, key, endKey []byte, limit int64) (keys [][]byte, vs [][]byte) {
	if limit <= 0 {
		limit = math.MaxInt64
	}

	var isMatch func(b []byte) bool
	if len(endKey) > 0 {
		isMatch = func(b []byte) bool { return bytes.Compare(b, endKey) < 0 }
	} else {
		isMatch = func(b []byte) bool { return bytes.Equal(b, key) }
		limit = 1
	}

	_ = c.db.View(func(tx *bolt.Tx) error {
		cursor := tx.Bucket(bucketType.Name()).Cursor()

		for ck, cv := cursor.Seek(key); ck != nil && isMatch(ck); ck, cv = cursor.Next() {
			vs = append(vs, cv)
			keys = append(keys, ck)
			if limit == int64(len(keys)) {
				break
			}
		}

		return nil
	})

	return keys, vs
}

func (c *client) ForEach(bucketType Bucket, visitor func(k, v []byte) error) {
	_ = c.db.View(func(tx *bolt.Tx) error {
		return tx.Bucket(bucketType.Name()).ForEach(visitor)
	})
}

func (c *client) Close() error {
	return c.db.Close()
}

type Backend interface {
   // ReadTx returns a read transaction. It is replaced by ConcurrentReadTx in the main data path, see #10523.
   ReadTx() ReadTx
   BatchTx() BatchTx
   // ConcurrentReadTx returns a non-blocking read transaction.
   ConcurrentReadTx() ReadTx

   Snapshot() Snapshot
   Hash(ignores func(bucketName, keyName []byte) bool) (uint32, error)
   // Size returns the current size of the backend physically allocated.
   // The backend can hold DB space that is not utilized at the moment,
   // since it can conduct pre-allocation or spare unused space for recycling.
   // Use SizeInUse() instead for the actual DB size.
   Size() int64
   // SizeInUse returns the current size of the backend logically in use.
   // Since the backend can manage free space in a non-byte unit such as
   // number of pages, the returned value can be not exactly accurate in bytes.
   SizeInUse() int64
   // OpenReadTxN returns the number of currently open read transactions in the backend.
   OpenReadTxN() int64
   Defrag() error
   ForceCommit()
   Close() error
}


func newBackend(bcfg BackendConfig) *backend {
	
	b := &backend{
		bopts: bopts,
		db:    db,

		batchInterval: bcfg.BatchInterval,
		batchLimit:    bcfg.BatchLimit,

		readTx: &readTx{
			baseReadTx: baseReadTx{
				buf: txReadBuffer{
					txBuffer:   txBuffer{make(map[BucketID]*bucketBuffer)},
					bufVersion: 0,
				},
				buckets: make(map[BucketID]*bolt.Bucket),
				txWg:    new(sync.WaitGroup),
				txMu:    new(sync.RWMutex),
			},
		},
		txReadBufferCache: txReadBufferCache{
			mu:         sync.Mutex{},
			bufVersion: 0,
			buf:        nil,
		},

	}

	b.batchTx = newBatchTxBuffered(b)


	go b.run()
	return b
}


func (b *backend) run() {
	defer close(b.donec)
	t := time.NewTimer(b.batchInterval)
	defer t.Stop()
	for {
		select {
		case <-t.C:
		case <-b.stopc:
			b.batchTx.CommitAndStop()
			return
		}
		if b.batchTx.safePending() != 0 {
			b.batchTx.Commit()
		}
		t.Reset(b.batchInterval)
	}
}

type ReadTx interface {
   Lock()
   Unlock()
   RLock()
   RUnlock()

   UnsafeRange(bucket Bucket, key, endKey []byte, limit int64) (keys [][]byte, vals [][]byte)
   UnsafeForEach(bucket Bucket, visitor func(k, v []byte) error) error
}

type baseReadTx struct {
	// mu protects accesses to the txReadBuffer
	mu  sync.RWMutex
	buf txReadBuffer

	// TODO: group and encapsulate {txMu, tx, buckets, txWg}, as they share the same lifecycle.
	// txMu protects accesses to buckets and tx on Range requests.
	txMu    *sync.RWMutex
	tx      *bolt.Tx
	buckets map[BucketID]*bolt.Bucket
	// txWg protects tx from being rolled back at the end of a batch interval until all reads using this tx are done.
	txWg *sync.WaitGroup
}

func (baseReadTx *baseReadTx) UnsafeForEach(bucket Bucket, visitor func(k, v []byte) error) error {
   dups := make(map[string]struct{})
   getDups := func(k, v []byte) error {
      dups[string(k)] = struct{}{}
      return nil
   }
   visitNoDup := func(k, v []byte) error {
      if _, ok := dups[string(k)]; ok {
         return nil
      }
      return visitor(k, v)
   }
   if err := baseReadTx.buf.ForEach(bucket, getDups); err != nil {
      return err
   }
   baseReadTx.txMu.Lock()
   err := unsafeForEach(baseReadTx.tx, bucket, visitNoDup)
   baseReadTx.txMu.Unlock()
   if err != nil {
      return err
   }
   return baseReadTx.buf.ForEach(bucket, visitor)
}

func (baseReadTx *baseReadTx) UnsafeRange(bucketType Bucket, key, endKey []byte, limit int64) ([][]byte, [][]byte) {
   if endKey == nil {
      // forbid duplicates for single keys
      limit = 1
   }
   if limit <= 0 {
      limit = math.MaxInt64
   }
   if limit > 1 && !bucketType.IsSafeRangeBucket() {
      panic("do not use unsafeRange on non-keys bucket")
   }
   keys, vals := baseReadTx.buf.Range(bucketType, key, endKey, limit)
   if int64(len(keys)) == limit {
      return keys, vals
   }

   // find/cache bucket
   bn := bucketType.ID()
   baseReadTx.txMu.RLock()
   bucket, ok := baseReadTx.buckets[bn]
   baseReadTx.txMu.RUnlock()
   lockHeld := false
   if !ok {
      baseReadTx.txMu.Lock()
      lockHeld = true
      bucket = baseReadTx.tx.Bucket(bucketType.Name())
      baseReadTx.buckets[bn] = bucket
   }

   // ignore missing bucket since may have been created in this batch
   if bucket == nil {
      if lockHeld {
         baseReadTx.txMu.Unlock()
      }
      return keys, vals
   }
   if !lockHeld {
      baseReadTx.txMu.Lock()
   }
   c := bucket.Cursor()
   baseReadTx.txMu.Unlock()

   k2, v2 := unsafeRange(c, key, endKey, limit-int64(len(keys)))
   return append(k2, keys...), append(v2, vals...)
}

type BatchTx interface {
   ReadTx
   UnsafeCreateBucket(bucket Bucket)
   UnsafeDeleteBucket(bucket Bucket)
   UnsafePut(bucket Bucket, key []byte, value []byte)
   UnsafeSeqPut(bucket Bucket, key []byte, value []byte)
   UnsafeDelete(bucket Bucket, key []byte)
   // Commit commits a previous tx and begins a new writable one.
   Commit()
   // CommitAndStop commits the previous tx and does not create a new one.
   CommitAndStop()
}

type batchTx struct {
	sync.Mutex
	tx      *bolt.Tx
	backend *backend

	pending int
}

type batchTxBuffered struct {
	batchTx
	buf txWriteBuffer
}

func (t *batchTx) Commit() {
   t.Lock()
   t.commit(false)
   t.Unlock()
}

func (t *batchTxBuffered) Commit() {
	t.Lock()
	t.commit(false)
	t.Unlock()
}

func (t *batchTx) CommitAndStop() {
	t.Lock()
	t.commit(true)
	t.Unlock()
}

func (t *batchTx) commit(stop bool) {
	// commit the last tx
	if t.tx != nil {
		if t.pending == 0 && !stop {
			return
		}

		start := time.Now()

		// gofail: var beforeCommit struct{}
		err := t.tx.Commit()
		// gofail: var afterCommit struct{}

		rebalanceSec.Observe(t.tx.Stats().RebalanceTime.Seconds())
		spillSec.Observe(t.tx.Stats().SpillTime.Seconds())
		writeSec.Observe(t.tx.Stats().WriteTime.Seconds())
		commitSec.Observe(time.Since(start).Seconds())
		atomic.AddInt64(&t.backend.commits, 1)

		t.pending = 0
		if err != nil {
			t.backend.lg.Fatal("failed to commit tx", zap.Error(err))
		}
	}
	if !stop {
		t.tx = t.backend.begin(true)
	}
}

func (t *batchTxBuffered) UnsafePut(bucket Bucket, key []byte, value []byte) {
   t.batchTx.UnsafePut(bucket, key, value)
   t.buf.put(bucket, key, value)
}

func (t *batchTxBuffered) UnsafeSeqPut(bucket Bucket, key []byte, value []byte) {
   t.batchTx.UnsafeSeqPut(bucket, key, value)
   t.buf.putSeq(bucket, key, value)
}