wx-win-agent/control-plane/store.go

package controlplane

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"
)

// StoreOptions controls durability and bounded retention for the JSON store.
// A Store still has one active writer; the lock file enables active/passive failover.
type StoreOptions struct {
	BackupDir      string
	BackupCount    int
	BackupInterval time.Duration
	TaskRetention  time.Duration
	EventRetention time.Duration
	AuditRetention time.Duration
}

// Store is a small durable JSON store for the control-plane MVP.
// The file is an implementation detail; callers only observe transactional methods.
type Store struct {
	path       string
	lockFile   *os.File
	options    StoreOptions
	lastBackup time.Time
	mu         sync.Mutex
	state      PersistedState
	closed     bool
}

func OpenStore(path string, options ...StoreOptions) (*Store, error) {
	if path == "" {
		return nil, errors.New("data file is required")
	}
	cleanPath := filepath.Clean(path)
	storeOptions := StoreOptions{}
	if len(options) > 0 {
		storeOptions = options[0]
	}
	if storeOptions.BackupDir == "" {
		storeOptions.BackupDir = cleanPath + ".backups"
	}
	if storeOptions.BackupCount == 0 {
		storeOptions.BackupCount = 7
	}
	if storeOptions.BackupCount < -1 {
		return nil, errors.New("backup count must be -1 or non-negative")
	}
	if err := os.MkdirAll(filepath.Dir(cleanPath), 0o700); err != nil {
		return nil, fmt.Errorf("create data directory: %w", err)
	}
	lockFile, err := os.OpenFile(cleanPath+".lock", os.O_CREATE|os.O_RDWR, 0o600)
	if err != nil {
		return nil, fmt.Errorf("open data lock: %w", err)
	}
	if err := syscall.Flock(int(lockFile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil {
		_ = lockFile.Close()
		if errors.Is(err, syscall.EWOULDBLOCK) {
			return nil, fmt.Errorf("data file is already in use: %s", cleanPath)
		}
		return nil, fmt.Errorf("lock data file: %w", err)
	}

	store := &Store{
		path:     cleanPath,
		lockFile: lockFile,
		options:  storeOptions,
		state:    PersistedState{Nodes: map[string]Node{}, Tasks: map[string]Task{}, Events: []StoredEvent{}, Audit: []AuditEntry{}, AIFlows: map[string]AIFlow{}, AIRuns: map[string]AIRun{}, AIPullTasks: map[string]AIPullTask{}, AIMessageKeys: map[string]string{}},
	}
	data, err := os.ReadFile(store.path)
	if errors.Is(err, os.ErrNotExist) {
		return store, nil
	}
	if err != nil {
		_ = store.Close()
		return nil, fmt.Errorf("read data file: %w", err)
	}
	if len(data) == 0 {
		return store, nil
	}
	if err := json.Unmarshal(data, &store.state); err != nil {
		_ = store.Close()
		return nil, fmt.Errorf("decode data file: %w", err)
	}
	store.ensureMaps()
	return store, nil
}

func (s *Store) Snapshot() PersistedState {
	s.mu.Lock()
	defer s.mu.Unlock()
	return cloneState(s.state)
}

func (s *Store) Mutate(fn func(*PersistedState) error) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return errors.New("store is closed")
	}
	s.ensureMaps()
	if err := fn(&s.state); err != nil {
		return err
	}
	return s.saveLocked()
}

func (s *Store) Read(fn func(PersistedState) error) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return errors.New("store is closed")
	}
	return fn(cloneState(s.state))
}

// Close releases the single-writer lock. A second control-plane process can then
// be promoted by the supervisor using the same data directory.
func (s *Store) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil
	}
	s.closed = true
	if s.lockFile == nil {
		return nil
	}
	unlockErr := syscall.Flock(int(s.lockFile.Fd()), syscall.LOCK_UN)
	closeErr := s.lockFile.Close()
	if unlockErr != nil {
		return fmt.Errorf("unlock data file: %w", unlockErr)
	}
	if closeErr != nil {
		return fmt.Errorf("close data lock: %w", closeErr)
	}
	return nil
}

func (s *Store) ensureMaps() {
	if s.state.Nodes == nil {
		s.state.Nodes = map[string]Node{}
	}
	if s.state.Tasks == nil {
		s.state.Tasks = map[string]Task{}
	}
	if s.state.Events == nil {
		s.state.Events = []StoredEvent{}
	}
	if s.state.Audit == nil {
		s.state.Audit = []AuditEntry{}
	}
	if s.state.AIFlows == nil {
		s.state.AIFlows = map[string]AIFlow{}
	}
	if s.state.AIRuns == nil {
		s.state.AIRuns = map[string]AIRun{}
	}
	if s.state.AIPullTasks == nil {
		s.state.AIPullTasks = map[string]AIPullTask{}
	}
	if s.state.AIMessageKeys == nil {
		s.state.AIMessageKeys = map[string]string{}
	}
}

