webui4frpc/internal/cluster/ring_balance.go

package cluster

import (
	"context"
	"log"
	"strconv"
	"time"

	"webui4frpc/internal/store"
)

// fwdKey identifies a forward for balancing purposes.
type fwdKey struct {
	local, remote string
	port          int
}

// planBalance decides whether one forward should move off an overloaded owner.
//
// The rule is deliberately conservative:
//   - never touch a stopped forward (disabled) — moving it would resurrect work
//     the user explicitly turned off;
//   - never touch a forward whose owner is not actually above the threshold;
//   - never move more than one per pass, so a transient load spike cannot cause a
//     burst of worker restarts (each migration drops live connections).
//
// It returns the chosen forward and the node that should take it.
// now / sustainedSec are accepted so the caller-side hysteresis (which lives in
// BalanceTracker) can pass its bookkeeping through without changing this
// signature, but this function itself is stateless about time: it answers
// "given the current snapshot, what (if anything) should move right now".
// minForwards is an absolute floor: never migrate unless the overloaded node
// holds at least this many forwards.
//
// The ratio test alone is misleading on small rings. With 3 nodes and a single
// forward, the owner scores ~105 (Forwards*100) and the two idle nodes ~5, so the
// owner sits at 20x the mean and trips any ratio — yet one forward on one node is
// the perfectly correct distribution, there is nothing to gain by moving it, and
// the move would only drop a live connection. The ratio detects *shape*; this
// floor says there must be something worth moving.
const minForwards = 2

func planBalance(s *State, minOverRatio float64) (*TopoEntry, string, bool) {
	// Candidate destinations: alive, non-leader-free (any), scored by Load.
	type cand struct {
		id       string
		score    float64
		forwards int
	}
	var dests []cand
	total := 0
	for i := range s.Nodes {
		n := &s.Nodes[i]
		if !n.Alive {
			continue
		}
		dests = append(dests, cand{n.ID, n.Load.Score(), n.Load.Forwards})
		total += n.Load.Forwards
	}
	if len(dests) < 2 {
		return nil, "", false // nothing to balance between
	}
	avg := float64(total) / float64(len(dests))
	threshold := avg * minOverRatio

	// Overloaded owners, worst first. The absolute floor is applied here so a node
	// holding one forward in a three-node ring is never considered hot.
	type over struct {
		id       string
		score    float64
		forwards int
	}
	var overs []over
	for _, d := range dests {
		if d.score > threshold && d.forwards >= minForwards {
			overs = append(overs, over{d.id, d.score, d.forwards})
		}
	}
	if len(overs) == 0 {
		return nil, "", false
	}
	// pick the worst (stable: compare score then id so ties are deterministic)
	worstIdx := 0
	for i := 1; i < len(overs); i++ {
		if overs[i].score > overs[worstIdx].score ||
			(overs[i].score == overs[worstIdx].score && overs[i].id < overs[worstIdx].id) {
			worstIdx = i
		}
	}
	worst := overs[worstIdx]

	// Lightest node overall is the destination, provided it is meaningfully
	// lighter than the worst owner.
	lightestIdx := 0
	for i := 1; i < len(dests); i++ {
		if dests[i].score < dests[lightestIdx].score ||
			(dests[i].score == dests[lightestIdx].score && dests[i].id < dests[lightestIdx].id) {
			lightestIdx = i
		}
	}
	dest := dests[lightestIdx]
	if dest.id == worst.id {
		return nil, "", false // the worst owner IS the lightest — nothing to do
	}
	// Require a real gap; migrating for a 1% difference just churns connections.
	if worst.score-dest.score < 100 {
		return nil, "", false // less than one forward's worth of difference
	}

	// Choose which forward to move: the cheapest to give up. Skip disabled ones.
	var pick *TopoEntry
	for _, te := range s.TopologyList() {
		if te.OwnerID != worst.id || !te.Active || te.Link.Disabled {
			continue
		}
		if pick == nil || te.Link.RemotePort < pick.Link.RemotePort {
			pick = te
		}
	}
	if pick == nil {
		return nil, "", false // everything it owns is stopped
	}
	return pick, dest.id, true
}

