EvoBGP/internal/store/memory_snapshot.go

package store

import (
	"context"
	"strings"
	"sync"
	"time"
)

// ModuleSnapshotUpdate computes a replacement from the current snapshot; nil deletes it.
// The callback must not perform store I/O or mutate its input.
type ModuleSnapshotUpdate func(current *ModulePrefixSnapshot) (*ModulePrefixSnapshot, error)

// ListModuleSnapshotSourceRows returns persisted rows for one source key.
func (m *Memory) ListModuleSnapshotSourceRows(ctx context.Context, tenantID, moduleID, source string) ([]PrefixRow, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	snap := m.moduleSnapshots[moduleSnapshotKey(tenantID, moduleID)]
	if snap == nil {
		return nil, nil
	}
	var out []PrefixRow
	for _, row := range snap.Prefixes {
		if row.Source == source {
			out = append(out, row)
		}
	}
	return out, nil
}

func (m *Memory) GetModulePrefixSnapshot(ctx context.Context, tenantID, moduleID string) (*ModulePrefixSnapshot, bool, error) {
	if err := ctx.Err(); err != nil {
		return nil, false, err
	}
	m.mu.RLock()
	defer m.mu.RUnlock()
	key := moduleSnapshotKey(tenantID, moduleID)
	snap, ok := m.moduleSnapshots[key]
	if !ok || snap == nil {
		return nil, false, nil
	}
	cp := make([]PrefixRow, len(snap.Prefixes))
	copy(cp, snap.Prefixes)
	return &ModulePrefixSnapshot{
		InputHash:   snap.InputHash,
		CollectedAt: snap.CollectedAt,
		Version:     snap.Version,
		Prefixes:    cp,
	}, true, nil
}

func (m *Memory) SetModulePrefixSnapshot(ctx context.Context, tenantID, moduleID, inputHash string, prefixes []PrefixRow) error {
	if strings.TrimSpace(tenantID) == "" || strings.TrimSpace(moduleID) == "" || strings.TrimSpace(inputHash) == "" {
		return ErrInvalidInput
	}
	return m.UpdateModulePrefixSnapshot(ctx, tenantID, moduleID, func(*ModulePrefixSnapshot) (*ModulePrefixSnapshot, error) {
		return &ModulePrefixSnapshot{InputHash: inputHash, Prefixes: prefixes}, nil
	})
}

func (m *Memory) DeleteModulePrefixSnapshot(ctx context.Context, tenantID, moduleID string) error {
	return m.UpdateModulePrefixSnapshot(ctx, tenantID, moduleID, func(*ModulePrefixSnapshot) (*ModulePrefixSnapshot, error) {
		return nil, nil
	})
}

