pangolin/server/internal/scheduler/scheduler.go

package scheduler

import (
	"context"
	"log/slog"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

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

// ─────────────────────────────────────────────────────────────────────────────
// Leader-election Redis key helper
// ─────────────────────────────────────────────────────────────────────────────

// leaderKey returns the Redis key for a loop's leader election.
// Format: sched:leader:{loopName}
func leaderKey(loop string) string { return "sched:leader:" + loop }

// ─────────────────────────────────────────────────────────────────────────────
// Config
// ─────────────────────────────────────────────────────────────────────────────

// Config holds all scheduler dependencies and tunable timing parameters.
// Zero values for duration fields use the production defaults listed in the
// field comments.  All wired components (Engine, Replacer, …) are required
// unless the field is explicitly marked optional.
type Config struct {
	// RDB is the Redis client used for leader-election keys.  Required.
	RDB *redis.Client

	// InstanceID uniquely identifies this process instance across the fleet.
	// If empty, a random UUID is generated by New.
	InstanceID string

	// ── Wired engine components ──────────────────────────────────────────────

	// Engine is the 15D detection engine.  Required.
	Engine DetectEngine

	// Replacer is the 15E replacement-orchestration Replacer.  Required.
	Replacer OrchestrateReplacer

	// Grayscale is the 15E grayscale warm-up driver.  Required.
	Grayscale OrchestrateGrayscale

	// Capacity is the 15F capacity/circuit-breaker service.
	// Optional — StubCapacityService is used when nil.
	Capacity CapacityService

	// ProbeStore reads probe-agent heartbeat liveness (15A).
	// Optional — probe disconnection checking is disabled when nil.
	ProbeStore ProbeStateReader

	// KnownProbeIDs is the set of probe-agent IDs expected to be alive.
	// When any ID in this list is absent from ProbeStore.AliveProbes, the
	// Notifier is called with reason "probe_disconnected".
	KnownProbeIDs []string

	// Prober is the 15C third-party dial-test driver.
	// Optional — 15C probing is disabled when nil.
	Prober ThirdPartyProber

	// Targets supplies the list of nodes to probe via Prober.
	// Required when Prober is non-nil.
	Targets TargetProvider

	// Notifier is the 15G alerting sink.
	// Optional — LogNotifier is used when nil.
	Notifier Notifier

	// ── Tick intervals (hot-reloadable from 15F config) ─────────────────────
	// Zero → production default (shown in comments).

	DetectInterval      time.Duration // 5 min
	OrchestrateInterval time.Duration // 30 s
	CapacityInterval    time.Duration // 90 s
	ThirdPartyInterval  time.Duration // 10 min — interval between 15C RunOnce calls

	// ── Leader-election timing ────────────────────────────────────────────────
	// Override only in tests (shorter values speed up convergence).

	LeaseTTL    time.Duration // 30 s  — Redis PEXPIRE on the leader key
	RenewPeriod time.Duration // 10 s  — how often the leader renews its lease
	RetryPeriod time.Duration // 5 s   — how often a non-leader tries to acquire
	TickTimeout time.Duration // 30 s  — per-tick context hard-timeout
}

// applyDefaults fills zero-value duration fields with production defaults.
func (c *Config) applyDefaults() {
	if c.InstanceID == "" {
		c.InstanceID = uuid.NewString()
	}
	if c.DetectInterval == 0 {
		c.DetectInterval = 5 * time.Minute
	}
	if c.OrchestrateInterval == 0 {
		c.OrchestrateInterval = 30 * time.Second
	}
	if c.CapacityInterval == 0 {
		c.CapacityInterval = 90 * time.Second
	}
	if c.ThirdPartyInterval == 0 {
		c.ThirdPartyInterval = 10 * time.Minute
	}
	if c.LeaseTTL == 0 {
		c.LeaseTTL = 30 * time.Second
	}
	if c.RenewPeriod == 0 {
		c.RenewPeriod = 10 * time.Second
	}
	if c.RetryPeriod == 0 {
		c.RetryPeriod = 5 * time.Second
	}
	if c.TickTimeout == 0 {
		c.TickTimeout = 30 * time.Second
	}
	if c.Notifier == nil {
		c.Notifier = LogNotifier{}
	}
	if c.Capacity == nil {
		c.Capacity = StubCapacityService{}
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Scheduler
// ─────────────────────────────────────────────────────────────────────────────

// Scheduler runs three leader-elected loops (DetectLoop / OrchestrateLoop /
// CapacityLoop) plus an optional unguarded 15C third-party-prober goroutine.
//
// Design rationale for three independent leader leases (one per loop):
//   - A slow Detect tick (e.g., DB latency spike) must not block Orchestrate
//     progress on other in-flight replacements.
//   - Each loop failing over independently avoids a partial-leader split where
//     a crashing replica holds only one of three leases.
//   - Raft-style leadership is overkill for these loosely coupled loops, which
//     are each idempotent and use Redis-side locking for concurrency safety.
type Scheduler struct {
	rdb  *redis.Client
	cfg  Config
}

// New creates a Scheduler from cfg, applying production defaults for any nil /
// zero fields.
func New(cfg Config) *Scheduler {
	cfg.applyDefaults()
	return &Scheduler{rdb: cfg.RDB, cfg: cfg}
}

// Run starts the scheduler with OS signal handling (SIGTERM / SIGINT) and
// blocks until the scheduler has shut down.  For production use from main().
func (s *Scheduler) Run(ctx context.Context) error {
	sigCtx, stop := signal.NotifyContext(ctx, os.Interrupt, syscall.SIGTERM)
	defer stop()
	return s.Start(sigCtx)
}

// Start starts the scheduler goroutines and blocks until ctx is cancelled.
// Unlike Run, it does not install OS signal handlers — useful for testing.
func (s *Scheduler) Start(ctx context.Context) error {
	var wg sync.WaitGroup

	// ── Three leader-elected loops ─────────────────────────────────────────────
	//
	// Each loop acquires its own independent Redis leader key so that a slow
	// tick in one loop does not delay the others.

	wg.Add(3)
	go func() {
		defer wg.Done()
		s.runLoop(ctx, "detect", s.cfg.DetectInterval, s.cfg.Engine.Tick)
	}()
	go func() {
		defer wg.Done()
		s.runLoop(ctx, "orchestrate", s.cfg.OrchestrateInterval, s.cfg.Replacer.Tick)
	}()
	go func() {
		defer wg.Done()
		s.runLoop(ctx, "capacity", s.cfg.CapacityInterval, s.capacityTick)
	}()

	// ── 15C third-party prober (no leader constraint) ─────────────────────────
	//
	// Design choice: run on ALL replicas without leader election because:
	//   1. Alibaba Cloud API calls are idempotent: same (node, ISP) pair creates
	//      a new one-shot task; the result is written to Redis via SET (overwrite)
	//      so duplicate writes produce the same snapshot with a slightly later
	//      ReceivedAt — harmless.
	//   2. Running on every replica improves data collection reliability: if one
	//      replica crashes mid-cycle, the others continue gathering snapshots.
	//   3. The detect engine (15D) aggregates snapshots regardless of source;
	//      receiving the same data twice per window has no logical effect.
	//
	// If billing for N×API calls becomes a concern, add leader election here by
	// wrapping the loop body in tryRunAsLeader("prober").
	if s.cfg.Prober != nil && s.cfg.Targets != nil {
		wg.Add(1)
		go func() {
			defer wg.Done()
			s.runThirdPartyProber(ctx)
		}()
	}

	// Wait for ctx cancellation (SIGTERM / test cancel).
	<-ctx.Done()
	slog.Info("scheduler: shutdown signal received — waiting for in-flight ticks to finish",
		"instance", s.cfg.InstanceID)

	// 30-second hard timeout: all loops must finish their current tick within
	// this window.  Leader keys are released by each loop's defer (see
	// runLeaderLoop → releaseLease) so followers can take over immediately.
	allStopped := make(chan struct{})
	go func() {
		wg.Wait()
		close(allStopped)
	}()

	shutdownTimeout := s.cfg.TickTimeout // reuse tick hard-timeout as shutdown budget
	select {
	case <-allStopped:
		slog.Info("scheduler: graceful shutdown complete", "instance", s.cfg.InstanceID)
	case <-time.After(shutdownTimeout):
		slog.Warn("scheduler: shutdown timeout exceeded — forcing exit",
			"timeout", shutdownTimeout, "instance", s.cfg.InstanceID)
	}
	return nil
}

// ─────────────────────────────────────────────────────────────────────────────
// Leader-elected loop driver
// ─────────────────────────────────────────────────────────────────────────────

// runLoop is the generic leader-elected loop driver.
//
//   - Non-leader: polls every RetryPeriod to acquire sched:leader:{name}.
//   - Leader: ticks at interval; a background goroutine renews the lease every
//     RenewPeriod.  Lost lease → return (exit from runLeaderLoop).
//
// The function loops continuously until ctx is Done.
func (s *Scheduler) runLoop(ctx context.Context, name string, interval time.Duration, tick func(context.Context) error) {
	key := leaderKey(name)

	for {
		// Try to acquire the leader lease (SET NX PX <leaseTTL ms>).
		ok, err := s.rdb.SetNX(ctx, key, s.cfg.InstanceID, s.cfg.LeaseTTL).Result()
		if ctx.Err() != nil {
			return // context cancelled — don't log spurious errors
		}
		if err != nil {
			slog.Error("scheduler: leader setnx error", "loop", name, "error", err)
		} else if ok {
			slog.Info("scheduler: acquired leader lease", "loop", name, "instance", s.cfg.InstanceID)
			s.runLeaderLoop(ctx, name, interval, tick, key)
			slog.Info("scheduler: relinquished leader", "loop", name, "instance", s.cfg.InstanceID)
		}

		// Not leader (or just stepped down): wait before retrying.
		select {
		case <-ctx.Done():
			return
		case <-time.After(s.cfg.RetryPeriod):
		}
	}
}

// runLeaderLoop ticks at interval while holding the leader lease.
// It starts a keepLease goroutine; when the lease is lost (renewal failure or
// key taken by a competitor), it returns so runLoop can retry.
//
// Graceful-exit guarantee: when ctx is cancelled, runLeaderLoop waits for any
// in-flight tick to finish (the tick is given its own TickTimeout context, not
// the parent ctx) before calling defer → releaseLease.  This ensures:
//   - No half-executed migrate/notify steps are left in Redis.
//   - The leader key is actively released (DEL) on clean shutdown, so the
//     follower can take over within RetryPeriod rather than LeaseTTL.
func (s *Scheduler) runLeaderLoop(
	ctx context.Context,
	name string,
	interval time.Duration,
	tick func(context.Context) error,
	leaseKey string,
) {
	// Renewal goroutine: signals via leaseLost when it can no longer renew.
	renewCtx, cancelRenew := context.WithCancel(ctx)
	leaseLost := make(chan struct{})
	go func() {
		defer close(leaseLost)
		s.keepLease(renewCtx, leaseKey, name)
	}()

	// On exit: stop the renewer and release the key so followers converge faster.
	defer func() {
		cancelRenew()
		<-leaseLost // wait for renewer before releasing
		// Use a fresh Background context: the original ctx may already be Done.
		s.releaseLease(context.Background(), leaseKey, name)
	}()

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			// Shutdown requested: exit (defer will release lease).
			// If a tick was already dispatched before this select was reached,
			// it is still running under its own TickTimeout context and will
			// complete before this goroutine is joined by Start's wg.Wait().
			return

		case <-leaseLost:
			slog.Warn("scheduler: leader lease lost — stepping down", "loop", name)
			return

		case <-ticker.C:
			// Run tick with a private context so SIGTERM / ctx cancel does not
			// abort it mid-flight.  Idempotency (15D optimistic lock, 15E
			// idempotency keys) ensures a tick that completes after the leader
			// key expires is safe.
			tickCtx, cancelTick := context.WithTimeout(context.Background(), s.cfg.TickTimeout)
			if err := tick(tickCtx); err != nil {
				slog.Error("scheduler: tick error", "loop", name, "error", err)
			}
			cancelTick()

			// If parent was cancelled while we were ticking, exit now instead
			// of waiting for the next ticker.C.
			if ctx.Err() != nil {
				return
			}
		}
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Lease renewal / release
// ─────────────────────────────────────────────────────────────────────────────

// Lua scripts are pre-evaluated to validate key ownership before each
// PEXPIRE / DEL so we never extend or delete a key owned by another instance.

const luaRenew = `
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('PEXPIRE', KEYS[1], ARGV[2])
else
    return 0
end`

const luaRelease = `
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('DEL', KEYS[1])
else
    return 0
end`

// keepLease renews the leader key every RenewPeriod using the Lua renewal
// script.  Returns when renewal fails (key gone or taken) or ctx is Done.
func (s *Scheduler) keepLease(ctx context.Context, key, name string) {
	ticker := time.NewTicker(s.cfg.RenewPeriod)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			res, err := s.rdb.Eval(
				ctx, luaRenew, []string{key},
				s.cfg.InstanceID,
				s.cfg.LeaseTTL.Milliseconds(),
			).Result()
			if err != nil {
				if ctx.Err() != nil {
					return // context cancelled: not a real failure
				}
				slog.Error("scheduler: lease renewal error", "loop", name, "error", err)
				return // treat Redis error as lease loss
			}
			if res.(int64) == 0 {
				slog.Warn("scheduler: lease renewal returned 0 — key gone or taken",
					"loop", name, "key", key)
				return
			}
		}
	}
}

// releaseLease deletes the leader key only if it still belongs to this
// instance, using the Lua release script.
func (s *Scheduler) releaseLease(ctx context.Context, key, name string) {
	res, err := s.rdb.Eval(ctx, luaRelease, []string{key}, s.cfg.InstanceID).Result()
	if err != nil {
		slog.Error("scheduler: release lease error", "loop", name, "error", err)
		return
	}
	if res.(int64) == 1 {
		slog.Info("scheduler: leader key released",
			"loop", name, "instance", s.cfg.InstanceID)
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Capacity tick (CapacityLoop work function)
// ─────────────────────────────────────────────────────────────────────────────

// capacityTick is the work function for CapacityLoop.  It runs:
//  1. 15F capacity / watermark check.
//  2. 15E grayscale weight advancement.
//  3. 15F circuit-breaker counter decay.
//  4. 15A probe-agent disconnection check → 15G alert.
//
// Each step is independent; a failure in one does not abort the others.
func (s *Scheduler) capacityTick(ctx context.Context) error {
	// 1. Capacity / watermark check (15F).
	if err := s.cfg.Capacity.Check(ctx); err != nil {
		slog.Error("scheduler: capacity check error", "error", err)
		// Non-fatal: continue.
	}

	// 2. Grayscale warm-up advancement (15E).
	if err := s.cfg.Grayscale.Advance(ctx); err != nil {
		slog.Error("scheduler: grayscale advance error", "error", err)
	}

	// 3. Circuit-breaker counter decay (15F).
	if err := s.cfg.Capacity.Decay(ctx); err != nil {
		slog.Error("scheduler: capacity decay error", "error", err)
	}

	// 4. Probe-agent disconnection check (15A probe store → 15G alert).
	if s.cfg.ProbeStore != nil && len(s.cfg.KnownProbeIDs) > 0 {
		s.checkProbeConnections(ctx)
	}

	return nil
}

// checkProbeConnections fires a Notifier alert for every known probe agent
// whose heartbeat key has expired (i.e. is absent from AliveProbes).
func (s *Scheduler) checkProbeConnections(ctx context.Context) {
	alive, err := s.cfg.ProbeStore.AliveProbes(ctx)
	if err != nil {
		slog.Error("scheduler: AliveProbes error", "error", err)
		return
	}

	aliveSet := make(map[string]struct{}, len(alive))
	for _, id := range alive {
		aliveSet[id] = struct{}{}
	}

	for _, id := range s.cfg.KnownProbeIDs {
		if _, ok := aliveSet[id]; !ok {
			if notifyErr := s.cfg.Notifier.NotifyFault(ctx, id, "probe_disconnected"); notifyErr != nil {
				slog.Error("scheduler: notify probe disconnect",
					"probe_id", id, "error", notifyErr)
			}
		}
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Third-party prober goroutine (15C, no leader constraint)
// ─────────────────────────────────────────────────────────────────────────────

// runThirdPartyProber fires Prober.RunOnce at ThirdPartyInterval on every
// replica (no leader election — see design comment in Start).
func (s *Scheduler) runThirdPartyProber(ctx context.Context) {
	ticker := time.NewTicker(s.cfg.ThirdPartyInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			targets, err := s.cfg.Targets.ListProbeTargets(ctx)
			if err != nil {
				slog.Error("scheduler: list probe targets for 15C", "error", err)
				continue
			}
			if len(targets) == 0 {
				continue
			}
			// Use a fresh context so ctx cancellation does not abort a running
			// prober call (the degradation contract: result is "no data" on
			// context error, which 15D treats as unknown, not failure).
			probeCtx, cancel := context.WithTimeout(context.Background(), s.cfg.TickTimeout)
			if err := s.cfg.Prober.RunOnce(probeCtx, targets); err != nil {
				slog.Error("scheduler: 15C prober error", "error", err)
			}
			cancel()
		}
	}
}