webui4frpc/internal/cluster/ring_balance.go

package cluster

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

	"webui4frpc/internal/store"
)

// The rebalance rule, deliberately simple and self-limiting:
//
//	release one forward when  Forwards >= minForwardsToShare  AND  Forwards > share
//
// where share = ceil(totalForwards / aliveNodes) — the even split.
//
// Every node runs this check about its OWN forwards. That is the whole point:
// a node's load only actually falls when its own topology entry goes away, so
// only the owner can produce the state change that makes the ring hand the
// forward to somebody else. A coordinator on another machine can see the
// imbalance but cannot make it true; it would have to reach in and mutate the
// owner's bookkeeping, which is exactly the sort of cross-node write this ring
// avoids everywhere else (revoke and restart are owner-directed for the same
// reason).
//
// The rule is self-limiting in both directions, so it needs no upper bound on
// how many forwards a node may shed:
//
//   - It stops at the average. 3 forwards over 3 nodes: the owner holds 2 > 1,
//     releases one, now holds 1 which is not > 1, and stops. Without the "> share"
//     bound an owner would shed EVERY forward it has.
//   - It cannot empty a node. A node holding fewer than minForwardsToShare never
//     releases, so there is always something left to serve.
//
// A release costs a live connection: stopping the worker drops whatever was
// flowing through the forward, and the new owner re-dials the local service.
// Because the rule above already requires a real imbalance, pacing is handled by
// a per-node cooldown (cooldownSec) rather than a "how long has this lasted"
// window — there is no transient to smooth over, only a churn rate to bound.
const (
	// minForwardsToShare: never release below this many owned forwards.
	minForwardsToShare = 2
	// balanceCooldownSec: minimum gap between two releases by one node. One
	// release per 10 minutes per node is a trickle; it is what stops a cluster
	// whose loads oscillate around the split point from restarting workers in a
	// loop.
	balanceCooldownSec = 600
	// releaseDeadlineSec: how long to wait for a published release to be
	// executed by its owner (us) before giving up on confirming it. A token
	// round is ~2s, so this is generous; it exists so a release lost to a
	// leader change or a joined ring cannot wedge the node forever.
	releaseDeadlineSec = 120
)

// Balance modes.
const (
	BalanceOff    = "off"
	BalanceDryRun = "dry-run"
	BalanceOn     = "on"
)

// pool is the replicated view the rule works from.
type pool struct {
	alive int // nodes that can take a forward
	mine  int // forwards THIS node owns (live count, not the replicated one)
	total int // forwards across the ring
	share int // ceil(total/alive) — the even split
}

// releaseState is one release this node has published and is waiting to see
// leave its own topology.
type releaseState struct {
	local, remote string
	port          int
	startedAt     int64
}

// Balancer lets each node shed its own share of the cluster's forwards.
//
// It publishes only ONE kind of thing: a release of a forward it owns. There is
// no targeted hand-off and no coordinator, because "drop it and let the ring's
// normal lowest-load claim pick it up" is already the ring's only mechanism for
// moving a forward — and it is the mechanism that, unlike a coordinated move,
// cannot be undone by the node that gave the forward up (the claim path runs
// after the node's load has already fallen).
type Balancer struct {
	engine *Engine
	mode   string

	// cooldownSec is the minimum gap between two releases by this node.
	cooldownSec int64
	// releasing is the forward this node gave up and is waiting to see leave.
	// While it is set, no second release is started: one at a time keeps the
	// connection churn bounded and the log readable.
	releasing *releaseState
	// lastReleaseAt is when the previous release was published (0 = never).
	lastReleaseAt int64
	// reported remembers the last dry-run candidate so the observe-only mode
	// logs a change of mind rather than the same line every tick.
	reported string
}

// NewBalancer builds a balancer. It starts in observe-only mode: releasing a
// forward drops a live connection, so acting on the ring is opt-in.
func NewBalancer(e *Engine, cooldownSec int64) *Balancer {
	b := &Balancer{engine: e, mode: BalanceDryRun, cooldownSec: cooldownSec}
	if b.cooldownSec <= 0 {
		b.cooldownSec = balanceCooldownSec
	}
	return b
}

// SetMode selects off / dry-run / on. An unrecognised value is treated as
// dry-run — a typo must not be able to turn releasing on.
func (b *Balancer) SetMode(m string) {
	switch m {
	case BalanceOff, BalanceDryRun, BalanceOn:
		b.mode = m
	default:
		if m != "" {
			log.Printf("ring[%s] unknown balance mode %q; using %s", b.engine.ID, m, BalanceDryRun)
		}
		b.mode = BalanceDryRun
	}
}

// Mode returns the active mode.
func (b *Balancer) Mode() string { return b.mode }

// SetCooldown overrides the per-node release cooldown.
func (b *Balancer) SetCooldown(sec int64) {
	if sec > 0 {
		b.cooldownSec = sec
	}
}

