pangolin/server/internal/scheduler/probe/store.go

package probe

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/redis/go-redis/v9"
)

// Redis TTL constants for the probe subsystem.
const (
	// snapshotTTL is how long a per-(node,vantage) probe snapshot is retained.
	// After this window the key expires and 15D must treat it as "no data".
	snapshotTTL = 30 * time.Minute

	// heartbeatTTL is how long a probe heartbeat key lives without renewal.
	// After this window the key expires — meaning "no recent data", NOT "failure".
	heartbeatTTL = 15 * time.Minute

	// seenTTL is the lifetime of a replay-prevention key.
	// Must exceed 2 × timeWindow (2 × 300 s = 600 s); 700 s adds a 100 s buffer.
	seenTTL = 700 * time.Second
)

// Store handles Redis reads and writes for the probe subsystem.
//
// # Key schema
//
//   probe:{nodeID}:{vantageKey}   → ProbeSnapshot JSON, TTL snapshotTTL (30 min)
//   probe:hb:{probeID}            → Unix timestamp string, TTL heartbeatTTL (15 min)
//   probe:seen:{probeID}:{ts}     → "1", TTL seenTTL (700 s), replay prevention
//
// # Missing-key semantics (IMPORTANT – must not be violated by callers)
//
//   A missing probe:hb:{probeID} key means "no recent heartbeat data".
//   It must NEVER be interpreted as "probe is down" or folded into a failure
//   signal.  Determination logic (15D) must treat an absent key as unknown,
//   not as a negative result.
//
// # Node-ID constraint
//
//   NodeIDs must not be the literal string "hb" or "seen", as those are used as
//   key-namespace prefixes.  In practice node IDs are UUIDs so this is safe.
type Store struct {
	rdb *redis.Client
}

// NewStore creates a Store backed by the given Redis client.
func NewStore(rdb *redis.Client) *Store {
	return &Store{rdb: rdb}
}

// --------------------------------------------------------------------------
// Key helpers
// --------------------------------------------------------------------------

// snapshotKey returns the Redis key for the latest probe snapshot
// for the given (nodeID, vantage) pair.
//
// Format: probe:{nodeID}:{country}:{region}:{isp}
func snapshotKey(nodeID string, v VantagePoint) string {
	return "probe:" + nodeID + ":" + vantageKey(v)
}

// vantageKey returns a canonical, colon-safe string for a VantagePoint.
// It is used as the trailing segment of a snapshot Redis key.
func vantageKey(v VantagePoint) string {
	return sanitizeKeySegment(v.Country) + ":" +
		sanitizeKeySegment(v.Region) + ":" +
		sanitizeKeySegment(v.ISP)
}

// sanitizeKeySegment replaces characters that would interfere with Redis key
// parsing (space, colon) with underscores so they are safe in compound keys.
func sanitizeKeySegment(s string) string {
	s = strings.ReplaceAll(s, " ", "_")
	s = strings.ReplaceAll(s, ":", "_")
	return s
}

// heartbeatKey returns the Redis key for a probe agent heartbeat.
func heartbeatKey(probeID string) string {
	return "probe:hb:" + probeID
}

// seenKey returns the Redis key used for replay-prevention on a
// (probeID, ts) tuple.
func seenKey(probeID, ts string) string {
	return "probe:seen:" + probeID + ":" + ts
}

// --------------------------------------------------------------------------
// Write path
// --------------------------------------------------------------------------

