kevo/pkg/transport/replica_persistence.go

package transport

import (
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"sync"
	"time"
)

var (
	// ErrPersistenceDisabled indicates persistence operations cannot be performed
	ErrPersistenceDisabled = errors.New("persistence is disabled")
	
	// ErrInvalidReplicaData indicates the stored replica data is invalid
	ErrInvalidReplicaData = errors.New("invalid replica data")
)

// PersistentReplicaInfo contains replica information that can be persisted
type PersistentReplicaInfo struct {
	ID            string       `json:"id"`
	Address       string       `json:"address"`
	Role          string       `json:"role"`
	LastSeen      int64        `json:"last_seen"`
	CurrentLSN    uint64       `json:"current_lsn"`
	Credentials   *ReplicaCredentials `json:"credentials,omitempty"`
}

// ReplicaPersistence manages persistence of replica information
type ReplicaPersistence struct {
	mu        sync.RWMutex
	dataDir   string
	enabled   bool
	autoSave  bool
	replicas  map[string]*PersistentReplicaInfo
	dirty     bool
	lastSave  time.Time
	saveTimer *time.Timer
}

// NewReplicaPersistence creates a new persistence manager
func NewReplicaPersistence(dataDir string, enabled bool, autoSave bool) (*ReplicaPersistence, error) {
	rp := &ReplicaPersistence{
		dataDir:   dataDir,
		enabled:   enabled,
		autoSave:  autoSave,
		replicas:  make(map[string]*PersistentReplicaInfo),
	}
	
	// Create data directory if it doesn't exist
	if enabled {
		if err := os.MkdirAll(dataDir, 0755); err != nil {
			return nil, err
		}
		
		// Load existing data
		if err := rp.Load(); err != nil {
			return nil, err
		}
		
		// Start auto-save timer if needed
		if autoSave {
			rp.saveTimer = time.AfterFunc(10*time.Second, rp.autoSaveFunc)
		}
	}
	
	return rp, nil
}

// autoSaveFunc is called periodically to save replica data
func (rp *ReplicaPersistence) autoSaveFunc() {
	rp.mu.RLock()
	dirty := rp.dirty
	rp.mu.RUnlock()
	
	if dirty {
		if err := rp.Save(); err != nil {
			// In a production system, this should be logged properly
			println("Error auto-saving replica data:", err.Error())
		}
	}
	
	// Reschedule timer
	rp.saveTimer.Reset(10 * time.Second)
}

// FromReplicaInfo converts a ReplicaInfo to a persistent form
func (rp *ReplicaPersistence) FromReplicaInfo(info *ReplicaInfo, creds *ReplicaCredentials) *PersistentReplicaInfo {
	return &PersistentReplicaInfo{
		ID:         info.ID,
		Address:    info.Address,
		Role:       string(info.Role),
		LastSeen:   info.LastSeen.UnixMilli(),
		CurrentLSN: info.CurrentLSN,
		Credentials: creds,
	}
}

// ToReplicaInfo converts from persistent form to ReplicaInfo
func (rp *ReplicaPersistence) ToReplicaInfo(pinfo *PersistentReplicaInfo) *ReplicaInfo {
	info := &ReplicaInfo{
		ID:         pinfo.ID,
		Address:    pinfo.Address,
		Role:       ReplicaRole(pinfo.Role),
		LastSeen:   time.UnixMilli(pinfo.LastSeen),
		CurrentLSN: pinfo.CurrentLSN,
	}
	return info
}

// Save persists all replica information to disk
func (rp *ReplicaPersistence) Save() error {
	if !rp.enabled {
		return ErrPersistenceDisabled
	}
	
	rp.mu.Lock()
	defer rp.mu.Unlock()
	
	// Nothing to save if no replicas or not dirty
	if len(rp.replicas) == 0 || !rp.dirty {
		return nil
	}
	
	// Save each replica to its own file for better concurrency
	for id, replica := range rp.replicas {
		filename := filepath.Join(rp.dataDir, "replica_"+id+".json")
		
		data, err := json.MarshalIndent(replica, "", "  ")
		if err != nil {
			return err
		}
		
		// Write to temp file first, then rename for atomic update
		tempFile := filename + ".tmp"
		if err := os.WriteFile(tempFile, data, 0644); err != nil {
			return err
		}
		
		if err := os.Rename(tempFile, filename); err != nil {
			return err
		}
	}
	
	rp.dirty = false
	rp.lastSave = time.Now()
	return nil
}