func (s *Store) saveLocked() error {
	if err := os.MkdirAll(filepath.Dir(s.path), 0o700); err != nil {
		return fmt.Errorf("create data directory: %w", err)
	}
	s.pruneLocked(time.Now().UTC())
	if err := s.backupLocked(); err != nil {
		return err
	}
	data, err := json.MarshalIndent(s.state, "", "  ")
	if err != nil {
		return fmt.Errorf("encode data file: %w", err)
	}
	temporary, err := os.CreateTemp(filepath.Dir(s.path), ".wxagent-control-plane-*")
	if err != nil {
		return fmt.Errorf("create temporary data file: %w", err)
	}
	temporaryName := temporary.Name()
	defer os.Remove(temporaryName)
	if err := temporary.Chmod(0o600); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("protect temporary data file: %w", err)
	}
	if _, err := temporary.Write(data); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("write data file: %w", err)
	}
	if err := temporary.Sync(); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("sync data file: %w", err)
	}
	if err := temporary.Close(); err != nil {
		return fmt.Errorf("close data file: %w", err)
	}
	if err := os.Rename(temporaryName, s.path); err != nil {
		return fmt.Errorf("replace data file: %w", err)
	}
	return syncDirectory(filepath.Dir(s.path))
}

func (s *Store) pruneLocked(now time.Time) {
	if s.options.TaskRetention > 0 {
		cutoff := now.Add(-s.options.TaskRetention)
		for id, task := range s.state.Tasks {
			if terminal(task.Status) && !task.UpdatedAt.IsZero() && task.UpdatedAt.Before(cutoff) {
				delete(s.state.Tasks, id)
			}
		}
	}
	if s.options.EventRetention > 0 {
		cutoff := now.Add(-s.options.EventRetention)
		kept := s.state.Events[:0]
		for _, event := range s.state.Events {
			if event.ReceivedAt.IsZero() || !event.ReceivedAt.Before(cutoff) {
				kept = append(kept, event)
			}
		}
		s.state.Events = kept
	}
	if s.options.AuditRetention > 0 {
		cutoff := now.Add(-s.options.AuditRetention)
		kept := s.state.Audit[:0]
		for _, entry := range s.state.Audit {
			if entry.At.IsZero() || !entry.At.Before(cutoff) {
				kept = append(kept, entry)
			}
		}
		s.state.Audit = kept
	}
}

func (s *Store) backupLocked() error {
	if s.options.BackupCount < 1 {
		return nil
	}
	now := time.Now().UTC()
	if s.options.BackupInterval > 0 && !s.lastBackup.IsZero() && now.Sub(s.lastBackup) < s.options.BackupInterval {
		return nil
	}
	data, err := os.ReadFile(s.path)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("read current data for backup: %w", err)
	}
	if err := os.MkdirAll(s.options.BackupDir, 0o700); err != nil {
		return fmt.Errorf("create backup directory: %w", err)
	}
	name := fmt.Sprintf("%s.%d.json", filepath.Base(s.path), time.Now().UTC().UnixNano())
	backupPath := filepath.Join(s.options.BackupDir, name)
	temporary, err := os.CreateTemp(s.options.BackupDir, ".wxagent-backup-*")
	if err != nil {
		return fmt.Errorf("create backup file: %w", err)
	}
	temporaryName := temporary.Name()
	defer os.Remove(temporaryName)
	if err := temporary.Chmod(0o600); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("protect backup file: %w", err)
	}
	if _, err := temporary.Write(data); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("write backup file: %w", err)
	}
	if err := temporary.Sync(); err != nil {
		_ = temporary.Close()
		return fmt.Errorf("sync backup file: %w", err)
	}
	if err := temporary.Close(); err != nil {
		return fmt.Errorf("close backup file: %w", err)
	}
	if err := os.Rename(temporaryName, backupPath); err != nil {
		return fmt.Errorf("publish backup file: %w", err)
	}
	if err := pruneBackups(s.options.BackupDir, filepath.Base(s.path), s.options.BackupCount); err != nil {
		return fmt.Errorf("prune backups: %w", err)
	}
	if err := syncDirectory(s.options.BackupDir); err != nil {
		return err
	}
	s.lastBackup = now
	return nil
}

func pruneBackups(directory, base string, keep int) error {
	entries, err := os.ReadDir(directory)
	if err != nil {
		return err
	}
	type backup struct {
		name string
		when time.Time
	}
	backups := make([]backup, 0, len(entries))
	prefix := base + "."
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasPrefix(entry.Name(), prefix) || !strings.HasSuffix(entry.Name(), ".json") {
			continue
		}
		info, err := entry.Info()
		if err != nil {
			return err
		}
		backups = append(backups, backup{name: entry.Name(), when: info.ModTime()})
	}
	sort.Slice(backups, func(i, j int) bool {
		if backups[i].when.Equal(backups[j].when) {
			return backups[i].name > backups[j].name
		}
		return backups[i].when.After(backups[j].when)
	})
	if len(backups) <= keep {
		return nil
	}
	for _, item := range backups[keep:] {
		if err := os.Remove(filepath.Join(directory, item.name)); err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}
	}
	return nil
}

func syncDirectory(path string) error {
	directory, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("open directory for sync: %w", err)
	}
	defer directory.Close()
	if err := directory.Sync(); err != nil && !errors.Is(err, syscall.EINVAL) {
		return fmt.Errorf("sync directory: %w", err)
	}
	return nil
}

func cloneState(source PersistedState) PersistedState {
	data, _ := json.Marshal(source)
	var copy PersistedState
	_ = json.Unmarshal(data, &copy)
	if copy.Nodes == nil {
		copy.Nodes = map[string]Node{}
	}
	if copy.Tasks == nil {
		copy.Tasks = map[string]Task{}
	}
	if copy.Events == nil {
		copy.Events = []StoredEvent{}
	}
	if copy.Audit == nil {
		copy.Audit = []AuditEntry{}
	}
	return copy
}