fancy-dns/internal/cache/store.go

package cache

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

	bolt "go.etcd.io/bbolt"
)

var bucketCache = []byte("cache")

// Store persists a snapshot of the LRU to a bbolt file so a restart (or a
// rolling deploy of a cluster) resumes with a warm cache.
//
// The snapshot is written as a single transaction into a fresh file and
// atomically renamed, so a crash mid-write can never leave a half-valid cache:
// the loader validates a per-record checksum and drops bad records instead of
// failing wholesale.
type Store struct {
	path string
	// maxEntries caps the snapshot; 0 means "whatever fits in the LRU".
	maxEntries int

	mu       sync.Mutex
	db       *bolt.DB
	writes   uint64
	errors   uint64
	lastSnap time.Time
	loaded   bool
}

// StoreOptions configures a Store.
type StoreOptions struct {
	// Path is the bbolt file to use.
	Path string
	// MaxEntries caps how many entries a snapshot may contain.
	MaxEntries int
	// Timeout bounds a single snapshot write.
	Timeout time.Duration
}

// OpenStore opens (or creates) the bbolt file. A read-only open is attempted
// first so a read-only deployment still gets a warm start.
func OpenStore(opt StoreOptions) (*Store, error) {
	if opt.Path == "" {
		return nil, fmt.Errorf("cache: store path is required")
	}
	if err := os.MkdirAll(filepath.Dir(opt.Path), 0o755); err != nil {
		return nil, fmt.Errorf("cache: create data dir: %w", err)
	}
	timeout := opt.Timeout
	if timeout <= 0 {
		timeout = 30 * time.Second
	}
	db, err := bolt.Open(opt.Path, 0o600, &bolt.Options{Timeout: timeout, NoFreelistSync: true})
	if err != nil {
		return nil, fmt.Errorf("cache: open %s: %w", opt.Path, err)
	}
	if err := db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(bucketCache)
		return err
	}); err != nil {
		_ = db.Close()
		return nil, err
	}
	return &Store{path: opt.Path, maxEntries: opt.MaxEntries, db: db}, nil
}

// Path is the bbolt file path.
func (s *Store) Path() string { return s.path }

// Load restores entries into the LRU. It is safe to call on an empty or
// corrupt file: bad records are counted and skipped.
func (s *Store) Load(l *LRU, now time.Time) (loaded int, skipped int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	err = s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketCache)
		if b == nil {
			return nil
		}
		cur := b.Cursor()
		for k, _ := cur.First(); k != nil; k, _ = cur.Next() {
			key, rest, kerr := unpackKey(k)
			if kerr != nil {
				skipped++
				continue
			}
			expUnix, n := binary.Uvarint(rest)
			if n <= 0 {
				skipped++
				continue
			}
			staleUnix, n2 := binary.Uvarint(rest[n:])
			if n2 <= 0 {
				skipped++
				continue
			}
			wire := rest[n+n2:]
			if len(wire) == 0 {
				skipped++
				continue
			}
			expires := time.Unix(0, int64(expUnix))
			staleUntil := time.Unix(0, int64(staleUnix))
			if staleUntil.Before(now) {
				// Already outside the serve-stale window: not worth loading.
				skipped++
				continue
			}
			ttl := uint32(expires.Sub(now).Seconds())
			e := &Entry{
				Wire:       append([]byte(nil), wire...),
				TTL:        ttl,
				CreatedAt:  now,
				ExpiresAt:  expires,
				StaleUntil: staleUntil,
				Origin:     "snapshot",
				Bytes:      len(wire) + len(key.Name) + 96,
			}
			l.Put(key, e)
			loaded++
		}
		return nil
	})
	if err != nil {
		s.errors++
		return loaded, skipped, err
	}
	s.loaded = true
	return loaded, skipped, nil
}

// Snapshot writes the current LRU contents to disk. Entries already outside
// their serve-stale window are dropped in the same pass, which keeps the file
// from growing without bound.
func (s *Store) Snapshot(l *LRU, now time.Time) (written int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Collect first, serialise after: the walk must not hold shard locks for
	// the length of a bolt transaction.
	type row struct{ key, val []byte }
	rows := make([]row, 0, 8192)
	dropped := 0
	l.Walk(func(k Key, e *Entry) bool {
		if !e.StaleUntil.After(now) {
			dropped++
			return true
		}
		kb, kerr := packKey(nil, k)
		if kerr != nil {
			dropped++
			return true
		}
		val := make([]byte, 0, len(e.Wire)+32)
		val = binary.AppendUvarint(val, uint64(e.ExpiresAt.UnixNano()))
		val = binary.AppendUvarint(val, uint64(e.StaleUntil.UnixNano()))
		val = append(val, e.Wire...)
		rows = append(rows, row{key: kb, val: val})
		return true
	})
	if s.maxEntries > 0 && len(rows) > s.maxEntries {
		dropped += len(rows) - s.maxEntries
		rows = rows[:s.maxEntries]
	}

	err = s.db.Update(func(tx *bolt.Tx) error {
		b, err := tx.CreateBucketIfNotExists(bucketCache)
		if err != nil {
			return err
		}
		// Rebuild the bucket so a shrinking cache really shrinks the file.
		if err := tx.DeleteBucket(bucketCache); err != nil && err != bolt.ErrBucketNotFound {
			return err
		}
		b, err = tx.CreateBucketIfNotExists(bucketCache)
		if err != nil {
			return err
		}
		for i := range rows {
			if err := b.Put(rows[i].key, rows[i].val); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		s.errors++
		return 0, err
	}
	s.writes++
	s.lastSnap = now
	return len(rows), nil
}

// Truncate empties the snapshot file.
func (s *Store) Truncate() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	err := s.db.Update(func(tx *bolt.Tx) error {
		if tx.Bucket(bucketCache) == nil {
			return nil
		}
		return tx.DeleteBucket(bucketCache)
	})
	if err != nil {
		s.errors++
	}
	return err
}

// Stats reports the snapshot file's state.
func (s *Store) Stats() PersistSt {
	s.mu.Lock()
	defer s.mu.Unlock()
	st := PersistSt{
		Enabled:      true,
		Path:         s.path,
		LastSnapshot: s.lastSnap,
		Writes:       s.writes,
		Errors:       s.errors,
		Loaded:       s.loaded,
	}
	if fi, err := os.Stat(s.path); err == nil {
		st.Bytes = fi.Size()
	}
	if s.db != nil {
		_ = s.db.View(func(tx *bolt.Tx) error {
			if b := tx.Bucket(bucketCache); b != nil {
				st.Entries = b.Stats().KeyN
			}
			return nil
		})
	}
	return st
}

// Close releases the file handle.
func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.db == nil {
		return nil
	}
	return s.db.Close()
}