feat(cluster): 负载均衡改为每节点自主「先撤销本地转发,再丢回环」

按用户要求重做:不再由 leader 单独调度,而是**每个节点对自己的转发跑同一条
规则**,超了就自己撤销一条、把任务丢回环上,由环原有的「最闲者认领」决定它
落到哪台。

## 为什么必须是 owner 自己动手(不是洁癖,是实测逼出来的)

真机探针(跑完已删)显示,真实 revoke 之后:

    TopologyOwner after revoke = "n1"     ← 条目还在(只标记不删除)
    n1 Forwards=3  score=310.0            ← 负载根本没降
    after step: revoking=true publishing=false

⇒ 一节点的负载**只有它自己的拓扑条目消失时才会下降**,而条目是本地簿记,没有任何
对端可以代它写。leader 从外面只能看见失衡、无法让失衡成为事实——它必须伸手去改
另一个节点的簿记,而这正是这个环到处都在避免的跨节点写(revoke / restart 定向往
owner 也是同一个理由)。

⇒ 而且旧实现**根本跑不起来**:phase 1 判据是「条目消失」,可 mark-don't-remove
让它永远不消失,于是每 120s 打一条 `revoking … timed out; abandoning`,
**phase 2 永不执行,一条转发都不会迁**。开 `W4F_BALANCE_MODE=on` 只会刷日志。

## 规则(替代原来的 ratio 阈值)

    当 Forwards >= 2 且 Forwards > ceil(总数/存活节点数) 时,撤销一条

- **自限**:3 条 / 3 台 → owner 2 > 1,卸一条 → 1,`1 > 1` 不成 → 停。
  没有上界的话 owner 会把转发**全部卸光**。
- **不会卸空**:低于下限(2)永不撤销。
- 平摊值向上取整(3/3=1,不是 0),否则 1 条 / 3 台的节点也算「超了」。
- **不再用 mem/net 参与判据**:实测 `.30` 空闲时 netPct 能到 55%,Score 会高于
  有转发的节点——用 Score 会把纯网络抖动当成负载。
- 节奏用**每节点冷却**(默认 600s)而不是「失衡持续多久」窗口:规则本身已经要求
  真实失衡,没有需要平滑的瞬态,只有需要限制的抖动频率。

## 为什么必须是新的 Release 任务类型,不能复用 Revoke

用户停止(Revoke)**保留拓扑条目**作为 disabled 标志的载体;均衡释放必须**删除
条目**,因为 (a) 负载由自己的条目数算出,条目不删负载不降;(b) 每个节点的
duplicate-guard 会拒绝「已有 owner 的转发」的创建任务——保留条目会让**所有节点
都拒领**,转发被永久搁死在刚放弃它的那台机器上。

Release 因此:`Handler.Release`(停 worker,但**不标记 stopped**)→ 删条目 →
`SubmitTask` 以**启用状态**重新发布。它需要自己的 Handler 方法,不能回落到
Revoke:Revoke 会顺手把 link 标成 stopped,而这条转发即将在别的节点上跑。

## 顺带修掉两个既有缺口(实测确认)

1. **`submitMigrate` 发布不出去**:`OnToken` 的 localPending 重合并只保留「从未
   publish 过」的本地 pending,而只有 `Add*` 系列会记 `published`。引擎里直接写
   `PendingTasks` 的任务会在下一次 adopt 时被静默丢弃。
   ⇒ 新 `ReleaseTask` 走 `AddPending`。
   (`SubmitRestart` 有同样的毛病,但它不在本次范围,未改动——见提交说明末尾。)
2. `Release` 需要绕过 duplicate-claim 守卫(同 Restart),否则「owner 本来就持有
   这条转发」这一前提会让每条 Release 都被丢掉。

## 测试(17 条,全绿)

规则:平摊值取整、超了才撤、撤到平摊就停、单节点不撤、低于下限不撤、不撤已停
转发、候选选择确定。
机制:Revoke 留载体 vs Release 删条目(这是两种语义的核心差异)、Release 只由
owner 执行且不消费他人的任务、owner 释放后环的普通认领把它接走。
模式:dry-run 不动任何东西、off 真的什么都不做、未知模式回落 dry-run、
冷却限制抖动频率、确认超时清理。
只读:Report 不变更任何状态、能解释「超了但不动」的原因、均衡节点给出真实答案。

**双向验证 3 条**(改坏→变红→还原→变绿):
- 去掉 Release 的 duplicate-guard 例外 → `the owner released 0 times, want 1`
- 去掉「不超过平摊值」上界 → `2 owned against share 2 must not release`
- Release 改成只标记不删除 → `n1 should now own 2, got 3`

go build / go vet / go test ./... 全绿(go clean -testcache 后逐包实跑)。

## 未做 / 待办

- `SubmitRestart` 的同类发布缺口未修(同样会被 adopt 丢弃),需要单独确认它
  在实机上是否真的丢过 restart。
- UI 尚未跟进新报告结构(下一步)。
This commit is contained in:
JianFeeeee
2026-10-06 14:03:03 +08:00
parent b40b53b40b
commit c69fb2caba
8 changed files with 989 additions and 860 deletions

View File