// Balancer moves forwards off an overloaded node, one at a time.
//
// It is a TWO-PHASE state machine, and the split is a safety requirement rather
// than a stylistic choice:
//
//	phase 1 (revoking): ask the current owner to stop. Wait until the topology
//	    entry is actually gone before publishing anything.
//	phase 2 (publishing): hand the forward to a NAMED destination.
//
// Publishing both at once does not work, for two independent reasons:
//
//  1. loadSnapshot() derives Forwards from ForwardsOwnedBy, so the moment the old
//     entry is dropped the old owner becomes the lightest node. A plain creation
//     task is claimed by "whoever is lightest" — so the node we just took the
//     forward FROM would immediately claim it back and the migration would be a
//     no-op that merely restarts the worker.
//  2. PendingList sorts by task id as a STRING (t10 < t5), so the revoke and the
//     submit could execute in either order within one pass.
//
// Migrate tasks solve (1) by carrying an explicit Target; this type solves (2) by
// refusing to publish phase 2 before phase 1 has completed.
type Balancer struct {
	engine *Engine

	// ratio is how far above the average a node's Score must be, relative to the
	// lightest node, before anything moves.
	ratio float64
	// sustainedSec is how long an imbalance must persist before we act. This is
	// what stops a transient CPU/network spike from restarting workers.
	sustainedSec int64

	// dryRun records the mode set at construction so Report can surface it.
	// SetDryRun updates it; main.go sets it once from W4F_BALANCE_MODE.
	dryRun bool
	// overSince records when the currently-overloaded node first crossed the
	// threshold. Zero means "not currently over".
	overSince map[string]int64
	// pending migration awaiting phase 1 (the forward to revoke).
	revoking *pendingMigration
	// migration awaiting phase 2 (the forward to publish, and where to).
	publishing *pendingMigration
}

type pendingMigration struct {
	local      string
	remote     string
	port       int
	fromOwner  string
	toOwner    string
	deadlineAt int64 // give up (re-publish for anyone) after this unix second
}

// NewBalancer builds a balancer. ratio <= 1 or sustainedSec < 0 are clamped to
// sane minimums so a bad configuration cannot turn into a migration storm.
func NewBalancer(e *Engine, ratio float64, sustainedSec int64) *Balancer {
	if ratio <= 1 {
		ratio = 2
	}
	if sustainedSec < 0 {
		sustainedSec = 0
	}
	return &Balancer{
		engine: e, ratio: ratio, sustainedSec: sustainedSec,
		dryRun:    true, // safe default: report, never move
		overSince: map[string]int64{},
	}
}

// SetDryRun selects the mode. Dry-run (the default) only reports.
func (b *Balancer) SetDryRun(v bool) { b.dryRun = v }

// DryRun reports whether the balancer is in observe-only mode.
func (b *Balancer) DryRun() bool { return b.dryRun }

// InFlight reports whether a migration is currently between phases.
func (b *Balancer) InFlight() bool {
	return b.revoking != nil || b.publishing != nil
}

