worker/internal/workercluster/store.go

package workercluster

import (
	"encoding/binary"
	"fmt"
	"os"
	"path/filepath"

	bolt "go.etcd.io/bbolt"

	"github.com/hashicorp/raft"
)

// BoltStore bundles a single bbolt DB used as both the Raft LogStore and
// the StableStore, plus a directory of snapshot files on disk.
//
// One DB for both stores keeps the on-disk layout simple: the bbolt file
// holds two buckets, "logs" and "stable", and the snapshot store is just
// `<dir>/snapshots/<id>.json` files alongside it. The library's API
// surfaces LogStore / StableStore / SnapshotStore separately so the
// application code does not have to care.
type BoltStore struct {
	db *bolt.DB

	logBucket    []byte
	stableBucket []byte

	snapshotDir string
}

const (
	defaultLogBucket    = "logs"
	defaultStableBucket = "stable"
)

// NewBoltStore opens or creates the bbolt-backed store rooted at dir.
// The directory is created if missing. The DB file lives at
// <dir>/raft.db; snapshots live in <dir>/snapshots/.
func NewBoltStore(dir string) (*BoltStore, error) {
	if err := os.MkdirAll(dir, 0o750); err != nil {
		return nil, fmt.Errorf("workercluster: mkdir %q: %w", dir, err)
	}
	if err := os.MkdirAll(filepath.Join(dir, "snapshots"), 0o750); err != nil {
		return nil, fmt.Errorf("workercluster: mkdir snapshots: %w", err)
	}

	db, err := bolt.Open(filepath.Join(dir, "raft.db"), 0o600, &bolt.Options{})
	if err != nil {
		return nil, fmt.Errorf("workercluster: open bolt: %w", err)
	}

	bs := &BoltStore{
		db:           db,
		logBucket:    []byte(defaultLogBucket),
		stableBucket: []byte(defaultStableBucket),
		snapshotDir:  filepath.Join(dir, "snapshots"),
	}

	if err := db.Update(func(tx *bolt.Tx) error {
		for _, b := range [][]byte{bs.logBucket, bs.stableBucket} {
			if _, err := tx.CreateBucketIfNotExists(b); err != nil {
				return err
			}
		}
		return nil
	}); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("workercluster: init buckets: %w", err)
	}

	return bs, nil
}

// Close releases the underlying bbolt handle.
func (s *BoltStore) Close() error {
	if s.db == nil {
		return nil
	}
	return s.db.Close()
}

// StableStore returns the StableStore half of the backing store.
func (s *BoltStore) StableStore() raft.StableStore {
	return &boltStableStore{store: s}
}

// LogStore returns the LogStore half of the backing store.
func (s *BoltStore) LogStore() raft.LogStore {
	return &boltLogStore{store: s}
}

// SnapshotStore returns the SnapshotStore half of the backing store.
func (s *BoltStore) SnapshotStore() raft.SnapshotStore {
	return &boltSnapshotStore{store: s}
}

// DB exposes the underlying bbolt handle for tests that want to poke
// at it directly. Production code should never need this.
func (s *BoltStore) DB() *bolt.DB {
	return s.db
}

// boltStableStore implements raft.StableStore on top of a BoltStore.
type boltStableStore struct {
	store *BoltStore
}

func (s *boltStableStore) Set(key, val []byte) error {
	return s.store.db.Update(func(tx *bolt.Tx) error {
		return tx.Bucket(s.store.stableBucket).Put(key, val)
	})
}

func (s *boltStableStore) Get(key []byte) ([]byte, error) {
	var out []byte
	err := s.store.db.View(func(tx *bolt.Tx) error {
		v := tx.Bucket(s.store.stableBucket).Get(key)
		if v != nil {
			// Copy out of the mmap region; bbolt may reuse the slice.
			out = append([]byte(nil), v...)
		}
		return nil
	})
	return out, err
}

func (s *boltStableStore) SetUint64(key []byte, val uint64) error {
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, val)
	return s.Set(key, buf)
}

func (s *boltStableStore) GetUint64(key []byte) (uint64, error) {
	v, err := s.Get(key)
	if err != nil {
		return 0, err
	}
	if len(v) == 0 {
		return 0, nil
	}
	if len(v) < 8 {
		return 0, nil
	}
	return binary.BigEndian.Uint64(v), nil
}

// boltLogStore implements raft.LogStore on top of a BoltStore.
//
// Layout: each log entry is stored under the 8-byte big-endian index
// key. FirstIndex scans for the lowest key, LastIndex for the highest.
type boltLogStore struct {
	store *BoltStore
}

func logKey(index uint64) []byte {
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, index)
	return buf
}

func (s *boltLogStore) FirstIndex() (uint64, error) {
	var idx uint64
	err := s.store.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket(s.store.logBucket).Cursor()
		k, _ := c.First()
		if k == nil {
			idx = 0
			return nil
		}
		idx = binary.BigEndian.Uint64(k)
		return nil
	})
	return idx, err
}

func (s *boltLogStore) LastIndex() (uint64, error) {
	var idx uint64
	err := s.store.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket(s.store.logBucket).Cursor()
		k, _ := c.Last()
		if k == nil {
			idx = 0
			return nil
		}
		idx = binary.BigEndian.Uint64(k)
		return nil
	})
	return idx, err
}

func (s *boltLogStore) GetLog(index uint64, log *raft.Log) error {
	err := s.store.db.View(func(tx *bolt.Tx) error {
		raw := tx.Bucket(s.store.logBucket).Get(logKey(index))
		if raw == nil {
			return raft.ErrLogNotFound
		}
		return decodeLog(raw, log)
	})
	return err
}

func (s *boltLogStore) StoreLog(log *raft.Log) error {
	return s.StoreLogs([]*raft.Log{log})
}

func (s *boltLogStore) StoreLogs(logs []*raft.Log) error {
	if len(logs) == 0 {
		return nil
	}
	return s.store.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(s.store.logBucket)
		for _, log := range logs {
			raw, err := encodeLog(log)
			if err != nil {
				return err
			}
			if err := b.Put(logKey(log.Index), raw); err != nil {
				return err
			}
		}
		return nil
	})
}

func (s *boltLogStore) DeleteRange(lo, hi uint64) error {
	return s.store.db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket(s.store.logBucket)
		for i := lo; i <= hi; i++ {
			if err := b.Delete(logKey(i)); err != nil {
				return err
			}
		}
		return nil
	})
}