@ -290,6 +290,21 @@ func main() {
log.Printf("ring[%s] revoked task %s: %s→%s:%d", selfID, tk.ID, tk.Local.Name, tk.Remote.Name, tk.Link.RemotePort) log.Printf("ring[%s] revoked task %s: %s→%s:%d", selfID, tk.ID, tk.Local.Name, tk.Remote.Name, tk.Link.RemotePort)
return nil return nil
}, },
// ReleaseFn stops this forward's worker for a REBALANCE, leaving the
// link ENABLED. This must not reuse RevokeFn: that marks the link
// stopped, and the released forward is about to run on another node —
// a "stopped" flag would make the claim path treat the re-published
// task as a disabled forward and refuse to start a worker for it,
// leaving the forward dead everywhere.
ReleaseFn: func(ctx context.Context, tk *cluster.Task) error {
key := process.WorkerKey(tk.Local.Name, tk.Remote.Name, tk.Link.RemotePort)
if _, running := pm.Status(key); running {
_ = pm.Stop(key)
}
log.Printf("ring[%s] released worker %s (rebalancing): %s→%s:%d", selfID, key,
tk.Local.Name, tk.Remote.Name, tk.Link.RemotePort)
return nil
},
// RestartFn brings an already-claimed forward back after a stop. It // RestartFn brings an already-claimed forward back after a stop. It
// must NOT run the claim bookkeeping (ClaimFn), which appends links and // must NOT run the claim bookkeeping (ClaimFn), which appends links and
// writes topology state: the forward is already established, and a stop // writes topology state: the forward is already established, and a stop
@ -392,7 +407,7 @@ func main() {
// Load balancer is created here (before the Handler) so the HTTP layer can // Load balancer is created here (before the Handler) so the HTTP layer can
// expose its read-only report. Mode/parameters are applied further down, once // expose its read-only report. Mode/parameters are applied further down, once
// the W4F_BALANCE_* environment has been parsed. // the W4F_BALANCE_* environment has been parsed.
balancer := cluster.NewBalancer(ring, 2, 1800) balancer := cluster.NewBalancer(ring, 600)
h := &httpapi.Handler{ h := &httpapi.Handler{
Store: st, Store: st,
@ -562,41 +577,40 @@ func main() {
go ring.WatchTokenLoss(clusterCtx) go ring.WatchTokenLoss(clusterCtx)
go ring.WatchLeader(clusterCtx) go ring.WatchLeader(clusterCtx)
// Load balancing. DRY-RUN BY DEFAULT. // Load balancing — SELF-DIRECTED, and dry-run by default.
// //
// The balancer moves forwards off an overloaded node, which means stopping a // Every node runs this check against its OWN forwards. There is no
// frpc worker and dropping its live connections. That is fine occasionally // coordinator: a node's load only actually falls when its own topology entry
// (a genuinely hot node) but destructive as a reflex, so it ships off and is // goes away, so only the owner can produce the state change that makes the
// opted into explicitly: // ring hand a forward to somebody else. The node gives a forward up and
// re-publishes it as an ordinary unclaimed task; the ring's normal
// lowest-load claim then places it. (The earlier design had the leader pick a
// victim and a destination — but it can only see the imbalance, not resolve
// it, so it had to reach into another node's bookkeeping to make its decision
// true.)
// //
// W4F_BALANCE_MODE=dry-run report what it would do, move nothing (default) // A release stops a worker and drops a live connection, so acting on the ring
// W4F_BALANCE_MODE=on actually migrate // is opt-in:
// W4F_BALANCE_RATIO=2 overload = Score above RATIO x the mean //
// W4F_BALANCE_SUSTAIN=1800 seconds an imbalance must persist // W4F_BALANCE_MODE=dry-run report the decision, release nothing (default)
// W4F_BALANCE_MODE=on actually release
// W4F_BALANCE_MODE=off do nothing at all
// W4F_BALANCE_COOLDOWN=600 min seconds between two releases by one node
// W4F_BALANCE_INTERVAL=10 check every N seconds // W4F_BALANCE_INTERVAL=10 check every N seconds
//
// Only the leader balances (WatchBalance no-ops elsewhere): moving forwards is
// a cluster-wide decision and two members doing it independently would race on
// the same token.
balMode := strings.ToLower(strings.TrimSpace(os.Getenv("W4F_BALANCE_MODE"))) balMode := strings.ToLower(strings.TrimSpace(os.Getenv("W4F_BALANCE_MODE")))
if balMode == "" { if balMode == "" {
balMode = "dry-run" balMode = cluster.BalanceDryRun
} }
balRatio := envFloat("W4F_BALANCE_RATIO", 2)
balSustain := int64(envFloat("W4F_BALANCE_SUSTAIN", 1800))
balInterval := time.Duration(envFloat("W4F_BALANCE_INTERVAL", 10)) * time.Second balInterval := time.Duration(envFloat("W4F_BALANCE_INTERVAL", 10)) * time.Second
if balInterval < time.Second { if balInterval < time.Second {
balInterval = 10 * time.Second balInterval = 10 * time.Second
} }
balancer.SetRatioAndSustain(balRatio, balSustain) balCooldown := int64(envFloat("W4F_BALANCE_COOLDOWN", 600))
if balMode == "on" { balancer.SetCooldown(balCooldown)
log.Printf("ring[%s] load balancing ENABLED (ratio=%.2f sustain=%ds interval=%s)", balancer.SetMode(balMode)
selfID, balRatio, balSustain, balInterval) log.Printf("ring[%s] load balancing %s (cooldown=%ds interval=%s, self-directed)",
} else { selfID, strings.ToUpper(balMode), balCooldown, balInterval)
log.Printf("ring[%s] load balancing in DRY-RUN (set W4F_BALANCE_MODE=on to enable; "+ go balancer.Watch(clusterCtx, balInterval)
"ratio=%.2f sustain=%ds interval=%s)", selfID, balRatio, balSustain, balInterval)
}
go balancer.WatchBalance(clusterCtx, balInterval, balMode != "on")
// Startup reconcile: every node (leader or not) checks its own topology // Startup reconcile: every node (leader or not) checks its own topology
// entries and re-claims any whose per-forward frpc worker process is // entries and re-claims any whose per-forward frpc worker process is

View File

@ -3,6 +3,7 @@ package cluster
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"strconv"
"strings" "strings"
"testing" "testing"
"time" "time"
@ -30,124 +31,6 @@ func mustState(t *testing.T, s State) *State {
return e.State() 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")
}
}
// migrateRing builds a 3-node ring where n1 owns every forward, plus n2/n3 idle. // migrateRing builds a 3-node ring where n1 owns every forward, plus n2/n3 idle.
func migrateRing(t *testing.T, ownCount int) (*Engine, *Engine, *Engine) { func migrateRing(t *testing.T, ownCount int) (*Engine, *Engine, *Engine) {
t.Helper() t.Helper()
@ -194,330 +77,545 @@ func migrateRing(t *testing.T, ownCount int) (*Engine, *Engine, *Engine) {
return n1, n2, n3 return n1, n2, n3
} }
// TestMigrateTaskIsAddressedNotClaimedByLightest is the core safety test. // ---- helpers ----
//
// After the revoke lands, the node we took the forward from becomes the
// lightest — an untargeted creation task would be claimed straight back, making
// the migration a no-op that merely restarts the worker. Migrate tasks must be
// delivered to their named target only.
//
// The forward is removed from the topology first (as the balancer's phase 1
// does), so this exercises the realistic post-revoke state.
func TestMigrateTaskIsAddressedNotClaimedByLightest(t *testing.T) {
n1, n2, n3 := migrateRing(t, 3)
n1.State().RemoveTopology("web", "frps1", 100) // phase 1 has landed
n1.submitMigrate(store.Local{Name: "web"}, store.Remote{Name: "frps1"}, // balanceRing builds a 3-node ring with `owned` forwards on n1 (the evaluation
store.Link{RemotePort: 100}, "n2") // node) and `others` on each of the other two, so a test can set up any
// distribution directly. The replicated Nodes[].Load is made consistent with the
// Give each node a DISTINCT load via its handler. loadSnapshot() refreshes // topology, because that is what a real token round leaves behind.
// mem/net from Handler.RuntimeLoad() on every OnToken, so seeding func balanceRing(t *testing.T, owned, others int) (*Engine, *Engine, *Engine) {
// State.Nodes[].Load by hand is silently overwritten — which is why an t.Helper()
// earlier version of this test had n2 and n3 tied and therefore passed even mk := func(id string, leader bool) *Engine {
// with the targeted delivery removed (the "lightest" node happened to be the h := &fakeHandler{load: Load{MemPct: 5, NetPct: 5}}
// target). n3 is now strictly the lightest, so an untargeted rule would hand return NewEngine(id, id+":7500", "u", "p", "0.1.3", nil, h,
// the task to n3 and fail here. func(ctx context.Context, next string, tk *Token) error { return nil },
claimedBy := map[string]int{} id+":7500", leader, "")
mk := func(id string, mem float64) Handler {
return &fakeHandler{load: Load{MemPct: mem, NetPct: 0},
claim: func(ctx context.Context, tk *Task) error {
claimedBy[id]++
return nil
}}
} }
n1.Handler = mk("n1", 50) n1, n2, n3 := mk("n1", true), mk("n2", false), mk("n3", false)
n2.Handler = mk("n2", 40) // the TARGET — heavier than n3
n3.Handler = mk("n3", 1) // strictly the lightest seed := State{
LeaderID: "n1",
Nodes: []Node{
{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true, Load: Load{Forwards: owned, MemPct: 5}},
{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{Forwards: others, MemPct: 5}},
{ID: "n3", Addr: "n3:7500", Alive: true, Load: Load{Forwards: others, MemPct: 5}},
},
Topology: map[string]*TopoEntry{},
PendingTasks: map[string]*Task{},
}
for _, e := range []*Engine{n1, n2, n3} {
var own State
b, _ := json.Marshal(seed)
if err := json.Unmarshal(b, &own); err != nil {
t.Fatal(err)
}
e.AdoptState(own)
}
// Build the whole topology in n1's state, then broadcast it. Adding each
// node's forwards to its OWN state and then re-adopting from n1 would discard
// the others (one state is the source of the broadcast), which is how an
// earlier version of this helper silently produced a 1-1-1 split regardless
// of the arguments.
add := func(owner, prefix string, count, basePort int) {
for i := 0; i < count; i++ {
n1.state.AddTopology(&Task{
ID: prefix + strconv.Itoa(i),
Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
Link: store.Link{RemotePort: basePort + i},
}, owner)
}
}
add("n1", "f", owned, 100)
add("n2", "g", others, 200)
add("n3", "h", others, 300)
var snap State var snap State
b, _ := json.Marshal(n1.State()) b, _ := json.Marshal(n1.state)
_ = json.Unmarshal(b, &snap) _ = json.Unmarshal(b, &snap)
for _, e := range []*Engine{n1, n2, n3} { for _, e := range []*Engine{n1, n2, n3} {
e.AdoptState(snap) e.AdoptState(snap)
} }
for i := 0; i < 3; i++ { return n1, n2, n3
for _, e := range []*Engine{n1, n2, n3} { }
if _, err := e.OnToken(context.Background(), &Token{
Cycle: int64(i + 1), State: *e.State(), SentAt: int64(i + 1), // ---- the rule ----
}); err != nil {
t.Fatal(err) // TestShareIsTheEvenSplit pins the arithmetic, including the ceiling. An even
} // split of 3 over 3 is 1 (not 1.0 floored to 0): with a floor, a node holding
// one of three forwards would still be "over its share" and shed it.
func TestShareIsTheEvenSplit(t *testing.T) {
for _, tc := range []struct{ total, alive, want int }{
{3, 3, 1}, {4, 3, 2}, {6, 2, 3}, {1, 3, 1}, {0, 3, 0}, {7, 3, 3},
} {
got := (tc.total + tc.alive - 1) / tc.alive
if got != tc.want {
t.Errorf("share(%d over %d) = %d, want %d", tc.total, tc.alive, got, tc.want)
} }
} }
if claimedBy["n1"] != 0 {
t.Errorf("n1 claimed %d migrate task(s); the old owner must never take it back",
claimedBy["n1"])
}
if claimedBy["n3"] != 0 {
t.Errorf("n3 (the LIGHTEST node) claimed %d migrate task(s); the task is addressed "+
"to n2 and must ignore load entirely", claimedBy["n3"])
}
if claimedBy["n2"] != 1 {
t.Errorf("the target n2 claimed %d time(s), want exactly 1", claimedBy["n2"])
}
} }
// TestMigrateRefusedWhileEntryStillOwned pins the duplicate guard: a migrate that // TestOwnerReleasesWhenOverItsShare: 3 forwards on n1, none elsewhere → n1 holds
// arrives before the revoke completed must NOT create a second topology entry // 3 against a share of 1 and must give one up.
// for the same forward. This is the failure that would otherwise leave two nodes func TestOwnerReleasesWhenOverItsShare(t *testing.T) {
// running workers for one remotePort — which frps rejects with n1, _, _ := balanceRing(t, 3, 0)
// "proxy already exists" (observed 14k times before that was fixed). b := NewBalancer(n1, 0)
func TestMigrateRefusedWhileEntryStillOwned(t *testing.T) {
n1, _, _ := migrateRing(t, 3)
claimed := 0
n1.Handler = &fakeHandler{load: Load{MemPct: 5, NetPct: 5},
claim: func(ctx context.Context, tk *Task) error { claimed++; return nil }}
// Entry still owned by n1 — the revoke has NOT landed. pick, p, blocked := b.assess(1000)
n1.submitMigrate(store.Local{Name: "web"}, store.Remote{Name: "frps1"}, if pick == nil {
store.Link{RemotePort: 100}, "n1") t.Fatalf("n1 holds 3 of 3 forwards and should release one (blocked=%q)", blocked)
if _, err := n1.OnToken(context.Background(), &Token{Cycle: 1, State: *n1.State(), SentAt: 1}); err != nil {
t.Fatal(err)
} }
if claimed != 0 { if blocked != "" {
t.Errorf("Claim ran %d time(s) while the forward was still owned; "+ t.Errorf("nothing should block the first release, got %q", blocked)
"this would create a duplicate proxy", claimed)
} }
if n := len(n1.State().TopologyList()); n != 3 { if p.mine != 3 || p.share != 1 || p.total != 3 || p.alive != 3 {
t.Errorf("topology has %d entries, want 3 (no duplicate created)", n) t.Errorf("pool = %+v, want mine=3 share=1 total=3 alive=3", p)
}
if pick.OwnerID != "n1" {
t.Errorf("must release a forward n1 owns, got owner %q", pick.OwnerID)
} }
} }
// TestBalancerWaitsForRevokeBeforePublishing covers the two-phase guarantee: no // TestReleasingStopsAtTheShare is the bound that keeps the rule from emptying a
// migrate task may exist while the old entry is still present, because that is // node: after shedding one, n1 holds 1 of 3 over 3 nodes, which is no longer more
// what would let the old owner race the new destination. // than the share, and it must stop.
func TestBalancerWaitsForRevokeBeforePublishing(t *testing.T) { func TestReleasingStopsAtTheShare(t *testing.T) {
n1, _, _ := migrateRing(t, 3) n1, _, _ := balanceRing(t, 1, 1) // even 1-1-1 split
b := NewBalancer(n1, 2, 0) b := NewBalancer(n1, 0)
// First step starts phase 1 (revoke) and publishes nothing else. pick, p, blocked := b.assess(1000)
_ = b.Step(1000) if pick != nil {
if b.revoking == nil { t.Errorf("an even 1-1-1 split must not release anything, got %s:%d (blocked=%q)",
t.Fatal("first step should have started a revoke") pick.Local.Name, pick.Link.RemotePort, blocked)
} }
for _, tk := range n1.State().PendingList() { if p.mine != 1 || p.share != 1 {
if tk.Migrate { t.Errorf("pool = %+v, want mine=1 share=1", p)
t.Fatal("a migrate task was published before the revoke completed")
}
}
// Phase 1 completes (the owner drops the entry) → phase 2 publishes.
n1.State().RemoveTopology("web", "frps1", 100)
_ = b.Step(1010)
if b.revoking != nil {
t.Error("revoke phase should have finished")
}
found := false
for _, tk := range n1.State().PendingList() {
if tk.Migrate && tk.Target != "" {
found = true
}
}
if !found {
t.Fatal("after the revoke completed, a targeted migrate task should exist")
} }
} }
// TestBalancerHysteresis: an imbalance must persist before anything moves, so a // TestSingleNodeNeverReleases: a standalone node is the whole cluster, so its
// momentary CPU spike never restarts a worker. // share is everything it holds. Without this the node would shed forwards onto
func TestBalancerHysteresis(t *testing.T) { // nobody (they would sit unclaimed until it re-claimed them).
n1, _, _ := migrateRing(t, 3) func TestSingleNodeNeverReleases(t *testing.T) {
b := NewBalancer(n1, 2, 600) // must be over for 10 minutes e := newTestEngine("solo", true)
e.state.Nodes = []Node{{ID: "solo", Addr: "solo:7500", Alive: true, IsLeader: true}}
_ = b.Step(1000) // first over-threshold observation: only starts the timer
for _, tk := range n1.State().PendingList() {
if tk.Revoke || tk.Migrate {
t.Fatal("acted before the sustained window elapsed")
}
}
// Still inside the window.
_ = b.Step(1300)
if b.revoking != nil {
t.Fatal("acted 5 minutes in; the window is 10")
}
// Past the window.
_ = b.Step(1700)
if b.revoking == nil {
t.Fatal("did not act after the imbalance persisted past the window")
}
}
// TestBalancerQuietWhenBalanced: the threshold means a healthy ring is never
// touched, which is the whole reason for not running this every hour blindly.
func TestBalancerQuietWhenBalanced(t *testing.T) {
n1, _, _ := migrateRing(t, 1)
b := NewBalancer(n1, 2, 0)
for i := 0; i < 5; i++ { for i := 0; i < 5; i++ {
if msg := b.Step(int64(1000 + i*600)); msg != "" { e.state.AddTopology(&Task{
t.Fatalf("a balanced ring produced action: %q", msg) ID: "f" + strconv.Itoa(i), Local: store.Local{Name: "web"},
Remote: store.Remote{Name: "frps1"}, Link: store.Link{RemotePort: 100 + i},
}, "solo")
}
b := NewBalancer(e, 0)
if pick, _, blocked := b.assess(1000); pick != nil {
t.Errorf("a single-node ring must not release; got blocked=%q", blocked)
}
}
// TestBelowFloorNeverReleases: two forwards on a node is the floor. A node
// holding exactly two at an even split of two must keep both — one release per
// node per pass is the granularity, and shedding below the floor would empty it.
func TestBelowFloorNeverReleases(t *testing.T) {
n1, _, _ := balanceRing(t, 2, 2) // 6 total over 3 nodes → share 2
b := NewBalancer(n1, 0)
if pick, p, blocked := b.assess(1000); pick != nil {
t.Errorf("2 owned against share %d must not release (blocked=%q, mine=%d)",
p.share, blocked, p.mine)
}
}
// TestNeverReleasesAStoppedForward: a stopped forward's entry is the carrier of
// the user's stop decision. Releasing it would hand another node a STOPPED
// forward as if it were work to run.
func TestNeverReleasesAStoppedForward(t *testing.T) {
n1, _, _ := balanceRing(t, 3, 0)
// Stop two of n1's three; the only live one must be chosen.
for _, te := range n1.state.ForwardsOwnedBy("n1") {
if te.Link.RemotePort == 100 || te.Link.RemotePort == 101 {
n1.state.UpdateTopologyDisabled("web", "frps1", te.Link.RemotePort, true)
} }
} }
b := NewBalancer(n1, 0)
pick, _, _ := b.assess(1000)
if pick == nil {
t.Fatal("n1 still owns one live forward and is over its share")
}
if pick.Link.RemotePort != 102 {
t.Errorf("released port %d; a stopped forward (100/101) must never be chosen",
pick.Link.RemotePort)
}
// And with every forward stopped there is nothing to release.
n1.state.UpdateTopologyDisabled("web", "frps1", 102, true)
if pick, _, blocked := b.assess(1000); pick != nil {
t.Errorf("all forwards stopped: nothing may be released (blocked=%q)", blocked)
}
}
// TestCandidateChoiceIsDeterministic: the same state must pick the same forward,
// so the decision is explainable and a re-evaluation does not churn a different
// connection each tick.
func TestCandidateChoiceIsDeterministic(t *testing.T) {
n1, _, _ := balanceRing(t, 4, 0)
a := releaseCandidate(n1.State(), "n1")
c := releaseCandidate(n1.State(), "n1")
if a == nil || c == nil {
t.Fatal("expected a candidate")
}
if a.TaskID != c.TaskID || a.Link.RemotePort != c.Link.RemotePort {
t.Errorf("candidate changed between identical evaluations: %s:%d vs %s:%d",
a.TaskID, a.Link.RemotePort, c.TaskID, c.Link.RemotePort)
}
if a.Link.RemotePort != 100 {
t.Errorf("expected the lowest port (100), got %d", a.Link.RemotePort)
}
}
// ---- the runtime signal that makes this safe ----
// TestRevokeLeavesACarrierAndReleaseRemovesIt is the reason the release path has
// to be a distinct task kind.
//
// A user stop keeps the topology entry as the carrier of the disabled flag, so
// the entry is STILL THERE and still attributed to the owner. A rebalance is the
// opposite: the entry must go, because (a) the node's load is derived from
// owned entries so it cannot fall while the entry exists, and (b) every node's
// duplicate guard refuses a creation task for a forward that already has an owner
// — so the re-published task would be dropped by everyone and the forward
// stranded on the node that gave it up.
func TestRevokeLeavesACarrierAndReleaseRemovesIt(t *testing.T) {
n1, _, _ := balanceRing(t, 3, 0)
// A stop marks the entry and keeps it.
n1.state.UpdateTopologyDisabled("web", "frps1", 100, true)
tk := &Task{Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
Link: store.Link{RemotePort: 100}}
if owner := n1.state.TopologyOwner(tk); owner != "n1" {
t.Fatalf("a stop must keep the entry as the flag's carrier, owner=%q", owner)
}
if got := len(n1.state.ForwardsOwnedBy("n1")); got != 3 {
t.Fatalf("a stop must not change the owner's forward count, got %d", got)
}
// A release removes the entry (leaving the other two alone).
if !n1.RemoveTopologyEntry("web", "frps1", 100) {
t.Fatal("release must remove the entry")
}
if owner := n1.state.TopologyOwner(tk); owner != "" {
t.Errorf("after a release the forward must have no owner, got %q", owner)
}
if got := len(n1.state.ForwardsOwnedBy("n1")); got != 2 {
t.Errorf("a release must reduce the owner's count, got %d", got)
}
if got := len(n1.state.TopologyList()); got != 2 {
t.Errorf("the other forwards must be untouched, got %d entries", got)
}
}
// TestReleaseTaskExecutesOnlyOnTheOwner: like a revoke, a release is
// owner-directed. Executing it elsewhere would be an attempt to move somebody
// else's load, which cannot work — the owner's own count is what must fall.
func TestReleaseTaskExecutesOnlyOnTheOwner(t *testing.T) {
n1, n2, _ := balanceRing(t, 3, 0)
released := map[string]int{}
for _, e := range []*Engine{n1, n2} {
h := e.Handler.(*fakeHandler)
h.release = func(ctx context.Context, tk *Task) error {
released[e.ID]++
return nil
}
}
tk := n1.ReleaseTask(store.Local{Name: "web"}, store.Remote{Name: "frps1"},
store.Link{RemotePort: 100})
if tk == nil {
t.Fatal("ReleaseTask returned nil")
}
// Each Engine keeps its OWN state, so a task published on n1 is not visible
// to n2 until a token round carries it. Broadcast it the way the ring does.
var carry State
cb, _ := json.Marshal(n1.state)
_ = json.Unmarshal(cb, &carry)
n2.AdoptState(carry)
n2.state.PendingTasks[tk.ID] = tk
// n2 is the lightest node and does NOT own the forward: it must leave the
// task alone.
if err := n2.runCommands(context.Background(), &Token{}); err != nil {
t.Fatalf("n2 runCommands: %v", err)
}
if released["n2"] != 0 {
t.Errorf("n2 released %d forward(s); a release is owner-directed only", released["n2"])
}
if _, still := n2.state.PendingTasks[tk.ID]; !still {
t.Error("n2 consumed a release it does not own; it must ride the token to the owner")
}
// n1 owns it: it must release, drop the entry, and re-publish the forward.
if err := n1.runCommands(context.Background(), &Token{}); err != nil {
t.Fatalf("n1 runCommands: %v", err)
}
if released["n1"] != 1 {
t.Fatalf("the owner released %d times, want 1", released["n1"])
}
if owner := n1.state.TopologyOwner(&Task{
Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
Link: store.Link{RemotePort: 100},
}); owner != "" {
t.Errorf("the released forward still has owner %q", owner)
}
// And it is back in the ring as an ordinary creation task.
var republished *Task
for _, p := range n1.state.PendingList() {
if !p.Release && !p.Revoke && !p.Restart &&
p.Link.RemotePort == 100 && p.Local.Name == "web" {
republished = p
}
}
if republished == nil {
t.Fatal("the released forward was not re-published as an unclaimed task")
}
if republished.Link.Disabled {
t.Error("the re-published forward must be ENABLED, not marked stopped")
}
}
// ---- dry-run / off ----
// TestDryRunReleasesNothing: dry-run is the default, and its whole value is that
// it can run on a live cluster without disturbing anything.
func TestDryRunReleasesNothing(t *testing.T) {
n1, _, _ := balanceRing(t, 3, 0)
b := NewBalancer(n1, 0)
b.SetMode(BalanceDryRun)
msg := b.Step(1000)
if !strings.Contains(msg, "would release") {
t.Fatalf("dry-run should report what it would do, got %q", msg)
}
if b.InFlight() { if b.InFlight() {
t.Error("balancer started work on a balanced ring") t.Error("dry-run started a release")
}
for _, tk := range n1.state.PendingList() {
if tk.Release {
t.Error("dry-run published a release task")
}
}
if got := len(n1.state.TopologyList()); got != 3 {
t.Errorf("dry-run changed the topology: %d entries, want 3", got)
}
// A second identical tick must not repeat the same line (log hygiene).
if again := b.Step(1010); again != "" {
t.Errorf("dry-run repeated an unchanged decision: %q", again)
} }
} }
// TestDryRunNeverTouchesTheRing: dry-run is the default, and its whole value is // TestModeOffDoesNothing reinforces that OFF means off, not "dry-run".
// that it can run against a live cluster without disturbing anything. It must func TestModeOffDoesNothing(t *testing.T) {
// report the migration it would make and publish nothing at all. n1, _, _ := balanceRing(t, 3, 0)
func TestDryRunNeverTouchesTheRing(t *testing.T) { b := NewBalancer(n1, 0)
n1, _, _ := migrateRing(t, 3) b.SetMode(BalanceOff)
b := NewBalancer(n1, 2, 0) if msg := b.Step(1000); msg != "" {
t.Errorf("mode off reported %q", msg)
// Satisfy the hysteresis so the dry-run has something to report.
b.overSince = map[string]int64{"n1": 1000}
msg := b.stepDryRun(5000)
if msg == "" {
t.Fatal("dry-run reported nothing for a 3-vs-0 imbalance")
}
if !strings.Contains(msg, "would migrate") {
t.Fatalf("dry-run message should be a WOULD report, got %q", msg)
}
// Nothing may have been published.
for _, tk := range n1.State().PendingList() {
if tk.Revoke || tk.Migrate {
t.Fatalf("dry-run published a %v task; it must not touch the ring", tk)
}
} }
if b.InFlight() { if b.InFlight() {
t.Error("dry-run started an in-flight migration") t.Error("mode off started a release")
} }
if n := len(n1.State().TopologyList()); n != 3 { // A typo must not enable releasing.
t.Errorf("topology changed under dry-run: %d entries, want 3", n) b.SetMode("yes-please")
if b.Mode() != BalanceDryRun {
t.Errorf("an unknown mode must fall back to dry-run, got %q", b.Mode())
} }
} }
// TestOnlyLeaderBalances: moving forwards is a cluster-wide decision. If every // ---- end-to-end through the ring machinery ----
// member decided independently they would race on the same token, each revoking
// the same forward. WatchBalance must therefore no-op on non-leaders.
func TestOnlyLeaderBalances(t *testing.T) {
n1, n2, _ := migrateRing(t, 3) // n1 is the leader, n2 is not
if !n1.IsLeader() || n2.IsLeader() {
t.Fatal("precondition broken: expected n1 leader, n2 not")
}
before := len(n1.State().PendingList()) // TestOwnerReleaseThenRingClaimsElsewhere is the whole design in one test: the
ctx, cancel := context.WithCancel(context.Background()) // owner gives the forward up, and the ordinary lowest-load claim then places it
defer cancel() // on an idle node — exactly what the self-directed rule depends on.
// A non-leader balancer with hysteresis already satisfied must still do func TestOwnerReleaseThenRingClaimsElsewhere(t *testing.T) {
// nothing, because Step is only reached through WatchBalance's leader gate. n1, n2, n3 := balanceRing(t, 3, 0)
// Verify the gate itself by running the loop briefly and inspecting state.
bal := NewBalancer(n2, 2, 0)
bal.overSince = map[string]int64{"n1": 1000}
go bal.WatchBalance(ctx, 10*time.Millisecond, false)
time.Sleep(60 * time.Millisecond)
cancel()
if n2.IsLeader() { // n2 and n3 must look idle to the claim path (they do: Forwards=0).
t.Fatal("precondition broken during the run") claimed := map[string]int{}
} for _, e := range []*Engine{n1, n2, n3} {
if got := len(n2.State().PendingList()); got != 0 { h := e.Handler.(*fakeHandler)
t.Errorf("a non-leader published %d task(s); balancing must be leader-only", got) h.claim = func(ctx context.Context, tk *Task) error {
} claimed[e.ID]++
_ = before return nil
}
// TestReportDoesNotDisturbHysteresis: Report is a read-only view that an operator
// can call at will. If it touched the overSince timer, a curious operator
// polling the endpoint would keep resetting the sustained-imbalance clock and
// balancing would never fire. It must observe without perturbing.
func TestReportDoesNotDisturbHysteresis(t *testing.T) {
n1, _, _ := migrateRing(t, 3)
b := NewBalancer(n1, 2, 1800)
// First Step only starts the timer (no action yet).
_ = b.Step(1000)
started, seen := b.overSince["n1"]
if !seen {
t.Fatal("precondition: the overload timer should have been started")
}
// Poll the report repeatedly, as an operator would.
for i := 0; i < 5; i++ {
rep := b.Report(int64(2000 + i*600))
if len(rep.Nodes) != 3 {
t.Fatalf("report listed %d nodes, want 3", len(rep.Nodes))
} }
} }
if got := b.overSince["n1"]; got != started { b := NewBalancer(n1, 0)
t.Errorf("Report reset the overload timer: %d -> %d", started, got) b.SetMode(BalanceOn)
}
// And the real path must still work: past the window it acts. // n1 is over its share → it releases.
_ = b.Step(1000 + 1801) msg := b.Step(1000)
if b.revoking == nil { if !strings.Contains(msg, "releasing") {
t.Error("polling the report suppressed the migration it should have allowed") t.Fatalf("expected a release, got %q", msg)
} }
} if err := n1.runCommands(context.Background(), &Token{}); err != nil {
t.Fatalf("n1 runCommands: %v", err)
// TestReportOnFollowerHasNoPlan: only the leader acts, so a follower must not
// claim it would migrate anything — otherwise an operator querying a random node
// gets an answer that does not reflect what the cluster will actually do.
func TestReportOnFollowerHasNoPlan(t *testing.T) {
_, n2, _ := migrateRing(t, 3) // n2 is not the leader
b := NewBalancer(n2, 2, 0)
b.overSince = map[string]int64{"n1": 1000}
rep := b.Report(9999)
if rep.IsLeader {
t.Fatal("precondition: n2 must not be the leader")
} }
if rep.WouldMigrate != nil { // Find the task the owner re-published for the forward it gave up.
t.Errorf("a follower reported a migration plan: %+v", rep.WouldMigrate) var tk *Task
} for _, p := range n1.state.PendingList() {
if len(rep.Nodes) == 0 { if !p.Release && p.Link.RemotePort == 100 && p.Local.Name == "web" {
t.Error("a follower should still report the node scores (that part is informational)") tk = p
}
}
// TestReportShowsPlanOnceHysteresisSatisfied pins the operator-facing contract:
// before the window elapses there is no plan; after it, the plan appears with
// the concrete forward and both endpoints.
func TestReportShowsPlanOnceHysteresisSatisfied(t *testing.T) {
n1, _, _ := migrateRing(t, 3)
b := NewBalancer(n1, 2, 1800)
// Start the timer but do not satisfy the window.
b.overSince = map[string]int64{"n1": 1000}
if rep := b.Report(1500); rep.WouldMigrate != nil {
t.Errorf("plan reported before the sustained window elapsed: %+v", rep.WouldMigrate)
}
rep := b.Report(1000 + 1801)
if rep.WouldMigrate == nil {
t.Fatal("no plan reported even though the imbalance is past the window")
}
m := rep.WouldMigrate
if m.From == "" || m.To == "" || m.Local == "" || m.Port == 0 {
t.Fatalf("plan is incomplete: %+v", m)
}
if m.From == m.To {
t.Fatalf("plan moves a forward to the node it already lives on: %+v", m)
}
if !rep.Enabled == false && rep.DryRun {
// sanity: default is observe-only
if !rep.DryRun {
t.Error("balancer should default to dry-run")
} }
} }
if tk == nil {
t.Fatal("the owner did not re-publish the forward it released")
}
if got := len(n1.state.ForwardsOwnedBy("n1")); got != 2 {
t.Fatalf("n1 should now own 2, got %d", got)
}
// The forward is unclaimed; the lowest node claims it. n2 is the lightest
// (Forwards=0 and the same mem/net), so it should take it. The task reaches
// n2 on the token — each engine holds its own state, so carry it over the way
// a real round does.
var carried State
cb, _ := json.Marshal(n1.state)
_ = json.Unmarshal(cb, &carried)
n2.AdoptState(carried)
n2.state.PendingTasks[tk.ID] = tk
if err := n2.runCommands(context.Background(), &Token{}); err != nil {
t.Fatalf("n2 runCommands: %v", err)
}
if claimed["n2"] == 0 {
t.Fatal("the idle node did not claim the released forward")
}
// And n1 no longer owns it.
for _, te := range n1.state.ForwardsOwnedBy("n1") {
if te.Link.RemotePort == 100 {
t.Error("n1 still owns the forward it released")
}
}
// The balancer sees the release land and stops holding the release state.
if msg := b.Step(1100); !strings.Contains(msg, "gave up") {
t.Errorf("expected the release to be confirmed, got %q", msg)
}
if b.InFlight() {
t.Error("the release should be complete")
}
} }
// TestReportSurfacesInFlightMigration: the two-phase machine must be visible while // TestCooldownBoundsTheChurn: a release costs a live connection, so one node may
// it runs, otherwise an operator sees a stalled plan and cannot tell a slow revoke // not shed again until the cooldown elapses — even if it is still over its share.
// from a stuck one. func TestCooldownBoundsTheChurn(t *testing.T) {
func TestReportSurfacesInFlightMigration(t *testing.T) { n1, _, _ := balanceRing(t, 4, 0)
n1, _, _ := migrateRing(t, 3) b := NewBalancer(n1, 600) // 10 minutes
b := NewBalancer(n1, 2, 0) b.SetMode(BalanceOn)
_ = b.Step(1000) // starts phase 1
if msg := b.Step(1000); !strings.Contains(msg, "releasing") {
t.Fatalf("first release expected, got %q", msg)
}
n1.RemoveTopologyEntry("web", "frps1", 100) // land it
if msg := b.Step(1010); !strings.Contains(msg, "gave up") {
t.Fatalf("expected confirmation, got %q", msg)
}
// Still over its share (3 of 4 over 3 → share 2), but inside the cooldown.
if pick, _, blocked := b.assess(1020); pick == nil || blocked != "cooldown" {
t.Errorf("expected the second release to be blocked by cooldown, got blocked=%q", blocked)
}
if b.InFlight() {
t.Error("cooldown must not leave a release in flight")
}
// Past the cooldown it may act again.
if pick, _, blocked := b.assess(1000 + 601); pick == nil || blocked != "" {
t.Errorf("past the cooldown the release should be allowed, got blocked=%q", blocked)
}
}
// TestReleaseDeadlineClearsAStuckConfirmation: if the release never lands (leader
// change, ring join), the node must not wedge forever.
func TestReleaseDeadlineClearsAStuckConfirmation(t *testing.T) {
n1, _, _ := balanceRing(t, 3, 0)
b := NewBalancer(n1, 0)
b.SetMode(BalanceOn)
_ = b.Step(1000) // release published, entry NOT removed (nothing executed it)
if !b.InFlight() { if !b.InFlight() {
t.Fatal("precondition: a migration should be in flight") t.Fatal("precondition: a release should be in flight")
} }
rep := b.Report(1001) if msg := b.Step(1000 + releaseDeadlineSec + 1); msg != "" {
if rep.InFlight == nil { t.Errorf("a timed-out release should clear silently, got %q", msg)
t.Fatal("report did not surface the in-flight migration")
} }
if rep.InFlight.Local == "" || rep.InFlight.To == "" { if b.InFlight() {
t.Fatalf("in-flight plan is incomplete: %+v", rep.InFlight) t.Error("the release state should have been cleared after the deadline")
}
}
// ---- read-only report ----
// TestReportDoesNotMutate — polling the endpoint must not start, delay or cancel
// a release. A read-only surface that perturbs what it reports is worse than none.
func TestReportDoesNotMutate(t *testing.T) {
n1, _, _ := balanceRing(t, 3, 0)
b := NewBalancer(n1, 600)
b.SetMode(BalanceOn)
for i := 0; i < 5; i++ {
rep := b.Report()
if rep.SelfForwards != 3 {
t.Fatalf("report changed the node's forward count: %d", rep.SelfForwards)
}
if rep.Mode != BalanceOn {
t.Fatalf("report changed the mode: %q", rep.Mode)
}
}
if b.InFlight() || b.lastReleaseAt != 0 {
t.Error("polling the report started a release")
}
if rep := b.Report(); rep.WouldRelease == nil {
t.Error("the report should still surface the pending decision")
}
}
// TestReportExplainsWhyNothingHappens pins the field that answers the question
// the log cannot: an over-share node that is holding back must say why.
func TestReportExplainsWhyNothingHappens(t *testing.T) {
n1, _, _ := balanceRing(t, 4, 0)
b := NewBalancer(n1, 600)
b.SetMode(BalanceOn)
rep := b.Report()
if !rep.SelfOverShare {
t.Fatalf("4 of 4 forwards over 3 nodes should be over the share (share=%d)", rep.Share)
}
if rep.WouldRelease == nil {
t.Fatal("an over-share node must surface the forward it would give up")
}
if rep.BlockedBy != "" {
t.Errorf("nothing should block the first release, got %q", rep.BlockedBy)
}
// Inside the cooldown the same node reports the block. Report() reads the
// real clock, so set lastReleaseAt relative to now rather than to a fake one.
b.lastReleaseAt = time.Now().Unix() - 10
n1.RemoveTopologyEntry("web", "frps1", 100)
rep = b.Report()
if rep.BlockedBy != "cooldown" {
t.Errorf("after a release the node should report a cooldown block, got %q", rep.BlockedBy)
}
}
// TestReportOnBalancedNodeIsQuiet: the healthy case must be a real answer, not an
// empty one — otherwise "balanced" and "balancing is switched off" look identical.
func TestReportOnBalancedNodeIsQuiet(t *testing.T) {
n1, _, _ := balanceRing(t, 1, 1)
b := NewBalancer(n1, 0)
b.SetMode(BalanceOn)
rep := b.Report()
if rep.WouldRelease != nil {
t.Errorf("a balanced node reported a release: %+v", rep.WouldRelease)
}
if rep.SelfOverShare {
t.Error("1 of 3 over 3 nodes is not over the share")
}
if rep.TotalForwards != 3 || rep.AliveNodes != 3 || rep.Share != 1 {
t.Errorf("report pool = total %d alive %d share %d, want 3/3/1",
rep.TotalForwards, rep.AliveNodes, rep.Share)
} }
} }

