worker/internal/workercluster/snapshot.go

package workercluster

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/hashicorp/raft"
)

// boltSnapshotStore stores raft snapshots as JSON files on disk under
// <storeDir>/snapshots/<id>.json. The store is concurrency-safe via
// the OS filesystem; no per-call mutex is needed.
type boltSnapshotStore struct {
	store *BoltStore
}

// snapshotRecord is the on-disk shape of a single snapshot.
type snapshotRecord struct {
	Meta  raft.SnapshotMeta `json:"meta"`
	State []byte            `json:"state"`
}

// Create opens a new snapshot sink at index/term. The library will write
// data to the sink; we close it by atomically renaming a temp file.
func (s *boltSnapshotStore) Create(version raft.SnapshotVersion, index, term uint64, configuration raft.Configuration,
	configurationIndex uint64, _ raft.Transport,
) (raft.SnapshotSink, error) {
	if version != raft.SnapshotVersionMax {
		return nil, fmt.Errorf("workercluster: unsupported snapshot version %d", version)
	}

	id := snapshotID(index, term)
	dir := s.store.snapshotDir
	tmp := filepath.Join(dir, id+".json.tmp")
	final := filepath.Join(dir, id+".json")

	f, err := os.Create(tmp)
	if err != nil {
		return nil, fmt.Errorf("workercluster: create snapshot tmp: %w", err)
	}

	sink := &boltSnapshotSink{
		file:  f,
		tmp:   tmp,
		final: final,
		meta: raft.SnapshotMeta{
			Version:            version,
			ID:                 id,
			Index:              index,
			Term:               term,
			Configuration:      configuration,
			ConfigurationIndex: configurationIndex,
		},
	}
	return sink, nil
}

// List returns all stored snapshots in descending index order.
func (s *boltSnapshotStore) List() ([]*raft.SnapshotMeta, error) {
	dir := s.store.snapshotDir
	entries, err := os.ReadDir(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}

	out := make([]*raft.SnapshotMeta, 0, len(entries))
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".json") {
			continue
		}
		rec, err := readSnapshotFile(filepath.Join(dir, e.Name()))
		if err != nil {
			continue
		}
		m := rec.Meta
		out = append(out, &m)
	}

	sort.Slice(out, func(i, j int) bool { return out[i].Index > out[j].Index })
	return out, nil
}

// Open returns a ReadCloser for the snapshot with the given id.
func (s *boltSnapshotStore) Open(id string) (*raft.SnapshotMeta, io.ReadCloser, error) {
	dir := s.store.snapshotDir
	rec, err := readSnapshotFile(filepath.Join(dir, id+".json"))
	if err != nil {
		return nil, nil, fmt.Errorf("workercluster: open snapshot %q: %w", id, err)
	}
	m := rec.Meta
	return &m, io.NopCloser(bytes.NewReader(rec.State)), nil
}

// boltSnapshotSink accumulates bytes written to it and renames the
// temporary file into place on Close. On Cancel the temp file is
// removed.
type boltSnapshotSink struct {
	file  *os.File
	tmp   string
	final string
	meta  raft.SnapshotMeta

	mu    sync.Mutex
	state bytes.Buffer
}

func (s *boltSnapshotSink) Write(p []byte) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.state.Write(p)
}

func (s *boltSnapshotSink) ID() string { return s.meta.ID }

func (s *boltSnapshotSink) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	rec := snapshotRecord{Meta: s.meta, State: s.state.Bytes()}
	raw, err := json.Marshal(rec)
	if err != nil {
		_ = s.file.Close()
		_ = os.Remove(s.tmp)
		return fmt.Errorf("workercluster: encode snapshot: %w", err)
	}
	if err := os.WriteFile(s.final, raw, 0o600); err != nil {
		_ = s.file.Close()
		_ = os.Remove(s.tmp)
		return fmt.Errorf("workercluster: write snapshot: %w", err)
	}
	_ = os.Remove(s.tmp)
	_ = s.file.Close()
	return nil
}

func (s *boltSnapshotSink) Cancel() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	_ = s.file.Close()
	_ = os.Remove(s.tmp)
	return nil
}

func readSnapshotFile(path string) (*snapshotRecord, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var rec snapshotRecord
	if err := json.Unmarshal(raw, &rec); err != nil {
		return nil, err
	}
	return &rec, nil
}

// snapshotID returns a deterministic, sortable snapshot identifier.
// FNV hash of the term+index keeps the file name short while staying
// unique per (index, term).
func snapshotID(index, term uint64) string {
	h := fnv.New64a()
	buf := make([]byte, 16)
	binary.BigEndian.PutUint64(buf[:8], index)
	binary.BigEndian.PutUint64(buf[8:], term)
	_, _ = h.Write(buf)
	return strconv.FormatUint(index, 10) + "-" + strconv.FormatUint(term, 10) + "-" + strconv.FormatUint(h.Sum64(), 16)
}

// encodeLog / decodeLog round-trip a raft.Log entry as msgpack-style
// JSON. The LogStore requires byte-stable encoding so a StoreLog +
// GetLog cycle returns the same values.
func encodeLog(log *raft.Log) ([]byte, error) {
	return json.Marshal(log)
}

func decodeLog(raw []byte, log *raft.Log) error {
	return json.Unmarshal(raw, log)
}