// UpdateModulePrefixSnapshot serializes a pure RMW callback and commits snapshot plus input hash together.
func (m *Memory) UpdateModulePrefixSnapshot(ctx context.Context, tenantID, moduleID string, update ModuleSnapshotUpdate) error {
	if update == nil || strings.TrimSpace(tenantID) == "" || strings.TrimSpace(moduleID) == "" {
		return ErrInvalidInput
	}
	unlock, err := m.lockModuleSnapshot(ctx, tenantID, moduleID)
	if err != nil {
		return err
	}
	defer unlock()
	current, _, err := m.GetModulePrefixSnapshot(ctx, tenantID, moduleID)
	if err != nil {
		return err
	}
	next, err := update(current)
	if err != nil {
		return err
	}
	if next != nil && strings.TrimSpace(next.InputHash) == "" {
		return ErrInvalidInput
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	mod := m.modules[moduleID]
	if mod == nil || mod.DeletedAt != nil {
		return ErrNotFound
	}
	if mod.TenantID != tenantID {
		return ErrTenantScope
	}
	key := moduleSnapshotKey(tenantID, moduleID)
	if next == nil {
		if m.moduleSnapshots[key] != nil {
			m.bumpRenderGenerationLocked(tenantID)
		}
		delete(m.moduleSnapshots, key)
	} else if !SameSnapshotContent(current, next) {
		m.bumpRenderGenerationLocked(tenantID)
		if m.moduleSnapshots == nil {
			m.moduleSnapshots = make(map[string]*moduleSnapshotRec)
		}
		m.nextModuleSnapshotVersion++
		m.moduleSnapshots[key] = &moduleSnapshotRec{InputHash: next.InputHash, CollectedAt: time.Now().UTC(), Version: m.nextModuleSnapshotVersion, Prefixes: append([]PrefixRow(nil), next.Prefixes...)}
	}
	if next == nil {
		mod.InputHash = ""
	} else {
		mod.InputHash = next.InputHash
	}
	return nil
}

func (m *Memory) SetModuleInputHash(ctx context.Context, tenantID, moduleID, hash string) error {
	unlock, err := m.lockModuleSnapshot(ctx, tenantID, moduleID)
	if err != nil {
		return err
	}
	defer unlock()
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	mod, ok := m.modules[moduleID]
	if !ok || mod.DeletedAt != nil {
		return ErrNotFound
	}
	if mod.TenantID != tenantID {
		return ErrTenantScope
	}
	mod.InputHash = strings.TrimSpace(hash)
	return nil
}

type snapshotLockEntry struct {
	gate chan struct{}
	refs int
}

// lockModuleSnapshot waits until the caller deadline, or EVOBGP_SNAPSHOT_TIMEOUT when unset.
func (m *Memory) lockModuleSnapshot(ctx context.Context, tenantID, moduleID string) (func(), error) {
	if strings.TrimSpace(tenantID) == "" || strings.TrimSpace(moduleID) == "" {
		return nil, ErrInvalidInput
	}
	if ctx == nil {
		ctx = context.Background()
	}
	ctx, cancel, _ := SnapshotWriteBudget(ctx)
	defer cancel()
	key := moduleSnapshotKey(tenantID, moduleID)
	m.snapshotLocksMu.Lock()
	if m.snapshotLocks == nil {
		m.snapshotLocks = make(map[string]*snapshotLockEntry)
	}
	entry := m.snapshotLocks[key]
	if entry == nil {
		entry = &snapshotLockEntry{gate: make(chan struct{}, 1)}
		entry.gate <- struct{}{}
		m.snapshotLocks[key] = entry
	}
	entry.refs++
	m.snapshotLocksMu.Unlock()
	select {
	case <-ctx.Done():
		m.releaseSnapshotLockRef(key, entry)
		return nil, ctx.Err()
	case <-entry.gate:
		if err := ctx.Err(); err != nil {
			entry.gate <- struct{}{}
			m.releaseSnapshotLockRef(key, entry)
			return nil, err
		}
	}
	var once sync.Once
	return func() {
		once.Do(func() {
			entry.gate <- struct{}{}
			m.releaseSnapshotLockRef(key, entry)
		})
	}, nil
}

func (m *Memory) releaseSnapshotLockRef(key string, entry *snapshotLockEntry) {
	m.snapshotLocksMu.Lock()
	defer m.snapshotLocksMu.Unlock()
	entry.refs--
	if entry.refs == 0 {
		delete(m.snapshotLocks, key)
	}
}

type moduleSnapshotRec struct {
	InputHash   string
	CollectedAt time.Time
	Version     int64
	Prefixes    []PrefixRow
}

// SameSnapshotContent compares persisted rows, ignoring timestamps, versions, and order.
// Rows are matched by (prefix, community, source).
func SameSnapshotContent(a, b *ModulePrefixSnapshot) bool {
	if a == nil || b == nil || a.InputHash != b.InputHash || len(a.Prefixes) != len(b.Prefixes) {
		return false
	}
	ap := append([]PrefixRow(nil), a.Prefixes...)
	bp := append([]PrefixRow(nil), b.Prefixes...)
	sortPrefixRows(ap)
	sortPrefixRows(bp)
	for i := range ap {
		if !samePrefixRow(ap[i], bp[i]) {
			return false
		}
	}
	return true
}

func moduleSnapshotKey(tenantID, moduleID string) string {
	return strings.TrimSpace(tenantID) + "\x00" + strings.TrimSpace(moduleID)
}