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

	// 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,
		overSince: map[string]int64{},
	}
}

// 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
}