webui4frpc/internal/cluster/ring_balance.go

package cluster

// 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".
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
	}
	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()})
		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.
	type over struct {
		id    string
		score float64
	}
	var overs []over
	for _, d := range dests {
		if d.score > threshold {
			overs = append(overs, over{d.id, d.score})
		}
	}
	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
}