// InFlight reports whether a release is awaiting confirmation.
func (b *Balancer) InFlight() bool { return b.releasing != nil }

// pool computes the ring-wide view this node derives for itself.
//
// Our own count is taken LIVE from the topology rather than from the replicated
// Nodes[].Load, because Nodes[].Load is only refreshed when a token arrives and
// the decision must see the effect of the release immediately — otherwise a node
// would keep believing it is over its share and shed another forward each round.
func (b *Balancer) pool() pool {
	st := &b.engine.state
	p := pool{mine: len(st.ForwardsOwnedBy(b.engine.ID))}
	for i := range st.Nodes {
		n := &st.Nodes[i]
		if !n.Alive {
			continue
		}
		p.alive++
		if n.ID == b.engine.ID {
			p.total += p.mine
			continue
		}
		p.total += n.Load.Forwards
	}
	if p.alive > 0 {
		p.share = (p.total + p.alive - 1) / p.alive
	}
	return p
}

// assess answers "should this node give one forward up, and which one?" together
// with the facts behind the answer, so the read-only report can explain a node
// that is over its share but not acting. It never mutates anything.
//
// blocked is a short reason: "single-node", "below-min", "balanced",
// "nothing-active", "cooldown" or "in-flight".
func (b *Balancer) assess(now int64) (pick *TopoEntry, p pool, blocked string) {
	p = b.pool()
	if !b.engine.IsMember() {
		return nil, p, "single-node"
	}
	if p.mine < minForwardsToShare {
		return nil, p, "below-min"
	}
	if p.mine <= p.share {
		return nil, p, "balanced"
	}
	pick = releaseCandidate(&b.engine.state, b.engine.ID)
	if pick == nil {
		return nil, p, "nothing-active"
	}
	// The guards below do not change the verdict ("this node is over its
	// share"), only whether now is the moment; the report shows both.
	if b.releasing != nil {
		return pick, p, "in-flight"
	}
	if b.lastReleaseAt > 0 && now-b.lastReleaseAt < b.cooldownSec {
		return pick, p, "cooldown"
	}
	return pick, p, ""
}

// releaseCandidate picks which of this node's forwards to give up.
//
// Deterministic (lowest remote port, then task id) so the choice is explainable
// and repeatable: a node that re-evaluates the same state picks the same
// forward. Stopped forwards are skipped — their entry is the carrier of a user's
// stop decision, and releasing one would hand a STOPPED forward to another node
// as if it were work to run.
func releaseCandidate(s *State, owner string) *TopoEntry {
	var mine []*TopoEntry
	for _, te := range s.TopologyList() {
		if te.OwnerID != owner || !te.Active || te.Link.Disabled {
			continue
		}
		mine = append(mine, te)
	}
	if len(mine) == 0 {
		return nil
	}
	sort.Slice(mine, func(i, j int) bool {
		if mine[i].Link.RemotePort != mine[j].Link.RemotePort {
			return mine[i].Link.RemotePort < mine[j].Link.RemotePort
		}
		return mine[i].TaskID < mine[j].TaskID
	})
	return mine[0]
}

// landed reports whether the forward this node released has actually left it.
//
// The test is "no longer owned by me", NOT "no entry exists": the ring's normal
// claim picks the forward up within a round or two, so an existence test would
// see the new owner's entry and wait forever for a state that has already
// happened.
func (b *Balancer) landed() bool {
	m := b.releasing
	if m == nil {
		return true
	}
	owner := b.engine.state.TopologyOwner(&Task{
		Local:  store.Local{Name: m.local},
		Remote: store.Remote{Name: m.remote},
		Link:   store.Link{RemotePort: m.port},
	})
	return owner != b.engine.ID
}

// Step advances the balancer by one cycle and returns a human-readable
// description when it did (or, in dry-run, would have done) something.
func (b *Balancer) Step(now int64) string {
	if b.releasing != nil {
		if b.landed() {
			done := b.releasing
			b.releasing = nil
			return "gave up " + done.local + ":" + strconv.Itoa(done.port) + " for rebalancing"
		}
		if now-b.releasing.startedAt > releaseDeadlineSec {
			log.Printf("ring[%s] release of %s:%d not confirmed within %ds; clearing (forward is still ours)",
				b.engine.ID, b.releasing.local, b.releasing.port, releaseDeadlineSec)
			b.releasing = nil
		}
		return ""
	}

	pick, p, blocked := b.assess(now)
	if pick == nil {
		b.reported = ""
		return ""
	}
	desc := pick.Local.Name + ":" + strconv.Itoa(pick.Link.RemotePort) +
		" (" + strconv.Itoa(p.mine) + " owned > " + strconv.Itoa(p.share) + " share of " +
		strconv.Itoa(p.total) + " over " + strconv.Itoa(p.alive) + " nodes)"

	switch b.mode {
	case BalanceOff:
		return ""
	case BalanceDryRun:
		// Log a change of mind, not the same line every tick.
		if b.reported == desc {
			return ""
		}
		b.reported = desc
		if blocked != "" {
			return "would release " + desc + " [blocked: " + blocked + "]"
		}
		return "would release " + desc
	}

	// Mode is ON. Blocked reasons still stop us.
	if blocked != "" {
		return ""
	}
	t := b.engine.ReleaseTask(pick.Local, pick.Remote, pick.Link)
	if t == nil {
		return ""
	}
	b.releasing = &releaseState{
		local: pick.Local.Name, remote: pick.Remote.Name,
		port: pick.Link.RemotePort, startedAt: now,
	}
	b.lastReleaseAt = now
	b.reported = ""
	return "releasing " + desc
}