// IsEnabled returns whether persistence is enabled
func (rp *ReplicaPersistence) IsEnabled() bool {
	return rp.enabled
}

// Load reads all persisted replica information from disk
func (rp *ReplicaPersistence) Load() error {
	if !rp.enabled {
		return ErrPersistenceDisabled
	}
	
	rp.mu.Lock()
	defer rp.mu.Unlock()
	
	// Clear existing data
	rp.replicas = make(map[string]*PersistentReplicaInfo)
	
	// Find all replica files
	pattern := filepath.Join(rp.dataDir, "replica_*.json")
	files, err := filepath.Glob(pattern)
	if err != nil {
		return err
	}
	
	// Load each file
	for _, file := range files {
		data, err := os.ReadFile(file)
		if err != nil {
			// Skip files with read errors
			continue
		}
		
		var replica PersistentReplicaInfo
		if err := json.Unmarshal(data, &replica); err != nil {
			// Skip files with parse errors
			continue
		}
		
		// Validate replica data
		if replica.ID == "" {
			continue
		}
		
		rp.replicas[replica.ID] = &replica
	}
	
	rp.dirty = false
	return nil
}

// SaveReplica persists a single replica's information
func (rp *ReplicaPersistence) SaveReplica(info *ReplicaInfo, creds *ReplicaCredentials) error {
	if !rp.enabled {
		return ErrPersistenceDisabled
	}
	
	if info == nil || info.ID == "" {
		return ErrInvalidReplicaData
	}
	
	pinfo := rp.FromReplicaInfo(info, creds)
	
	rp.mu.Lock()
	rp.replicas[info.ID] = pinfo
	rp.dirty = true
	// For immediate save option
	shouldSave := !rp.autoSave
	rp.mu.Unlock()
	
	// Save immediately if auto-save is disabled
	if shouldSave {
		return rp.Save()
	}
	
	return nil
}

// LoadReplica loads a single replica's information
func (rp *ReplicaPersistence) LoadReplica(id string) (*ReplicaInfo, *ReplicaCredentials, error) {
	if !rp.enabled {
		return nil, nil, ErrPersistenceDisabled
	}
	
	if id == "" {
		return nil, nil, ErrInvalidReplicaData
	}
	
	rp.mu.RLock()
	defer rp.mu.RUnlock()
	
	pinfo, exists := rp.replicas[id]
	if !exists {
		return nil, nil, nil // Not found but not an error
	}
	
	return rp.ToReplicaInfo(pinfo), pinfo.Credentials, nil
}

// DeleteReplica removes a replica's persisted information
func (rp *ReplicaPersistence) DeleteReplica(id string) error {
	if !rp.enabled {
		return ErrPersistenceDisabled
	}
	
	if id == "" {
		return ErrInvalidReplicaData
	}
	
	rp.mu.Lock()
	defer rp.mu.Unlock()
	
	// Remove from memory
	delete(rp.replicas, id)
	rp.dirty = true
	
	// Remove file
	filename := filepath.Join(rp.dataDir, "replica_"+id+".json")
	err := os.Remove(filename)
	// Ignore file not found errors
	if err != nil && !os.IsNotExist(err) {
		return err
	}
	
	return nil
}

// GetAllReplicas returns all persisted replicas
func (rp *ReplicaPersistence) GetAllReplicas() (map[string]*ReplicaInfo, map[string]*ReplicaCredentials, error) {
	if !rp.enabled {
		return nil, nil, ErrPersistenceDisabled
	}
	
	rp.mu.RLock()
	defer rp.mu.RUnlock()
	
	infoMap := make(map[string]*ReplicaInfo, len(rp.replicas))
	credsMap := make(map[string]*ReplicaCredentials, len(rp.replicas))
	
	for id, pinfo := range rp.replicas {
		infoMap[id] = rp.ToReplicaInfo(pinfo)
		credsMap[id] = pinfo.Credentials
	}
	
	return infoMap, credsMap, nil
}

// Close shuts down the persistence manager
func (rp *ReplicaPersistence) Close() error {
	if !rp.enabled {
		return nil
	}
	
	// Stop auto-save timer
	if rp.autoSave && rp.saveTimer != nil {
		rp.saveTimer.Stop()
	}
	
	// Save any pending changes
	return rp.Save()
}