// Step advances the state machine by one cycle. now is the current unix second.
// It returns a short human-readable description when something happened.
func (b *Balancer) Step(now int64) string {
	// --- phase 2: did phase 1 finish? ---
	if b.revoking != nil {
		if b.engine.state.TopologyOwner(&Task{
			Local:  store.Local{Name: b.revoking.local},
			Remote: store.Remote{Name: b.revoking.remote},
			Link:   store.Link{RemotePort: b.revoking.port},
		}) != "" {
			// Still owned: the revoke has not completed yet. Wait, unless it has
			// been far too long, in which case something is wrong with the ring.
			if now < b.revoking.deadlineAt {
				return ""
			}
			log.Printf("ring[%s] migrate revoke for %s→%s:%d timed out; abandoning",
				b.engine.ID, b.revoking.local, b.revoking.remote, b.revoking.port)
			b.revoking = nil
			return ""
		}
		// Entry gone — the old worker is stopped. Publish for the named target.
		m := b.revoking
		b.revoking = nil
		b.publishing = m
		b.publishMigration()
		return "revoke complete, migrating " + m.local + ":" + strconv.Itoa(m.port) + " → " + m.toOwner
	}

	// --- phase 1: wait for an in-flight publish, or start a new migration ---
	if b.publishing != nil {
		// The migrate task carries itself in the ring; nothing to do but let it
		// ride. Clear once the forward is back in the topology.
		if b.engine.state.TopologyOwner(&Task{
			Local:  store.Local{Name: b.publishing.local},
			Remote: store.Remote{Name: b.publishing.remote},
			Link:   store.Link{RemotePort: b.publishing.port},
		}) == b.publishing.toOwner {
			done := b.publishing
			b.publishing = nil
			return "migrated " + done.local + ":" + strconv.Itoa(done.port) + " to " + done.toOwner
		}
		if now >= b.publishing.deadlineAt {
			// Nobody claimed it. Submit an UNTARGETED task so whichever node is
			// lightest picks it up — the ring keeps the forward alive rather than
			// stranding it, which is the same guarantee the rest of the protocol
			// relies on.
			m := b.publishing
			b.publishing = nil
			b.engine.SubmitTask(store.Local{Name: m.local}, store.Remote{Name: m.remote},
				store.Link{RemotePort: m.port})
			return "migrate target did not claim " + m.local + ":" + strconv.Itoa(m.port) + "; re-opened for any node"
		}
		return ""
	}

	// --- nothing in flight: look for work ---
	pick, dest, hot := b.engine.findBalanceCandidate(b.ratio)
	if hot != "" {
		if _, seen := b.overSince[hot]; !seen {
			b.overSince[hot] = now
		}
	} else {
		b.overSince = map[string]int64{}
	}
	if pick == nil || dest == "" {
		return ""
	}
	if since, ok := b.overSince[pick.OwnerID]; !ok || now-since < int64(b.sustainedSec) {
		return ""
	}
	// Phase 1: publish the revoke and remember what we are waiting for.
	m := &pendingMigration{
		local: pick.Local.Name, remote: pick.Remote.Name, port: pick.Link.RemotePort,
		fromOwner: pick.OwnerID, toOwner: dest,
		// Generous deadline: a token round is ~2s, and the owner may be busy.
		deadlineAt: now + 120,
	}
	b.revoking = m
	b.engine.RevokeTask(pick.Local, pick.Remote, pick.Link)
	// The revoke task must be able to reach the owner even though the balancer
	// runs on the leader: it travels in the ring like any other command.
	return "revoking " + m.local + ":" + strconv.Itoa(m.port) + " from " + m.fromOwner
}

func (b *Balancer) publishMigration() {
	m := b.publishing
	if m == nil {
		return
	}
	b.engine.submitMigrate(store.Local{Name: m.local}, store.Remote{Name: m.remote},
		store.Link{RemotePort: m.port}, m.toOwner)
}

// findBalanceCandidate exposes planBalance to the engine's live state.
func (e *Engine) findBalanceCandidate(ratio float64) (*TopoEntry, string, string) {
	pick, dest, ok := planBalance(&e.state, ratio)
	if !ok {
		return nil, "", ""
	}
	// Report the overloaded node even when there is nothing movable, so the
	// hysteresis timer can start for a node whose forwards are all stopped.
	return pick, dest, pick.OwnerID
}

// WatchBalance drives the balancer on a ticker until ctx is cancelled.
//
// Only the LEADER runs this. Balancing is a cluster-wide decision (it moves
// forwards between members), and letting every member decide independently would
// mean several migrations racing for the same forward — the ring is a shared
// token, so two nodes must not both publish a revoke for the same entry.
//
// dryRun makes the balancer report what it WOULD do without publishing anything.
// That is the mode to use first on a live cluster: it answers "is my real data
// going to trip the threshold, and what would it move?" before anything on the
// ring is disturbed.
func (b *Balancer) WatchBalance(ctx context.Context, interval time.Duration, dryRun bool) {
	tick := time.NewTicker(interval)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-tick.C:
		}
		if !b.engine.IsLeader() {
			continue
		}
		now := time.Now().Unix()
		if dryRun {
			if msg := b.stepDryRun(now); msg != "" {
				log.Printf("ring[%s] balance (dry-run): %s", b.engine.ID, msg)
			}
			continue
		}
		if msg := b.Step(now); msg != "" {
			log.Printf("ring[%s] balance: %s", b.engine.ID, msg)
		}
	}
}

// stepDryRun reports the migration the balancer would start right now, without
// publishing a revoke. Hysteresis is still applied so the report reflects the
// same sustained-imbalance rule the real path uses.
func (b *Balancer) stepDryRun(now int64) string {
	if b.InFlight() {
		return ""
	}
	pick, dest, hot := b.engine.findBalanceCandidate(b.ratio)
	if hot == "" {
		b.overSince = map[string]int64{}
		return ""
	}
	if _, seen := b.overSince[hot]; !seen {
		b.overSince[hot] = now
		return ""
	}
	if pick == nil || dest == "" {
		return ""
	}
	if since, ok := b.overSince[pick.OwnerID]; !ok || now-since < int64(b.sustainedSec) {
		return ""
	}
	return "would migrate " + pick.Local.Name + ":" + strconv.Itoa(pick.Link.RemotePort) +
		" from " + pick.OwnerID + " to " + dest
}