// Watch drives Step on a ticker until ctx is cancelled.
//
// It runs on EVERY node, not just the leader. Balancing here is not a cluster
// decision being delegated to one node — it is each node applying the same
// replicated rule to its own forwards. Two nodes never contend for the same
// forward because a node can only release what it owns.
//
// The first check happens after one interval rather than immediately: a node
// that has just booted adopts the ring's topology and its own load in the first
// token round, and deciding before that would act on a half-known pool.
func (b *Balancer) Watch(ctx context.Context, interval time.Duration) {
	tick := time.NewTicker(interval)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-tick.C:
		}
		if msg := b.Step(time.Now().Unix()); msg != "" {
			log.Printf("ring[%s] balance: %s", b.engine.ID, msg)
		}
	}
}

// ---- read-only report for the API/UI ----

// BalanceReport is what this node thinks about balancing right now.
//
// It is per-node by nature: every member has its own answer, and the answer that
// matters is the one on the node doing (or refusing to do) the releasing. The
// peer list is included because "am I over my share?" is only meaningful next to
// what everyone else holds.
type BalanceReport struct {
	Mode     string `json:"mode"`    // off | dry-run | on
	Enabled  bool   `json:"enabled"` // mode != off
	DryRun   bool   `json:"dryRun"`  // mode != on
	IsMember bool   `json:"isMember"`
	Self     string `json:"self"`

	SelfForwards  int  `json:"selfForwards"`
	Share         int  `json:"share"`
	TotalForwards int  `json:"totalForwards"`
	AliveNodes    int  `json:"aliveNodes"`
	SelfOverShare bool `json:"selfOverShare"`

	// WouldRelease is what this node would give up (present whenever it is over
	// its share, including while blocked — see BlockedBy).
	WouldRelease *BalanceReleasePlan `json:"wouldRelease,omitempty"`
	// Releasing is the release this node has published and is awaiting.
	Releasing *BalanceReleasePlan `json:"releasing,omitempty"`
	// BlockedBy explains why an over-share node is not acting yet.
	BlockedBy string `json:"blockedBy,omitempty"`

	CooldownSec   int64 `json:"cooldownSec"`
	LastReleaseAt int64 `json:"lastReleaseAt,omitempty"`

	// Peers: every member's replicated load, for context.
	Nodes []BalanceNodeReport `json:"nodes"`
}

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 BalanceReleasePlan struct {
	Local      string `json:"local"`
	Remote     string `json:"remote"`
	RemotePort int    `json:"remotePort"`
	Owner      string `json:"owner"`
}

// Report renders the current view. It never mutates anything — in particular it
// does NOT touch lastReleaseAt or the release state, so an operator polling this
// cannot start, delay or cancel a release. (A read-only endpoint that perturbs
// the behaviour it reports is worse than no endpoint.)
func (b *Balancer) Report() BalanceReport {
	now := time.Now().Unix()
	rep := BalanceReport{
		Mode:        b.mode,
		Enabled:     b.mode != BalanceOff,
		DryRun:      b.mode != BalanceOn,
		IsMember:    b.engine.IsMember(),
		Self:        b.engine.ID,
		CooldownSec: b.cooldownSec,
	}
	p := b.pool()
	rep.SelfForwards, rep.Share = p.mine, p.share
	rep.TotalForwards, rep.AliveNodes = p.total, p.alive
	rep.SelfOverShare = p.mine > p.share

	if pick, _, blocked := b.assess(now); pick != nil {
		rep.WouldRelease = &BalanceReleasePlan{
			Local: pick.Local.Name, Remote: pick.Remote.Name,
			RemotePort: pick.Link.RemotePort, Owner: pick.OwnerID,
		}
		rep.BlockedBy = blocked
	}
	if m := b.releasing; m != nil {
		rep.Releasing = &BalanceReleasePlan{
			Local: m.local, Remote: m.remote, RemotePort: m.port, Owner: b.engine.ID,
		}
		rep.BlockedBy = "in-flight"
	}
	if b.lastReleaseAt > 0 {
		rep.LastReleaseAt = b.lastReleaseAt
	}
	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(),
		})
	}
	return rep
}