View File

@ -68,13 +68,19 @@ type Task struct {
// duplicate-claim guard would discard it. The owner applies this one // duplicate-claim guard would discard it. The owner applies this one
// unconditionally: re-mark the entry enabled and (re)spawn the worker. // unconditionally: re-mark the entry enabled and (re)spawn the worker.
Restart bool `json:"restart,omitempty"` Restart bool `json:"restart,omitempty"`
// Migrate marks a LOAD-BALANCING task and Target names the node that must take // Release marks a REBALANCING task: the OWNER of the forward gives it up so
// the forward over. Unlike a plain creation task — which any lowest-load node // the ring can place it elsewhere. Only the owning node may act on it — it is
// may claim — a migrate is addressed: it exists precisely because the node that // that node's own load which has to fall, and only it can make that true (its
// is currently lightest (usually the node we just took the forward FROM) would // topology entry is local bookkeeping no peer may write).
// otherwise claim it straight back, undoing the migration. //
Migrate bool `json:"migrate,omitempty"` // It is deliberately NOT the same thing as a user stop (Revoke). A stop keeps
Target string `json:"target,omitempty"` // the topology entry as the carrier of the disabled flag and flags the link
// stopped; a release REMOVES the entry and leaves the link enabled, because
// the forward is about to run somewhere else. Marking it stopped instead would
// make the re-published creation task look like a duplicate of an existing
// forward, and every node would refuse to claim it — stranding the forward on
// the node that just gave it up.
Release bool `json:"release,omitempty"`
// RemoveNode: node ID to remove from the ring; the node self-removes when // RemoveNode: node ID to remove from the ring; the node self-removes when
// the command reaches it via the token. // the command reaches it via the token.
RemoveNode string `json:"removeNode,omitempty"` RemoveNode string `json:"removeNode,omitempty"`

View File

@ -3,321 +3,316 @@ package cluster
import ( import (
"context" "context"
"log" "log"
"sort"
"strconv" "strconv"
"time" "time"
"webui4frpc/internal/store" "webui4frpc/internal/store"
) )
// fwdKey identifies a forward for balancing purposes. // The rebalance rule, deliberately simple and self-limiting:
type fwdKey struct { //
local, remote string // release one forward when Forwards >= minForwardsToShare AND Forwards > share
port int //
// 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
} }
// planBalance decides whether one forward should move off an overloaded owner. // releaseState is one release this node has published and is waiting to see
// // leave its own topology.
// The rule is deliberately conservative: type releaseState struct {
// - never touch a stopped forward (disabled) — moving it would resurrect work local, remote string
// the user explicitly turned off; port int
// - never touch a forward whose owner is not actually above the threshold; startedAt int64
// - 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) { // Balancer lets each node shed its own share of the cluster's forwards.
// Candidate destinations: alive, non-leader-free (any), scored by Load. //
type cand struct { // It publishes only ONE kind of thing: a release of a forward it owns. There is
id string // no targeted hand-off and no coordinator, because "drop it and let the ring's
score float64 // normal lowest-load claim pick it up" is already the ring's only mechanism for
forwards int // 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
} }
var dests []cand return b
total := 0 }
for i := range s.Nodes {
n := &s.Nodes[i] // 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 { if !n.Alive {
continue continue
} }
dests = append(dests, cand{n.ID, n.Load.Score(), n.Load.Forwards}) p.alive++
total += n.Load.Forwards if n.ID == b.engine.ID {
} p.total += p.mine
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 continue
} }
if pick == nil || te.Link.RemotePort < pick.Link.RemotePort { p.total += n.Load.Forwards
pick = te
}
} }
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 { if pick == nil {
return nil, "", false // everything it owns is stopped return nil, p, "nothing-active"
} }
return pick, dest.id, true // 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 {
// Balancer moves forwards off an overloaded node, one at a time. return pick, p, "in-flight"
//
// 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 { if b.lastReleaseAt > 0 && now-b.lastReleaseAt < b.cooldownSec {
sustainedSec = 0 return pick, p, "cooldown"
}
return &Balancer{
engine: e, ratio: ratio, sustainedSec: sustainedSec,
dryRun: true, // safe default: report, never move
overSince: map[string]int64{},
} }
return pick, p, ""
} }
// SetDryRun selects the mode. Dry-run (the default) only reports. // releaseCandidate picks which of this node's forwards to give up.
func (b *Balancer) SetDryRun(v bool) { b.dryRun = v } //
// Deterministic (lowest remote port, then task id) so the choice is explainable
// DryRun reports whether the balancer is in observe-only mode. // and repeatable: a node that re-evaluates the same state picks the same
func (b *Balancer) DryRun() bool { return b.dryRun } // 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
// InFlight reports whether a migration is currently between phases. // as if it were work to run.
func (b *Balancer) InFlight() bool { func releaseCandidate(s *State, owner string) *TopoEntry {
return b.revoking != nil || b.publishing != nil 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]
} }
// Step advances the state machine by one cycle. now is the current unix second. // landed reports whether the forward this node released has actually left it.
// It returns a short human-readable description when something happened. //
// 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 { func (b *Balancer) Step(now int64) string {
// --- phase 2: did phase 1 finish? --- if b.releasing != nil {
if b.revoking != nil { if b.landed() {
if b.engine.state.TopologyOwner(&Task{ done := b.releasing
Local: store.Local{Name: b.revoking.local}, b.releasing = nil
Remote: store.Remote{Name: b.revoking.remote}, return "gave up " + done.local + ":" + strconv.Itoa(done.port) + " for rebalancing"
Link: store.Link{RemotePort: b.revoking.port}, }
}) != "" { if now-b.releasing.startedAt > releaseDeadlineSec {
// Still owned: the revoke has not completed yet. Wait, unless it has log.Printf("ring[%s] release of %s:%d not confirmed within %ds; clearing (forward is still ours)",
// been far too long, in which case something is wrong with the ring. b.engine.ID, b.releasing.local, b.releasing.port, releaseDeadlineSec)
if now < b.revoking.deadlineAt { b.releasing = nil
return "" }
} 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 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 "" return ""
} }
// Entry gone — the old worker is stopped. Publish for the named target. b.reported = desc
m := b.revoking if blocked != "" {
b.revoking = nil return "would release " + desc + " [blocked: " + blocked + "]"
b.publishing = m }
b.publishMigration() return "would release " + desc
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 --- // Mode is ON. Blocked reasons still stop us.
if b.publishing != nil { if blocked != "" {
// 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 "" return ""
} }
t := b.engine.ReleaseTask(pick.Local, pick.Remote, pick.Link)
// --- nothing in flight: look for work --- if t == nil {
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 "" return ""
} }
if since, ok := b.overSince[pick.OwnerID]; !ok || now-since < int64(b.sustainedSec) { b.releasing = &releaseState{
return "" local: pick.Local.Name, remote: pick.Remote.Name,
port: pick.Link.RemotePort, startedAt: now,
} }
// Phase 1: publish the revoke and remember what we are waiting for. b.lastReleaseAt = now
m := &pendingMigration{ b.reported = ""
local: pick.Local.Name, remote: pick.Remote.Name, port: pick.Link.RemotePort, return "releasing " + desc
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() { // Watch drives Step on a ticker until ctx is cancelled.
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 // It runs on EVERY node, not just the leader. Balancing here is not a cluster
// forwards between members), and letting every member decide independently would // decision being delegated to one node — it is each node applying the same
// mean several migrations racing for the same forward — the ring is a shared // replicated rule to its own forwards. Two nodes never contend for the same
// token, so two nodes must not both publish a revoke for the same entry. // forward because a node can only release what it owns.
// //
// dryRun makes the balancer report what it WOULD do without publishing anything. // The first check happens after one interval rather than immediately: a node
// That is the mode to use first on a live cluster: it answers "is my real data // that has just booted adopts the ring's topology and its own load in the first
// going to trip the threshold, and what would it move?" before anything on the // token round, and deciding before that would act on a half-known pool.
// ring is disturbed. func (b *Balancer) Watch(ctx context.Context, interval time.Duration) {
func (b *Balancer) WatchBalance(ctx context.Context, interval time.Duration, dryRun bool) {
tick := time.NewTicker(interval) tick := time.NewTicker(interval)
defer tick.Stop() defer tick.Stop()
for { for {
@ -326,69 +321,46 @@ func (b *Balancer) WatchBalance(ctx context.Context, interval time.Duration, dry
return return
case <-tick.C: case <-tick.C:
} }
if !b.engine.IsLeader() { if msg := b.Step(time.Now().Unix()); msg != "" {
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) log.Printf("ring[%s] balance: %s", b.engine.ID, msg)
} }
} }
} }
// stepDryRun reports the migration the balancer would start right now, without // ---- read-only report for the API/UI ----
// 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. // BalanceReport is what this node thinks about balancing right now.
// //
// It exists because the balancer's decision is otherwise only visible in the // It is per-node by nature: every member has its own answer, and the answer that
// log, and only when an imbalance exists AND the hysteresis window has elapsed. // matters is the one on the node doing (or refusing to do) the releasing. The
// An operator asking "is this cluster balanced?" gets silence in exactly the // peer list is included because "am I over my share?" is only meaningful next to
// healthy case, which is indistinguishable from "balancing is switched off". // what everyone else holds.
type BalanceReport struct { type BalanceReport struct {
Enabled bool `json:"enabled"` Mode string `json:"mode"` // off | dry-run | on
DryRun bool `json:"dryRun"` Enabled bool `json:"enabled"` // mode != off
Ratio float64 `json:"ratio"` DryRun bool `json:"dryRun"` // mode != on
SustainSec int64 `json:"sustainSec"` IsMember bool `json:"isMember"`
// Only the leader acts; on a follower these stay false/empty. Self string `json:"self"`
IsLeader bool `json:"isLeader"`
// Current per-node scores. 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"` 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 { type BalanceNodeReport struct {
@ -401,24 +373,47 @@ type BalanceNodeReport struct {
Score float64 `json:"score"` Score float64 `json:"score"`
} }
type BalanceMigrationPlan struct { type BalanceReleasePlan struct {
Local string `json:"local"` Local string `json:"local"`
Remote string `json:"remote"` Remote string `json:"remote"`
Port int `json:"remotePort"` RemotePort int `json:"remotePort"`
From string `json:"from"` Owner string `json:"owner"`
To string `json:"to"`
} }
// Report renders the balancer's current view. It never mutates anything — // Report renders the current view. It never mutates anything — in particular it
// note that it does NOT touch overSince, so calling it cannot start or reset // does NOT touch lastReleaseAt or the release state, so an operator polling this
// the hysteresis timer that the real Step path uses. // cannot start, delay or cancel a release. (A read-only endpoint that perturbs
func (b *Balancer) Report(now int64) BalanceReport { // the behaviour it reports is worse than no endpoint.)
func (b *Balancer) Report() BalanceReport {
now := time.Now().Unix()
rep := BalanceReport{ rep := BalanceReport{
Enabled: !b.dryRun, Mode: b.mode,
DryRun: b.dryRun, Enabled: b.mode != BalanceOff,
Ratio: b.ratio, DryRun: b.mode != BalanceOn,
SustainSec: int64(b.sustainedSec), IsMember: b.engine.IsMember(),
IsLeader: b.engine.IsLeader(), 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 { for i := range b.engine.state.Nodes {
n := &b.engine.state.Nodes[i] n := &b.engine.state.Nodes[i]
@ -428,46 +423,5 @@ func (b *Balancer) Report(now int64) BalanceReport {
Score: n.Load.Score(), 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 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
}

View File

@ -24,6 +24,12 @@ type Handler interface {
// creation path (idempotency guard) and the duplicate-claim guard refuse to // creation path (idempotency guard) and the duplicate-claim guard refuse to
// re-create an existing forward. // re-create an existing forward.
Restart(ctx context.Context, tk *Task) error Restart(ctx context.Context, tk *Task) error
// Release stops this node's worker for a forward it is giving up for
// rebalancing, WITHOUT marking the link stopped. The forward is about to run
// on another node, so a stopped flag would be wrong: every node's duplicate
// guard reads it and would then refuse to claim the re-published task,
// stranding the forward on the node that just gave it up.
Release(ctx context.Context, tk *Task) error
RuntimeLoad() Load RuntimeLoad() Load
} }
@ -406,16 +412,24 @@ func (e *Engine) runCommands(ctx context.Context, tk *Token) error {
} }
continue // owned elsewhere — ride to the owner continue // owned elsewhere — ride to the owner
} }
if t.Migrate { if t.Release {
// Addressed to one node. The whole point is that "whoever is // Owner-directed, like a revoke: only the node holding the forward
// lightest" is the WRONG answer here: the node we just revoked the // may give it up. A release executed anywhere else would be an
// forward from is the one that became lightest, and it would claim // attempt to move somebody else's load, which cannot work — the
// the task straight back. // owner's own count is what has to fall, and its topology entry is
if t.Target == e.ID { // local bookkeeping no peer may write.
//
// A release whose forward is already gone is an idempotent no-op
// consumed by the lowest node, so it cannot ride forever if its
// owner departed first.
if owner := e.state.TopologyOwner(t); owner == e.ID {
target = t
break
} else if owner == "" && selfIsLowest {
target = t target = t
break break
} }
continue // ride to the target continue // owned elsewhere — ride to the owner
} }
if t.Restart { if t.Restart {
// Owner-directed, exactly like a revoke: only the node holding the // Owner-directed, exactly like a revoke: only the node holding the
@ -549,7 +563,11 @@ func (e *Engine) runCommands(ctx context.Context, tk *Token) error {
// exists precisely to re-activate a forward whose entry is still present // exists precisely to re-activate a forward whose entry is still present
// (a stop keeps the entry as the flag's carrier), so applying the guard // (a stop keeps the entry as the flag's carrier), so applying the guard
// would swallow every re-enable and leave the forward stopped forever. // would swallow every re-enable and leave the forward stopped forever.
if !claimed.Restart && !claimed.Migrate { // A RELEASE is the second deliberate exception, for the same reason as a
// restart: it is addressed to the node that ALREADY owns the forward, so
// the entry is still there by definition. Applying the guard would drop
// every release and a node could never shed anything.
if !claimed.Restart && !claimed.Release {
if owner := e.state.TopologyOwner(claimed); owner != "" { if owner := e.state.TopologyOwner(claimed); owner != "" {
log.Printf("ring[%s] drop duplicate task %s: %s→%s:%d already owned by %s", log.Printf("ring[%s] drop duplicate task %s: %s→%s:%d already owned by %s",
e.ID, claimed.ID, claimed.Local.Name, claimed.Remote.Name, e.ID, claimed.ID, claimed.Local.Name, claimed.Remote.Name,
@ -557,33 +575,46 @@ func (e *Engine) runCommands(ctx context.Context, tk *Token) error {
continue continue
} }
} }
if claimed.Migrate { if claimed.Release {
// We are the designated destination. The forward must NOT already be // Only reached when this node owns the forward (or it is already gone
// here: the two-phase migration guarantees the old entry was gone before // and we are the fallback consumer — see the selection loop).
// this task was published. If one somehow exists, drop the task rather owner := e.state.TopologyOwner(claimed)
// than creating a second topology entry for the same forward. if owner != e.ID {
if owner := e.state.TopologyOwner(claimed); owner != "" { // The forward vanished before we got here; nothing to give up.
log.Printf("ring[%s] drop migrate %s: already owned by %s", e.ID, claimed.ID, owner)
continue continue
} }
// Stop our worker. Deliberately NOT Handler.Revoke: that marks the
// link stopped, and the forward is about to run elsewhere — a
// "stopped" flag would make every node treat the re-published task as
// a duplicate and refuse to claim it.
if e.Handler != nil { if e.Handler != nil {
if err := e.Handler.Claim(ctx, claimed); err != nil { if err := e.Handler.Release(ctx, claimed); err != nil {
// Put it back so another cycle can retry — the forward is log.Printf("ring[%s] release %s: %v", e.ID, claimed.ID, err)
// currently owned by nobody, so losing this task would strand it.
e.state.PendingTasks[claimed.ID] = claimed
log.Printf("ring[%s] migrate %s failed: %v", e.ID, claimed.ID, err)
break
} }
} }
e.state.AddTopology(claimed, e.ID) // Drop the entry so the forward has no owner. This is what makes the
// load actually fall (ForwardsOwnedBy counts entries) and what lets a
// fresh creation task pass the duplicate guard above.
e.RemoveTopologyEntry(claimed.Local.Name, claimed.Remote.Name, claimed.Link.RemotePort)
// Hand it back to the ring as an ordinary unclaimed forward. It is
// published by THIS node so it rides the very token we are holding —
// no coordinator, no addressed hand-off.
link := claimed.Link
link.Disabled = false
if t := e.SubmitTask(claimed.Local, claimed.Remote, link); t != nil {
log.Printf("ring[%s] released %s: %s→%s:%d back to the ring as %s",
e.ID, claimed.ID, claimed.Local.Name, claimed.Remote.Name,
claimed.Link.RemotePort, t.ID)
} else {
log.Printf("ring[%s] released %s but the re-publish was deduped",
e.ID, claimed.ID)
}
if e.Log != nil { if e.Log != nil {
_, _ = e.Log.Append(e.ID, LogForwardAdd, map[string]any{ _, _ = e.Log.Append(e.ID, LogForwardStop, map[string]any{
"taskId": claimed.ID, "local": claimed.Local.Name, "remote": claimed.Remote.Name, "taskId": claimed.ID, "local": claimed.Local.Name,
"migratedFrom": e.state.TopologyOwner(claimed), "remote": claimed.Remote.Name, "released": true,
}) })
} }
log.Printf("ring[%s] migrated in %s: %s→%s:%d", e.ID, claimed.ID,
claimed.Local.Name, claimed.Remote.Name, claimed.Link.RemotePort)
continue continue
} }
if claimed.Restart { if claimed.Restart {
@ -1000,27 +1031,23 @@ func (e *Engine) SubmitRestart(local store.Local, remote store.Remote, link stor
return t return t
} }
// submitMigrate publishes a load-balancing task addressed to one specific node. // ReleaseTask publishes a rebalancing task asking the OWNER of this forward to
// give it up. The task rides the token like any other command; it is not
// addressed to a destination, because the ring's normal lowest-load claim
// decides where the forward goes next.
// //
// It is addressed on purpose. A plain creation task is claimed by "whichever node // It goes through AddPending rather than writing PendingTasks directly. That is
// is currently lightest", and after a revoke the node we just took the forward // not cosmetic: OnToken rebuilds local-only pending commands from the ids it
// from IS the lightest — so an untargeted task would be claimed straight back by // never published, and only the Add* helpers record them — a task injected by
// the old owner and the migration would achieve nothing but a worker restart. // hand was silently dropped at the next state adoption, so it never reached the
func (e *Engine) submitMigrate(local store.Local, remote store.Remote, link store.Link, target string) *Task { // ring at all.
func (e *Engine) ReleaseTask(local store.Local, remote store.Remote, link store.Link) *Task {
link.Disabled = false link.Disabled = false
t := &Task{ t := e.state.AddPending(local, remote, link)
ID: e.state.NextTaskID(), if t == nil {
Local: local, return nil
Remote: remote,
Link: link,
Created: time.Now().Unix(),
Migrate: true,
Target: target,
} }
if e.state.PendingTasks == nil { t.Release = true
e.state.PendingTasks = map[string]*Task{}
}
e.state.PendingTasks[t.ID] = t
return t return t
} }

View File

@ -12,6 +12,14 @@ type fakeHandler struct {
claim func(ctx context.Context, tk *Task) error claim func(ctx context.Context, tk *Task) error
revoke func(ctx context.Context, tk *Task) error revoke func(ctx context.Context, tk *Task) error
restart func(ctx context.Context, tk *Task) error restart func(ctx context.Context, tk *Task) error
release func(ctx context.Context, tk *Task) error
}
func (h *fakeHandler) Release(ctx context.Context, tk *Task) error {
if h.release != nil {
return h.release(ctx, tk)
}
return nil
} }
func (h *fakeHandler) Revoke(ctx context.Context, tk *Task) error { func (h *fakeHandler) Revoke(ctx context.Context, tk *Task) error {

View File

@ -21,6 +21,12 @@ type AppHandler struct {
// stopped flag + bring the worker back). Required because a stop keeps the // stopped flag + bring the worker back). Required because a stop keeps the
// topology entry, so the normal claim path will not re-create it. // topology entry, so the normal claim path will not re-create it.
RestartFn func(ctx context.Context, tk *Task) error RestartFn func(ctx context.Context, tk *Task) error
// ReleaseFn stops this node's worker for a forward being handed back to the
// ring, leaving the link ENABLED. When nil, Release falls back to RevokeFn —
// which also stops the worker, but additionally marks the link stopped; that
// is the pre-existing stop semantics and is wrong for a rebalance, so the app
// should provide this.
ReleaseFn func(ctx context.Context, tk *Task) error
} }
// RuntimeLoad implements Handler. // RuntimeLoad implements Handler.
@ -48,6 +54,22 @@ func (a *AppHandler) Revoke(ctx context.Context, tk *Task) error {
return a.RevokeFn(ctx, tk) return a.RevokeFn(ctx, tk)
} }
// Release implements Handler.
func (a *AppHandler) Release(ctx context.Context, tk *Task) error {
if a.ReleaseFn != nil {
return a.ReleaseFn(ctx, tk)
}
// Fall back to the stop path so a host that has not provided ReleaseFn still
// stops the worker (the forward keeps running otherwise). The link ends up
// flagged stopped, which the claim path may then refuse — the host is warned
// by the ring's own "re-publish was deduped" log line rather than silently
// leaving two workers up.
if a.RevokeFn != nil {
return a.RevokeFn(ctx, tk)
}
return nil
}
// Restart implements Handler. // Restart implements Handler.
func (a *AppHandler) Restart(ctx context.Context, tk *Task) error { func (a *AppHandler) Restart(ctx context.Context, tk *Task) error {
if a.RestartFn == nil { if a.RestartFn == nil {

View File

@ -655,7 +655,7 @@ func (h *Handler) handleClusterBalance(w http.ResponseWriter, r *http.Request) {
http.Error(w, "load balancer not enabled", http.StatusNotFound) http.Error(w, "load balancer not enabled", http.StatusNotFound)
return return
} }
writeJSON(w, http.StatusOK, h.Balancer.Report(time.Now().Unix())) writeJSON(w, http.StatusOK, h.Balancer.Report())
} }
// handleClusterJoin accepts a newcomer join request: the target node inserts // handleClusterJoin accepts a newcomer join request: the target node inserts