// BalanceReport is a read-only view of what the balancer sees and would do.
//
// It exists because the balancer's decision is otherwise only visible in the
// log, and only when an imbalance exists AND the hysteresis window has elapsed.
// An operator asking "is this cluster balanced?" gets silence in exactly the
// healthy case, which is indistinguishable from "balancing is switched off".
type BalanceReport struct {
	Enabled    bool    `json:"enabled"`
	DryRun     bool    `json:"dryRun"`
	Ratio      float64 `json:"ratio"`
	SustainSec int64   `json:"sustainSec"`
	// Only the leader acts; on a follower these stay false/empty.
	IsLeader bool `json:"isLeader"`
	// Current per-node scores.
	Nodes []BalanceNodeReport `json:"nodes"`
	// Overloaded (above ratio x mean) and sustained past the window, if any.
	OverloadedSince int64 `json:"overloadedSince,omitempty"`
	// The migration that would happen on the next Step, if any.
	WouldMigrate *BalanceMigrationPlan `json:"wouldMigrate,omitempty"`
	// Set when a migration is between phases.
	InFlight *BalanceMigrationPlan `json:"inFlight,omitempty"`
}

type BalanceNodeReport struct {
	ID       string  `json:"id"`
	Alive    bool    `json:"alive"`
	IsLeader bool    `json:"isLeader"`
	Forwards int     `json:"forwards"`
	MemPct   float64 `json:"memPct"`
	NetPct   float64 `json:"netPct"`
	Score    float64 `json:"score"`
}

type BalanceMigrationPlan struct {
	Local  string `json:"local"`
	Remote string `json:"remote"`
	Port   int    `json:"remotePort"`
	From   string `json:"from"`
	To     string `json:"to"`
}

// Report renders the balancer's current view. It never mutates anything —
// note that it does NOT touch overSince, so calling it cannot start or reset
// the hysteresis timer that the real Step path uses.
func (b *Balancer) Report(now int64) BalanceReport {
	rep := BalanceReport{
		Enabled:    !b.dryRun,
		DryRun:     b.dryRun,
		Ratio:      b.ratio,
		SustainSec: int64(b.sustainedSec),
		IsLeader:   b.engine.IsLeader(),
	}
	for i := range b.engine.state.Nodes {
		n := &b.engine.state.Nodes[i]
		rep.Nodes = append(rep.Nodes, BalanceNodeReport{
			ID: n.ID, Alive: n.Alive, IsLeader: n.IsLeader,
			Forwards: n.Load.Forwards, MemPct: n.Load.MemPct, NetPct: n.Load.NetPct,
			Score: n.Load.Score(),
		})
	}
	if m := b.revoking; m != nil {
		rep.InFlight = &BalanceMigrationPlan{
			Local: m.local, Remote: m.remote, Port: m.port, From: m.fromOwner, To: m.toOwner,
		}
	}
	if m := b.publishing; m != nil {
		rep.InFlight = &BalanceMigrationPlan{
			Local: m.local, Remote: m.remote, Port: m.port, From: m.fromOwner, To: m.toOwner,
		}
	}
	// Only a leader would act, so only report a plan there.
	if rep.IsLeader {
		pick, dest, hot := b.engine.findBalanceCandidate(b.ratio)
		if hot != "" {
			rep.OverloadedSince = b.overSince[hot]
		}
		if pick != nil && dest != "" {
			plan := &BalanceMigrationPlan{
				Local: pick.Local.Name, Remote: pick.Remote.Name,
				Port: pick.Link.RemotePort, From: pick.OwnerID, To: dest,
			}
			// surface it as "would migrate" only once the window is satisfied
			if since, ok := b.overSince[pick.OwnerID]; ok && now-since >= int64(b.sustainedSec) {
				rep.WouldMigrate = plan
			}
		}
	}
	return rep
}

// SetRatioAndSustain updates the thresholds after construction, so main.go can
// create the balancer early (the HTTP handler needs it) and apply the environment
// configuration once it has parsed it.
func (b *Balancer) SetRatioAndSustain(ratio float64, sustainedSec int64) {
	if ratio <= 1 {
		ratio = 2
	}
	if sustainedSec < 0 {
		sustainedSec = 0
	}
	b.ratio = ratio
	b.sustainedSec = sustainedSec
}