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()) }