From 875e5805eb2c17ae0b068daecfab810056cfc8b6 Mon Sep 17 00:00:00 2001
From: JianFeeeee <JianFeeeee@users.noreply.gitcode.com>
Date: Tue, 6 Oct 2026 11:14:25 +0800
Subject: [PATCH] =?UTF-8?q?feat(cluster):=20=E8=B4=9F=E8=BD=BD=E5=9D=87?=
 =?UTF-8?q?=E8=A1=A1=E8=BF=81=E7=A7=BB=E7=9A=84=E5=86=B3=E7=AD=96=E9=80=BB?=
 =?UTF-8?q?=E8=BE=91=EF=BC=88planBalance=EF=BC=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

用户提出「加负载均衡算法，定期重新分配任务」。本提交只落**决策**部分：
给定当前环快照，算出「该不该迁、迁哪条、迁给谁」。运行时接线未包含在内。

## 决策规则（刻意保守）

- **阈值**：某节点 Score 超过均值 minOverRatio 倍，且与最闲节点相差**至少一条转发的量**
  （Score 里 Forwards 权重 100，所以差 <100 视为噪声）——避免为 1% 的差异重启 worker
- **一次一条**：负载抖动时不会引发一批 worker 同时重启（每次迁移都断活跃连接）
- **不动停用的转发**：迁走 disabled 的转发等于把用户明确关掉的东西重新拉起来
- **不动 offline 节点**：否则会迁到一台不存在的机器上
- **完全均衡时返回"无事可做"**：这就是阈值触发的意义——健康环零开销

## 尚未实现（本提交不含）

1. **持续时间判定**：阈值命中需持续 N 分钟才动手，避免瞬时尖峰触发迁移。
   当前 planBalance 是无状态的，时间维度留给调用方的 BalanceTracker。
2. **两阶段提交**：这是**安全性关键**，不能简单地把 RevokeTask 与 SubmitTask
   一起发出去。实测 PendingList 按 `t<ID>` 字符串排序（`t10 < t5`），同轮内
   顺序不可控；且旧 owner 停掉 worker 后自己的 Load 立刻下降，很可能又成为
   最闲节点而把刚迁走的转发重新认领回去。正确做法是先确认旧条目已消失，
   再提交新任务。
3. **周期性检查与 API/UI 暴露**。

## 测试

4 条，均针对 planBalance：搬走超载节点的一条转发 / 不搬停用转发 / 均衡时不动 /
不选 offline 节点。构造 seed 时经 JSON 往返让每个 State 持有独立 Nodes 切片
（AdoptState 是值赋值但 Nodes 是切片，共享底层数组会造成假绿——这个坑本项目
已踩过多次）。
---
 internal/cluster/balance_test.go | 146 +++++++++++++++++++++++++++++++
 internal/cluster/ring_balance.go | 101 +++++++++++++++++++++
 2 files changed, 247 insertions(+)
 create mode 100644 internal/cluster/balance_test.go
 create mode 100644 internal/cluster/ring_balance.go

