browser/plugins/examples/scum-server-plugin/companion/snapshot_sequences.go

package companion

import (
	"encoding/json"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const SnapshotSequenceFilename = "snapshot-sequences.json"

// SnapshotSequenceStore is plugin-local state used only to keep the plugin's
// typed snapshot stream monotonic. It never stores or examines raw log text.
type SnapshotSequenceStore interface {
	Next(snapshotType string, streamKey string, floor uint64) (uint64, error)
}

type MemorySnapshotSequenceStore struct {
	mu        sync.Mutex
	sequences map[string]uint64
}

func NewMemorySnapshotSequenceStore() *MemorySnapshotSequenceStore {
	return &MemorySnapshotSequenceStore{sequences: map[string]uint64{}}
}

func (store *MemorySnapshotSequenceStore) Next(snapshotType string, streamKey string, floor uint64) (uint64, error) {
	if store == nil {
		return 0, fmt.Errorf("snapshot sequence store is not configured")
	}
	key, err := snapshotSequenceKey(snapshotType, streamKey)
	if err != nil {
		return 0, err
	}
	store.mu.Lock()
	defer store.mu.Unlock()
	next, err := nextSnapshotSequence(store.sequences[key], floor)
	if err != nil {
		return 0, err
	}
	store.sequences[key] = next
	return next, nil
}

type FileSnapshotSequenceStore struct {
	filename string
	mu       sync.Mutex
}

type snapshotSequenceDocument struct {
	SchemaVersion int               `json:"schemaVersion"`
	Sequences     map[string]uint64 `json:"sequences"`
}

func NewFileSnapshotSequenceStore(filename string) (*FileSnapshotSequenceStore, error) {
	filename = strings.TrimSpace(filename)
	if filename == "" || filepath.Base(filename) != filename || filename == "." {
		return nil, fmt.Errorf("snapshot sequence filename is invalid")
	}
	return &FileSnapshotSequenceStore{filename: filename}, nil
}

func (store *FileSnapshotSequenceStore) Next(snapshotType string, streamKey string, floor uint64) (uint64, error) {
	if store == nil {
		return 0, fmt.Errorf("snapshot sequence store is not configured")
	}
	key, err := snapshotSequenceKey(snapshotType, streamKey)
	if err != nil {
		return 0, err
	}
	store.mu.Lock()
	defer store.mu.Unlock()
	document, err := store.read()
	if err != nil {
		return 0, err
	}
	next, err := nextSnapshotSequence(document.Sequences[key], floor)
	if err != nil {
		return 0, err
	}
	document.Sequences[key] = next
	if err := store.write(document); err != nil {
		return 0, err
	}
	return next, nil
}

func (store *FileSnapshotSequenceStore) read() (snapshotSequenceDocument, error) {
	document := snapshotSequenceDocument{SchemaVersion: 1, Sequences: map[string]uint64{}}
	body, err := os.ReadFile(store.filename)
	if os.IsNotExist(err) {
		return document, nil
	}
	if err != nil {
		return snapshotSequenceDocument{}, fmt.Errorf("read snapshot sequences: %w", err)
	}
	if err := json.Unmarshal(body, &document); err != nil {
		return snapshotSequenceDocument{}, fmt.Errorf("decode snapshot sequences: %w", err)
	}
	if document.SchemaVersion != 1 || document.Sequences == nil {
		return snapshotSequenceDocument{}, fmt.Errorf("snapshot sequence state is invalid")
	}
	return document, nil
}

func (store *FileSnapshotSequenceStore) write(document snapshotSequenceDocument) error {
	body, err := json.Marshal(document)
	if err != nil {
		return err
	}
	temporary, err := os.CreateTemp(".", ".snapshot-sequences-")
	if err != nil {
		return fmt.Errorf("create snapshot sequence state: %w", err)
	}
	temporaryName := temporary.Name()
	defer os.Remove(temporaryName)
	if _, err := temporary.Write(body); err != nil {
		temporary.Close()
		return fmt.Errorf("write snapshot sequence state: %w", err)
	}
	if err := temporary.Chmod(0o600); err != nil {
		temporary.Close()
		return fmt.Errorf("protect snapshot sequence state: %w", err)
	}
	if err := temporary.Close(); err != nil {
		return fmt.Errorf("close snapshot sequence state: %w", err)
	}
	if err := os.Rename(temporaryName, store.filename); err != nil {
		return fmt.Errorf("commit snapshot sequence state: %w", err)
	}
	return nil
}

func snapshotSequenceKey(snapshotType string, streamKey string) (string, error) {
	snapshotType = strings.TrimSpace(snapshotType)
	streamKey = strings.TrimSpace(streamKey)
	if snapshotType == "" || streamKey == "" || len(snapshotType) > 80 || len(streamKey) > 80 {
		return "", fmt.Errorf("snapshot stream identity is invalid")
	}
	return snapshotType + "\x00" + streamKey, nil
}

func nextSnapshotSequence(current uint64, floor uint64) (uint64, error) {
	if current == math.MaxUint64 {
		return 0, fmt.Errorf("snapshot sequence is exhausted")
	}
	next := current + 1
	if floor > next {
		next = floor
	}
	if next == 0 {
		return 0, fmt.Errorf("snapshot sequence is exhausted")
	}
	return next, nil
}

func snapshotSequenceFloor(observedAt time.Time) uint64 {
	if observedAt.IsZero() || observedAt.UnixNano() <= 0 {
		return 1
	}
	return uint64(observedAt.UnixNano())
}