// SaveReports persists one ProbeSnapshot per (nodeID, vantage) pair from the
// given report batch and updates the probe heartbeat.
// All writes are issued in a single pipeline for efficiency.
func (s *Store) SaveReports(ctx context.Context, probeID string, vantage VantagePoint, reports []NodeReport) error {
	now := time.Now().Unix()
	pipe := s.rdb.Pipeline()

	for _, rep := range reports {
		snap := ProbeSnapshot{
			ProbeID:    probeID,
			Vantage:    vantage,
			Report:     rep,
			ReceivedAt: now,
		}
		data, err := json.Marshal(snap)
		if err != nil {
			return fmt.Errorf("probe: marshal snapshot for node %s: %w", rep.NodeID, err)
		}
		pipe.Set(ctx, snapshotKey(rep.NodeID, vantage), data, snapshotTTL)
	}

	// Update heartbeat: value is the ingest timestamp so readers can compute
	// staleness without needing a separate TTL query.
	pipe.Set(ctx, heartbeatKey(probeID), strconv.FormatInt(now, 10), heartbeatTTL)

	if _, err := pipe.Exec(ctx); err != nil {
		return fmt.Errorf("probe: redis pipeline exec: %w", err)
	}
	return nil
}

// CheckAndMarkSeen atomically checks whether the (probeID, ts) pair has been
// seen before, and marks it as seen if not.
//
//   Returns (true, nil)  – this is a replay; the caller should handle idempotently.
//   Returns (false, nil) – first time seen; the caller should process normally.
func (s *Store) CheckAndMarkSeen(ctx context.Context, probeID, ts string) (bool, error) {
	key := seenKey(probeID, ts)
	// SET … NX EX: atomically set iff the key does not exist.
	set, err := s.rdb.SetNX(ctx, key, "1", seenTTL).Result()
	if err != nil {
		return false, fmt.Errorf("probe: seen check: %w", err)
	}
	// SetNX returns true if the key was newly created (not a replay).
	return !set, nil
}

// --------------------------------------------------------------------------
// Read interfaces consumed by 15D (determination logic)
// --------------------------------------------------------------------------

// SnapshotsByNode returns all cached probe snapshots for the given nodeID,
// indexed by vantage key string (Country:Region:ISP).
//
// A missing key — and therefore an empty map — means "no recent probe data
// from any vantage point".  Callers (15D) must treat this as "unknown", not
// as a failure signal.
//
// Uses SCAN + GET in two passes.  For the expected scale (dozens of vantages
// per node) this is efficient enough; 15D may add caching on top if needed.
func (s *Store) SnapshotsByNode(ctx context.Context, nodeID string) (map[string]ProbeSnapshot, error) {
	pattern := "probe:" + nodeID + ":*"

	var keys []string
	iter := s.rdb.Scan(ctx, 0, pattern, 0).Iterator()
	for iter.Next(ctx) {
		keys = append(keys, iter.Val())
	}
	if err := iter.Err(); err != nil {
		return nil, fmt.Errorf("probe: scan snapshots for node %s: %w", nodeID, err)
	}

	result := make(map[string]ProbeSnapshot, len(keys))
	prefix := "probe:" + nodeID + ":"

	for _, k := range keys {
		val, err := s.rdb.Get(ctx, k).Result()
		if err == redis.Nil {
			// Key expired between SCAN and GET – not an error.
			continue
		}
		if err != nil {
			return nil, fmt.Errorf("probe: get snapshot %s: %w", k, err)
		}
		var snap ProbeSnapshot
		if err := json.Unmarshal([]byte(val), &snap); err != nil {
			// Corrupted data – log-worthy but non-fatal; skip this entry.
			continue
		}
		vk := strings.TrimPrefix(k, prefix)
		result[vk] = snap
	}
	return result, nil
}

// AliveProbes returns the IDs of probe agents that have sent a heartbeat
// within the last heartbeatTTL window (15 min).
//
// An empty result means "no probes have reported recently" – not that all
// probes are down.  Missing heartbeat keys must never be folded into a
// failure signal by callers.
func (s *Store) AliveProbes(ctx context.Context) ([]string, error) {
	var ids []string
	iter := s.rdb.Scan(ctx, 0, "probe:hb:*", 0).Iterator()
	for iter.Next(ctx) {
		k := iter.Val()
		ids = append(ids, strings.TrimPrefix(k, "probe:hb:"))
	}
	if err := iter.Err(); err != nil {
		return nil, fmt.Errorf("probe: scan heartbeats: %w", err)
	}
	return ids, nil
}