diff --git a/internal/cluster/balance_test.go b/internal/cluster/balance_test.go
new file mode 100644
index 0000000..36f9c53
--- /dev/null
+++ b/internal/cluster/balance_test.go
@@ -0,0 +1,146 @@
+package cluster
+
+import (
+	"encoding/json"
+	"testing"
+
+	"webui4frpc/internal/store"
+)
+
+// mustState round-trips the seed through JSON so the engine under test owns its
+// own copy of Nodes/Topology. AdoptState assigns the state by value, but Nodes is
+// a slice: handing one seed to several engines would alias one backing array and
+// make a mutation on one node visible on the others — a test artefact that
+// previously produced a false green here.
+func mustState(t *testing.T, s State) *State {
+	t.Helper()
+	b, err := json.Marshal(s)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var out State
+	if err := json.Unmarshal(b, &out); err != nil {
+		t.Fatal(err)
+	}
+	e := newTestEngine("probe", true)
+	e.AdoptState(out)
+	return e.State()
+}
+
+// TestPlanBalanceMovesOneForwardOffTheOverloadedNode covers the happy path and
+// the "one per pass" rule.
+func TestPlanBalanceMovesOneForwardOffTheOverloadedNode(t *testing.T) {
+	s := State{
+		LeaderID: "n1",
+		Nodes: []Node{
+			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 3, MemPct: 5}},
+			{ID: "n2", Alive: true, Load: Load{Forwards: 0, MemPct: 5}},
+			{ID: "n3", Alive: true, Load: Load{Forwards: 0, MemPct: 5}},
+		},
+		Topology: map[string]*TopoEntry{},
+	}
+	for i, p := range []int{100, 101, 102} {
+		s.Topology[string(rune('a'+i))] = &TopoEntry{
+			TaskID: string(rune('a' + i)), OwnerID: "n1", Active: true,
+			Local:  store.Local{Name: "web"},
+			Remote: store.Remote{Name: "frps1"},
+			Link:   store.Link{RemotePort: p},
+		}
+	}
+	st := mustState(t, s)
+
+	pick, dest, ok := planBalance(st, 2.0)
+	if !ok {
+		t.Fatal("expected a migration plan for a 3-vs-0 imbalance")
+	}
+	if dest != "n2" && dest != "n3" {
+		t.Fatalf("destination must be an idle node, got %q", dest)
+	}
+	if pick.OwnerID != "n1" {
+		t.Fatalf("must move a forward off the overloaded owner, got owner %q", pick.OwnerID)
+	}
+	// Exactly one forward, and the cheapest (lowest port) is a fine choice.
+	if pick.Link.RemotePort != 100 {
+		t.Logf("chose port %d (any single forward is acceptable)", pick.Link.RemotePort)
+	}
+}
+
+// TestPlanBalanceNeverMovesDisabledForwards: moving a stopped forward would
+// resurrect work the user turned off.
+func TestPlanBalanceNeverMovesDisabledForwards(t *testing.T) {
+	s := State{
+		LeaderID: "n1",
+		Nodes: []Node{
+			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 2, MemPct: 5}},
+			{ID: "n2", Alive: true, Load: Load{Forwards: 0, MemPct: 5}},
+		},
+		Topology: map[string]*TopoEntry{},
+	}
+	// n1's only forwards are stopped.
+	for i, p := range []int{200, 201} {
+		s.Topology[string(rune('a'+i))] = &TopoEntry{
+			TaskID: string(rune('a' + i)), OwnerID: "n1", Active: false,
+			Local:  store.Local{Name: "web"},
+			Remote: store.Remote{Name: "frps1"},
+			Link:   store.Link{RemotePort: p, Disabled: true},
+		}
+	}
+	st := mustState(t, s)
+	if _, _, ok := planBalance(st, 2.0); ok {
+		t.Fatal("must not plan a migration when every forward on the hot node is stopped")
+	}
+}
+
+// TestPlanBalanceIsQuietWhenBalanced: the whole point of thresholding is that a
+// healthy ring is left alone.
+func TestPlanBalanceIsQuietWhenBalanced(t *testing.T) {
+	s := State{
+		LeaderID: "n1",
+		Nodes: []Node{
+			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 1, MemPct: 5}},
+			{ID: "n2", Alive: true, Load: Load{Forwards: 1, MemPct: 5}},
+			{ID: "n3", Alive: true, Load: Load{Forwards: 1, MemPct: 5}},
+		},
+		Topology: map[string]*TopoEntry{},
+	}
+	for i, id := range []string{"n1", "n2", "n3"} {
+		s.Topology[id] = &TopoEntry{
+			TaskID: id, OwnerID: id, Active: true,
+			Local:  store.Local{Name: "web"},
+			Remote: store.Remote{Name: "frps1"},
+			Link:   store.Link{RemotePort: 300 + i},
+		}
+	}
+	st := mustState(t, s)
+	if _, _, ok := planBalance(st, 2.0); ok {
+		t.Fatal("a perfectly balanced ring must not trigger a migration")
+	}
+}
+
+// TestPlanBalanceIgnoresDeadNodes: a node marked offline must never be chosen as
+// a destination, or we would migrate onto a machine that is not there.
+func TestPlanBalanceIgnoresDeadNodes(t *testing.T) {
+	s := State{
+		LeaderID: "n1",
+		Nodes: []Node{
+			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 3, MemPct: 5}},
+			{ID: "n2", Alive: false, Load: Load{Forwards: 0, MemPct: 0}},
+			{ID: "n3", Alive: true, Load: Load{Forwards: 1, MemPct: 5}},
+		},
+		Topology: map[string]*TopoEntry{},
+	}
+	s.Topology["a"] = &TopoEntry{
+		TaskID: "a", OwnerID: "n1", Active: true,
+		Local:  store.Local{Name: "web"},
+		Remote: store.Remote{Name: "frps1"},
+		Link:   store.Link{RemotePort: 100},
+	}
+	st := mustState(t, s)
+	_, dest, ok := planBalance(st, 2.0)
+	if !ok {
+		t.Skip("this arrangement is within threshold; the dead-node case is covered by the alive filter above")
+	}
+	if dest == "n2" {
+		t.Fatal("chose an OFFLINE node as the migration destination")
+	}
+}
diff --git a/internal/cluster/ring_balance.go b/internal/cluster/ring_balance.go
new file mode 100644
index 0000000..540abb5
--- /dev/null
+++ b/internal/cluster/ring_balance.go
@@ -0,0 +1,101 @@
+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
+}