diff --git a/cmd/webui4frpc/main.go b/cmd/webui4frpc/main.go index b66804a..73fa04b 100644 --- a/cmd/webui4frpc/main.go +++ b/cmd/webui4frpc/main.go @@ -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) 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 // must NOT run the claim bookkeeping (ClaimFn), which appends links and // 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 // expose its read-only report. Mode/parameters are applied further down, once // the W4F_BALANCE_* environment has been parsed. - balancer := cluster.NewBalancer(ring, 2, 1800) + balancer := cluster.NewBalancer(ring, 600) h := &httpapi.Handler{ Store: st, @@ -562,41 +577,40 @@ func main() { go ring.WatchTokenLoss(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 - // frpc worker and dropping its live connections. That is fine occasionally - // (a genuinely hot node) but destructive as a reflex, so it ships off and is - // opted into explicitly: + // Every node runs this check against its OWN forwards. There is no + // coordinator: 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 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) - // W4F_BALANCE_MODE=on actually migrate - // W4F_BALANCE_RATIO=2 overload = Score above RATIO x the mean - // W4F_BALANCE_SUSTAIN=1800 seconds an imbalance must persist + // A release stops a worker and drops a live connection, so acting on the ring + // is opt-in: + // + // 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 - // - // 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"))) 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 if balInterval < time.Second { balInterval = 10 * time.Second } - balancer.SetRatioAndSustain(balRatio, balSustain) - if balMode == "on" { - log.Printf("ring[%s] load balancing ENABLED (ratio=%.2f sustain=%ds interval=%s)", - selfID, balRatio, balSustain, balInterval) - } else { - log.Printf("ring[%s] load balancing in DRY-RUN (set W4F_BALANCE_MODE=on to enable; "+ - "ratio=%.2f sustain=%ds interval=%s)", selfID, balRatio, balSustain, balInterval) - } - go balancer.WatchBalance(clusterCtx, balInterval, balMode != "on") + balCooldown := int64(envFloat("W4F_BALANCE_COOLDOWN", 600)) + balancer.SetCooldown(balCooldown) + balancer.SetMode(balMode) + log.Printf("ring[%s] load balancing %s (cooldown=%ds interval=%s, self-directed)", + selfID, strings.ToUpper(balMode), balCooldown, balInterval) + go balancer.Watch(clusterCtx, balInterval) // Startup reconcile: every node (leader or not) checks its own topology // entries and re-claims any whose per-forward frpc worker process is diff --git a/internal/cluster/balance_test.go b/internal/cluster/balance_test.go index a5c9236..d77d8b4 100644 --- a/internal/cluster/balance_test.go +++ b/internal/cluster/balance_test.go @@ -3,6 +3,7 @@ package cluster import ( "context" "encoding/json" + "strconv" "strings" "testing" "time" @@ -30,124 +31,6 @@ func mustState(t *testing.T, s State) *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. func migrateRing(t *testing.T, ownCount int) (*Engine, *Engine, *Engine) { t.Helper() @@ -194,330 +77,545 @@ func migrateRing(t *testing.T, ownCount int) (*Engine, *Engine, *Engine) { return n1, n2, n3 } -// TestMigrateTaskIsAddressedNotClaimedByLightest is the core safety test. -// -// 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 +// ---- helpers ---- - n1.submitMigrate(store.Local{Name: "web"}, store.Remote{Name: "frps1"}, - store.Link{RemotePort: 100}, "n2") - - // Give each node a DISTINCT load via its handler. loadSnapshot() refreshes - // mem/net from Handler.RuntimeLoad() on every OnToken, so seeding - // State.Nodes[].Load by hand is silently overwritten — which is why an - // earlier version of this test had n2 and n3 tied and therefore passed even - // with the targeted delivery removed (the "lightest" node happened to be the - // target). n3 is now strictly the lightest, so an untargeted rule would hand - // the task to n3 and fail here. - claimedBy := map[string]int{} - 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 - }} +// balanceRing builds a 3-node ring with `owned` forwards on n1 (the evaluation +// 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 +// topology, because that is what a real token round leaves behind. +func balanceRing(t *testing.T, owned, others int) (*Engine, *Engine, *Engine) { + t.Helper() + mk := func(id string, leader bool) *Engine { + h := &fakeHandler{load: Load{MemPct: 5, NetPct: 5}} + return NewEngine(id, id+":7500", "u", "p", "0.1.3", nil, h, + func(ctx context.Context, next string, tk *Token) error { return nil }, + id+":7500", leader, "") } - n1.Handler = mk("n1", 50) - n2.Handler = mk("n2", 40) // the TARGET — heavier than n3 - n3.Handler = mk("n3", 1) // strictly the lightest + n1, n2, n3 := mk("n1", true), mk("n2", false), mk("n3", false) + + 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 - b, _ := json.Marshal(n1.State()) + b, _ := json.Marshal(n1.state) _ = json.Unmarshal(b, &snap) for _, e := range []*Engine{n1, n2, n3} { e.AdoptState(snap) } - for i := 0; i < 3; i++ { - 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), - }); err != nil { - t.Fatal(err) - } + return n1, n2, n3 +} + +// ---- the rule ---- + +// 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 -// arrives before the revoke completed must NOT create a second topology entry -// for the same forward. This is the failure that would otherwise leave two nodes -// running workers for one remotePort — which frps rejects with -// "proxy already exists" (observed 14k times before that was fixed). -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 }} +// TestOwnerReleasesWhenOverItsShare: 3 forwards on n1, none elsewhere → n1 holds +// 3 against a share of 1 and must give one up. +func TestOwnerReleasesWhenOverItsShare(t *testing.T) { + n1, _, _ := balanceRing(t, 3, 0) + b := NewBalancer(n1, 0) - // Entry still owned by n1 — the revoke has NOT landed. - n1.submitMigrate(store.Local{Name: "web"}, store.Remote{Name: "frps1"}, - store.Link{RemotePort: 100}, "n1") - - if _, err := n1.OnToken(context.Background(), &Token{Cycle: 1, State: *n1.State(), SentAt: 1}); err != nil { - t.Fatal(err) + pick, p, blocked := b.assess(1000) + if pick == nil { + t.Fatalf("n1 holds 3 of 3 forwards and should release one (blocked=%q)", blocked) } - if claimed != 0 { - t.Errorf("Claim ran %d time(s) while the forward was still owned; "+ - "this would create a duplicate proxy", claimed) + if blocked != "" { + t.Errorf("nothing should block the first release, got %q", blocked) } - if n := len(n1.State().TopologyList()); n != 3 { - t.Errorf("topology has %d entries, want 3 (no duplicate created)", n) + if p.mine != 3 || p.share != 1 || p.total != 3 || p.alive != 3 { + 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 -// migrate task may exist while the old entry is still present, because that is -// what would let the old owner race the new destination. -func TestBalancerWaitsForRevokeBeforePublishing(t *testing.T) { - n1, _, _ := migrateRing(t, 3) - b := NewBalancer(n1, 2, 0) +// TestReleasingStopsAtTheShare is the bound that keeps the rule from emptying a +// node: after shedding one, n1 holds 1 of 3 over 3 nodes, which is no longer more +// than the share, and it must stop. +func TestReleasingStopsAtTheShare(t *testing.T) { + n1, _, _ := balanceRing(t, 1, 1) // even 1-1-1 split + b := NewBalancer(n1, 0) - // First step starts phase 1 (revoke) and publishes nothing else. - _ = b.Step(1000) - if b.revoking == nil { - t.Fatal("first step should have started a revoke") + pick, p, blocked := b.assess(1000) + if pick != nil { + t.Errorf("an even 1-1-1 split must not release anything, got %s:%d (blocked=%q)", + pick.Local.Name, pick.Link.RemotePort, blocked) } - for _, tk := range n1.State().PendingList() { - if tk.Migrate { - 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") + if p.mine != 1 || p.share != 1 { + t.Errorf("pool = %+v, want mine=1 share=1", p) } } -// TestBalancerHysteresis: an imbalance must persist before anything moves, so a -// momentary CPU spike never restarts a worker. -func TestBalancerHysteresis(t *testing.T) { - n1, _, _ := migrateRing(t, 3) - b := NewBalancer(n1, 2, 600) // must be over for 10 minutes - - _ = 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) +// TestSingleNodeNeverReleases: a standalone node is the whole cluster, so its +// share is everything it holds. Without this the node would shed forwards onto +// nobody (they would sit unclaimed until it re-claimed them). +func TestSingleNodeNeverReleases(t *testing.T) { + e := newTestEngine("solo", true) + e.state.Nodes = []Node{{ID: "solo", Addr: "solo:7500", Alive: true, IsLeader: true}} for i := 0; i < 5; i++ { - if msg := b.Step(int64(1000 + i*600)); msg != "" { - t.Fatalf("a balanced ring produced action: %q", msg) + e.state.AddTopology(&Task{ + 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() { - 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 -// that it can run against a live cluster without disturbing anything. It must -// report the migration it would make and publish nothing at all. -func TestDryRunNeverTouchesTheRing(t *testing.T) { - n1, _, _ := migrateRing(t, 3) - b := NewBalancer(n1, 2, 0) - - // 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) - } +// TestModeOffDoesNothing reinforces that OFF means off, not "dry-run". +func TestModeOffDoesNothing(t *testing.T) { + n1, _, _ := balanceRing(t, 3, 0) + b := NewBalancer(n1, 0) + b.SetMode(BalanceOff) + if msg := b.Step(1000); msg != "" { + t.Errorf("mode off reported %q", msg) } 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 { - t.Errorf("topology changed under dry-run: %d entries, want 3", n) + // A typo must not enable releasing. + 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 -// 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") - } +// ---- end-to-end through the ring machinery ---- - before := len(n1.State().PendingList()) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - // A non-leader balancer with hysteresis already satisfied must still do - // nothing, because Step is only reached through WatchBalance's leader gate. - // 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() +// TestOwnerReleaseThenRingClaimsElsewhere is the whole design in one test: the +// owner gives the forward up, and the ordinary lowest-load claim then places it +// on an idle node — exactly what the self-directed rule depends on. +func TestOwnerReleaseThenRingClaimsElsewhere(t *testing.T) { + n1, n2, n3 := balanceRing(t, 3, 0) - if n2.IsLeader() { - t.Fatal("precondition broken during the run") - } - if got := len(n2.State().PendingList()); got != 0 { - t.Errorf("a non-leader published %d task(s); balancing must be leader-only", got) - } - _ = before -} - -// 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)) + // n2 and n3 must look idle to the claim path (they do: Forwards=0). + claimed := map[string]int{} + for _, e := range []*Engine{n1, n2, n3} { + h := e.Handler.(*fakeHandler) + h.claim = func(ctx context.Context, tk *Task) error { + claimed[e.ID]++ + return nil } } - if got := b.overSince["n1"]; got != started { - t.Errorf("Report reset the overload timer: %d -> %d", started, got) - } + b := NewBalancer(n1, 0) + b.SetMode(BalanceOn) - // And the real path must still work: past the window it acts. - _ = b.Step(1000 + 1801) - if b.revoking == nil { - t.Error("polling the report suppressed the migration it should have allowed") + // n1 is over its share → it releases. + msg := b.Step(1000) + if !strings.Contains(msg, "releasing") { + t.Fatalf("expected a release, got %q", msg) } -} - -// 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 err := n1.runCommands(context.Background(), &Token{}); err != nil { + t.Fatalf("n1 runCommands: %v", err) } - if rep.WouldMigrate != nil { - t.Errorf("a follower reported a migration plan: %+v", rep.WouldMigrate) - } - if len(rep.Nodes) == 0 { - t.Error("a follower should still report the node scores (that part is informational)") - } -} - -// 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") + // Find the task the owner re-published for the forward it gave up. + var tk *Task + for _, p := range n1.state.PendingList() { + if !p.Release && p.Link.RemotePort == 100 && p.Local.Name == "web" { + tk = p } } + 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 -// it runs, otherwise an operator sees a stalled plan and cannot tell a slow revoke -// from a stuck one. -func TestReportSurfacesInFlightMigration(t *testing.T) { - n1, _, _ := migrateRing(t, 3) - b := NewBalancer(n1, 2, 0) - _ = b.Step(1000) // starts phase 1 +// TestCooldownBoundsTheChurn: a release costs a live connection, so one node may +// not shed again until the cooldown elapses — even if it is still over its share. +func TestCooldownBoundsTheChurn(t *testing.T) { + n1, _, _ := balanceRing(t, 4, 0) + b := NewBalancer(n1, 600) // 10 minutes + b.SetMode(BalanceOn) + + 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() { - t.Fatal("precondition: a migration should be in flight") + t.Fatal("precondition: a release should be in flight") } - rep := b.Report(1001) - if rep.InFlight == nil { - t.Fatal("report did not surface the in-flight migration") + if msg := b.Step(1000 + releaseDeadlineSec + 1); msg != "" { + t.Errorf("a timed-out release should clear silently, got %q", msg) } - if rep.InFlight.Local == "" || rep.InFlight.To == "" { - t.Fatalf("in-flight plan is incomplete: %+v", rep.InFlight) + if b.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) } } diff --git a/internal/cluster/ring.go b/internal/cluster/ring.go index 5227bec..79cfa8a 100644 --- a/internal/cluster/ring.go +++ b/internal/cluster/ring.go @@ -68,13 +68,19 @@ type Task struct { // duplicate-claim guard would discard it. The owner applies this one // unconditionally: re-mark the entry enabled and (re)spawn the worker. Restart bool `json:"restart,omitempty"` - // Migrate marks a LOAD-BALANCING task and Target names the node that must take - // the forward over. Unlike a plain creation task — which any lowest-load node - // may claim — a migrate is addressed: it exists precisely because the node that - // is currently lightest (usually the node we just took the forward FROM) would - // otherwise claim it straight back, undoing the migration. - Migrate bool `json:"migrate,omitempty"` - Target string `json:"target,omitempty"` + // Release marks a REBALANCING task: the OWNER of the forward gives it up so + // the ring can place it elsewhere. Only the owning node may act on it — it is + // that node's own load which has to fall, and only it can make that true (its + // topology entry is local bookkeeping no peer may write). + // + // It is deliberately NOT the same thing as a user stop (Revoke). A stop keeps + // 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 // the command reaches it via the token. RemoveNode string `json:"removeNode,omitempty"` diff --git a/internal/cluster/ring_balance.go b/internal/cluster/ring_balance.go index 844ca5e..6e42623 100644 --- a/internal/cluster/ring_balance.go +++ b/internal/cluster/ring_balance.go @@ -3,321 +3,316 @@ package cluster import ( "context" "log" + "sort" "strconv" "time" "webui4frpc/internal/store" ) -// fwdKey identifies a forward for balancing purposes. -type fwdKey struct { - local, remote string - port int +// The rebalance rule, deliberately simple and self-limiting: +// +// release one forward when Forwards >= minForwardsToShare AND Forwards > share +// +// where share = ceil(totalForwards / aliveNodes) — the even split. +// +// Every node runs this check about its OWN forwards. That is the whole point: +// a node's load only actually falls when its own topology entry goes away, so +// only the owner can produce the state change that makes the ring hand the +// forward to somebody else. A coordinator on another machine can see the +// imbalance but cannot make it true; it would have to reach in and mutate the +// owner's bookkeeping, which is exactly the sort of cross-node write this ring +// avoids everywhere else (revoke and restart are owner-directed for the same +// reason). +// +// The rule is self-limiting in both directions, so it needs no upper bound on +// how many forwards a node may shed: +// +// - It stops at the average. 3 forwards over 3 nodes: the owner holds 2 > 1, +// releases one, now holds 1 which is not > 1, and stops. Without the "> share" +// bound an owner would shed EVERY forward it has. +// - It cannot empty a node. A node holding fewer than minForwardsToShare never +// releases, so there is always something left to serve. +// +// A release costs a live connection: stopping the worker drops whatever was +// flowing through the forward, and the new owner re-dials the local service. +// Because the rule above already requires a real imbalance, pacing is handled by +// a per-node cooldown (cooldownSec) rather than a "how long has this lasted" +// window — there is no transient to smooth over, only a churn rate to bound. +const ( + // minForwardsToShare: never release below this many owned forwards. + minForwardsToShare = 2 + // balanceCooldownSec: minimum gap between two releases by one node. One + // release per 10 minutes per node is a trickle; it is what stops a cluster + // whose loads oscillate around the split point from restarting workers in a + // loop. + balanceCooldownSec = 600 + // releaseDeadlineSec: how long to wait for a published release to be + // executed by its owner (us) before giving up on confirming it. A token + // round is ~2s, so this is generous; it exists so a release lost to a + // leader change or a joined ring cannot wedge the node forever. + releaseDeadlineSec = 120 +) + +// Balance modes. +const ( + BalanceOff = "off" + BalanceDryRun = "dry-run" + BalanceOn = "on" +) + +// pool is the replicated view the rule works from. +type pool struct { + alive int // nodes that can take a forward + mine int // forwards THIS node owns (live count, not the replicated one) + total int // forwards across the ring + share int // ceil(total/alive) — the even split } -// planBalance decides whether one forward should move off an overloaded owner. -// -// The rule is deliberately conservative: -// - never touch a stopped forward (disabled) — moving it would resurrect work -// the user explicitly turned off; -// - never touch a forward whose owner is not actually above the threshold; -// - never move more than one per pass, so a transient load spike cannot cause a -// burst of worker restarts (each migration drops live connections). -// -// It returns the chosen forward and the node that should take it. -// now / sustainedSec are accepted so the caller-side hysteresis (which lives in -// BalanceTracker) can pass its bookkeeping through without changing this -// signature, but this function itself is stateless about time: it answers -// "given the current snapshot, what (if anything) should move right now". -// minForwards is an absolute floor: never migrate unless the overloaded node -// holds at least this many forwards. -// -// The ratio test alone is misleading on small rings. With 3 nodes and a single -// forward, the owner scores ~105 (Forwards*100) and the two idle nodes ~5, so the -// owner sits at 20x the mean and trips any ratio — yet one forward on one node is -// the perfectly correct distribution, there is nothing to gain by moving it, and -// the move would only drop a live connection. The ratio detects *shape*; this -// floor says there must be something worth moving. -const minForwards = 2 +// releaseState is one release this node has published and is waiting to see +// leave its own topology. +type releaseState struct { + local, remote string + port int + startedAt int64 +} -func planBalance(s *State, minOverRatio float64) (*TopoEntry, string, bool) { - // Candidate destinations: alive, non-leader-free (any), scored by Load. - type cand struct { - id string - score float64 - forwards int +// Balancer lets each node shed its own share of the cluster's forwards. +// +// It publishes only ONE kind of thing: a release of a forward it owns. There is +// no targeted hand-off and no coordinator, because "drop it and let the ring's +// normal lowest-load claim pick it up" is already the ring's only mechanism for +// moving a forward — and it is the mechanism that, unlike a coordinated move, +// cannot be undone by the node that gave the forward up (the claim path runs +// after the node's load has already fallen). +type Balancer struct { + engine *Engine + mode string + + // cooldownSec is the minimum gap between two releases by this node. + cooldownSec int64 + // releasing is the forward this node gave up and is waiting to see leave. + // While it is set, no second release is started: one at a time keeps the + // connection churn bounded and the log readable. + releasing *releaseState + // lastReleaseAt is when the previous release was published (0 = never). + lastReleaseAt int64 + // reported remembers the last dry-run candidate so the observe-only mode + // logs a change of mind rather than the same line every tick. + reported string +} + +// NewBalancer builds a balancer. It starts in observe-only mode: releasing a +// forward drops a live connection, so acting on the ring is opt-in. +func NewBalancer(e *Engine, cooldownSec int64) *Balancer { + b := &Balancer{engine: e, mode: BalanceDryRun, cooldownSec: cooldownSec} + if b.cooldownSec <= 0 { + b.cooldownSec = balanceCooldownSec } - var dests []cand - total := 0 - for i := range s.Nodes { - n := &s.Nodes[i] + return b +} + +// SetMode selects off / dry-run / on. An unrecognised value is treated as +// dry-run — a typo must not be able to turn releasing on. +func (b *Balancer) SetMode(m string) { + switch m { + case BalanceOff, BalanceDryRun, BalanceOn: + b.mode = m + default: + if m != "" { + log.Printf("ring[%s] unknown balance mode %q; using %s", b.engine.ID, m, BalanceDryRun) + } + b.mode = BalanceDryRun + } +} + +// Mode returns the active mode. +func (b *Balancer) Mode() string { return b.mode } + +// SetCooldown overrides the per-node release cooldown. +func (b *Balancer) SetCooldown(sec int64) { + if sec > 0 { + b.cooldownSec = sec + } +} + +// InFlight reports whether a release is awaiting confirmation. +func (b *Balancer) InFlight() bool { return b.releasing != nil } + +// pool computes the ring-wide view this node derives for itself. +// +// Our own count is taken LIVE from the topology rather than from the replicated +// Nodes[].Load, because Nodes[].Load is only refreshed when a token arrives and +// the decision must see the effect of the release immediately — otherwise a node +// would keep believing it is over its share and shed another forward each round. +func (b *Balancer) pool() pool { + st := &b.engine.state + p := pool{mine: len(st.ForwardsOwnedBy(b.engine.ID))} + for i := range st.Nodes { + n := &st.Nodes[i] if !n.Alive { continue } - dests = append(dests, cand{n.ID, n.Load.Score(), n.Load.Forwards}) - total += n.Load.Forwards - } - if len(dests) < 2 { - return nil, "", false // nothing to balance between - } - avg := float64(total) / float64(len(dests)) - threshold := avg * minOverRatio - - // Overloaded owners, worst first. The absolute floor is applied here so a node - // holding one forward in a three-node ring is never considered hot. - type over struct { - id string - score float64 - forwards int - } - var overs []over - for _, d := range dests { - if d.score > threshold && d.forwards >= minForwards { - overs = append(overs, over{d.id, d.score, d.forwards}) - } - } - if len(overs) == 0 { - return nil, "", false - } - // pick the worst (stable: compare score then id so ties are deterministic) - worstIdx := 0 - for i := 1; i < len(overs); i++ { - if overs[i].score > overs[worstIdx].score || - (overs[i].score == overs[worstIdx].score && overs[i].id < overs[worstIdx].id) { - worstIdx = i - } - } - worst := overs[worstIdx] - - // Lightest node overall is the destination, provided it is meaningfully - // lighter than the worst owner. - lightestIdx := 0 - for i := 1; i < len(dests); i++ { - if dests[i].score < dests[lightestIdx].score || - (dests[i].score == dests[lightestIdx].score && dests[i].id < dests[lightestIdx].id) { - lightestIdx = i - } - } - dest := dests[lightestIdx] - if dest.id == worst.id { - return nil, "", false // the worst owner IS the lightest — nothing to do - } - // Require a real gap; migrating for a 1% difference just churns connections. - if worst.score-dest.score < 100 { - return nil, "", false // less than one forward's worth of difference - } - - // Choose which forward to move: the cheapest to give up. Skip disabled ones. - var pick *TopoEntry - for _, te := range s.TopologyList() { - if te.OwnerID != worst.id || !te.Active || te.Link.Disabled { + p.alive++ + if n.ID == b.engine.ID { + p.total += p.mine continue } - if pick == nil || te.Link.RemotePort < pick.Link.RemotePort { - pick = te - } + p.total += n.Load.Forwards } + if p.alive > 0 { + p.share = (p.total + p.alive - 1) / p.alive + } + return p +} + +// assess answers "should this node give one forward up, and which one?" together +// with the facts behind the answer, so the read-only report can explain a node +// that is over its share but not acting. It never mutates anything. +// +// blocked is a short reason: "single-node", "below-min", "balanced", +// "nothing-active", "cooldown" or "in-flight". +func (b *Balancer) assess(now int64) (pick *TopoEntry, p pool, blocked string) { + p = b.pool() + if !b.engine.IsMember() { + return nil, p, "single-node" + } + if p.mine < minForwardsToShare { + return nil, p, "below-min" + } + if p.mine <= p.share { + return nil, p, "balanced" + } + pick = releaseCandidate(&b.engine.state, b.engine.ID) if pick == nil { - return nil, "", false // everything it owns is stopped + return nil, p, "nothing-active" } - return pick, dest.id, true -} - -// Balancer moves forwards off an overloaded node, one at a time. -// -// It is a TWO-PHASE state machine, and the split is a safety requirement rather -// than a stylistic choice: -// -// phase 1 (revoking): ask the current owner to stop. Wait until the topology -// entry is actually gone before publishing anything. -// phase 2 (publishing): hand the forward to a NAMED destination. -// -// Publishing both at once does not work, for two independent reasons: -// -// 1. loadSnapshot() derives Forwards from ForwardsOwnedBy, so the moment the old -// entry is dropped the old owner becomes the lightest node. A plain creation -// task is claimed by "whoever is lightest" — so the node we just took the -// forward FROM would immediately claim it back and the migration would be a -// no-op that merely restarts the worker. -// 2. PendingList sorts by task id as a STRING (t10 < t5), so the revoke and the -// submit could execute in either order within one pass. -// -// Migrate tasks solve (1) by carrying an explicit Target; this type solves (2) by -// refusing to publish phase 2 before phase 1 has completed. -type Balancer struct { - engine *Engine - - // ratio is how far above the average a node's Score must be, relative to the - // lightest node, before anything moves. - ratio float64 - // sustainedSec is how long an imbalance must persist before we act. This is - // what stops a transient CPU/network spike from restarting workers. - sustainedSec int64 - - // 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 + // The guards below do not change the verdict ("this node is over its + // share"), only whether now is the moment; the report shows both. + if b.releasing != nil { + return pick, p, "in-flight" } - if sustainedSec < 0 { - sustainedSec = 0 - } - return &Balancer{ - engine: e, ratio: ratio, sustainedSec: sustainedSec, - dryRun: true, // safe default: report, never move - overSince: map[string]int64{}, + if b.lastReleaseAt > 0 && now-b.lastReleaseAt < b.cooldownSec { + return pick, p, "cooldown" } + return pick, p, "" } -// SetDryRun selects the mode. Dry-run (the default) only reports. -func (b *Balancer) SetDryRun(v bool) { b.dryRun = v } - -// DryRun reports whether the balancer is in observe-only mode. -func (b *Balancer) DryRun() bool { return b.dryRun } - -// InFlight reports whether a migration is currently between phases. -func (b *Balancer) InFlight() bool { - return b.revoking != nil || b.publishing != nil +// releaseCandidate picks which of this node's forwards to give up. +// +// Deterministic (lowest remote port, then task id) so the choice is explainable +// and repeatable: a node that re-evaluates the same state picks the same +// forward. Stopped forwards are skipped — their entry is the carrier of a user's +// stop decision, and releasing one would hand a STOPPED forward to another node +// as if it were work to run. +func releaseCandidate(s *State, owner string) *TopoEntry { + var mine []*TopoEntry + for _, te := range s.TopologyList() { + if te.OwnerID != owner || !te.Active || te.Link.Disabled { + continue + } + mine = append(mine, te) + } + if len(mine) == 0 { + return nil + } + sort.Slice(mine, func(i, j int) bool { + if mine[i].Link.RemotePort != mine[j].Link.RemotePort { + return mine[i].Link.RemotePort < mine[j].Link.RemotePort + } + return mine[i].TaskID < mine[j].TaskID + }) + return mine[0] } -// Step advances the state machine by one cycle. now is the current unix second. -// It returns a short human-readable description when something happened. +// landed reports whether the forward this node released has actually left it. +// +// The test is "no longer owned by me", NOT "no entry exists": the ring's normal +// claim picks the forward up within a round or two, so an existence test would +// see the new owner's entry and wait forever for a state that has already +// happened. +func (b *Balancer) landed() bool { + m := b.releasing + if m == nil { + return true + } + owner := b.engine.state.TopologyOwner(&Task{ + Local: store.Local{Name: m.local}, + Remote: store.Remote{Name: m.remote}, + Link: store.Link{RemotePort: m.port}, + }) + return owner != b.engine.ID +} + +// Step advances the balancer by one cycle and returns a human-readable +// description when it did (or, in dry-run, would have done) something. func (b *Balancer) Step(now int64) string { - // --- phase 2: did phase 1 finish? --- - if b.revoking != nil { - if b.engine.state.TopologyOwner(&Task{ - Local: store.Local{Name: b.revoking.local}, - Remote: store.Remote{Name: b.revoking.remote}, - Link: store.Link{RemotePort: b.revoking.port}, - }) != "" { - // Still owned: the revoke has not completed yet. Wait, unless it has - // been far too long, in which case something is wrong with the ring. - if now < b.revoking.deadlineAt { - return "" - } - log.Printf("ring[%s] migrate revoke for %s→%s:%d timed out; abandoning", - b.engine.ID, b.revoking.local, b.revoking.remote, b.revoking.port) - b.revoking = nil + if b.releasing != nil { + if b.landed() { + done := b.releasing + b.releasing = nil + return "gave up " + done.local + ":" + strconv.Itoa(done.port) + " for rebalancing" + } + if now-b.releasing.startedAt > releaseDeadlineSec { + log.Printf("ring[%s] release of %s:%d not confirmed within %ds; clearing (forward is still ours)", + b.engine.ID, b.releasing.local, b.releasing.port, releaseDeadlineSec) + b.releasing = nil + } + return "" + } + + pick, p, blocked := b.assess(now) + if pick == nil { + b.reported = "" + return "" + } + desc := pick.Local.Name + ":" + strconv.Itoa(pick.Link.RemotePort) + + " (" + strconv.Itoa(p.mine) + " owned > " + strconv.Itoa(p.share) + " share of " + + strconv.Itoa(p.total) + " over " + strconv.Itoa(p.alive) + " nodes)" + + switch b.mode { + case BalanceOff: + return "" + case BalanceDryRun: + // Log a change of mind, not the same line every tick. + if b.reported == desc { return "" } - // Entry gone — the old worker is stopped. Publish for the named target. - m := b.revoking - b.revoking = nil - b.publishing = m - b.publishMigration() - return "revoke complete, migrating " + m.local + ":" + strconv.Itoa(m.port) + " → " + m.toOwner + b.reported = desc + if blocked != "" { + return "would release " + desc + " [blocked: " + blocked + "]" + } + return "would release " + desc } - // --- phase 1: wait for an in-flight publish, or start a new migration --- - if b.publishing != nil { - // The migrate task carries itself in the ring; nothing to do but let it - // ride. Clear once the forward is back in the topology. - if b.engine.state.TopologyOwner(&Task{ - Local: store.Local{Name: b.publishing.local}, - Remote: store.Remote{Name: b.publishing.remote}, - Link: store.Link{RemotePort: b.publishing.port}, - }) == b.publishing.toOwner { - done := b.publishing - b.publishing = nil - return "migrated " + done.local + ":" + strconv.Itoa(done.port) + " to " + done.toOwner - } - if now >= b.publishing.deadlineAt { - // Nobody claimed it. Submit an UNTARGETED task so whichever node is - // lightest picks it up — the ring keeps the forward alive rather than - // stranding it, which is the same guarantee the rest of the protocol - // relies on. - m := b.publishing - b.publishing = nil - b.engine.SubmitTask(store.Local{Name: m.local}, store.Remote{Name: m.remote}, - store.Link{RemotePort: m.port}) - return "migrate target did not claim " + m.local + ":" + strconv.Itoa(m.port) + "; re-opened for any node" - } + // Mode is ON. Blocked reasons still stop us. + if blocked != "" { return "" } - - // --- nothing in flight: look for work --- - pick, dest, hot := b.engine.findBalanceCandidate(b.ratio) - if hot != "" { - if _, seen := b.overSince[hot]; !seen { - b.overSince[hot] = now - } - } else { - b.overSince = map[string]int64{} - } - if pick == nil || dest == "" { + t := b.engine.ReleaseTask(pick.Local, pick.Remote, pick.Link) + if t == nil { return "" } - if since, ok := b.overSince[pick.OwnerID]; !ok || now-since < int64(b.sustainedSec) { - return "" + b.releasing = &releaseState{ + 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. - m := &pendingMigration{ - local: pick.Local.Name, remote: pick.Remote.Name, port: pick.Link.RemotePort, - fromOwner: pick.OwnerID, toOwner: dest, - // Generous deadline: a token round is ~2s, and the owner may be busy. - deadlineAt: now + 120, - } - b.revoking = m - b.engine.RevokeTask(pick.Local, pick.Remote, pick.Link) - // The revoke task must be able to reach the owner even though the balancer - // runs on the leader: it travels in the ring like any other command. - return "revoking " + m.local + ":" + strconv.Itoa(m.port) + " from " + m.fromOwner + b.lastReleaseAt = now + b.reported = "" + return "releasing " + desc } -func (b *Balancer) publishMigration() { - m := b.publishing - if m == nil { - return - } - b.engine.submitMigrate(store.Local{Name: m.local}, store.Remote{Name: m.remote}, - store.Link{RemotePort: m.port}, m.toOwner) -} - -// findBalanceCandidate exposes planBalance to the engine's live state. -func (e *Engine) findBalanceCandidate(ratio float64) (*TopoEntry, string, string) { - pick, dest, ok := planBalance(&e.state, ratio) - if !ok { - return nil, "", "" - } - // Report the overloaded node even when there is nothing movable, so the - // hysteresis timer can start for a node whose forwards are all stopped. - return pick, dest, pick.OwnerID -} - -// WatchBalance drives the balancer on a ticker until ctx is cancelled. +// Watch drives Step on a ticker until ctx is cancelled. // -// Only the LEADER runs this. Balancing is a cluster-wide decision (it moves -// forwards between members), and letting every member decide independently would -// mean several migrations racing for the same forward — the ring is a shared -// token, so two nodes must not both publish a revoke for the same entry. +// It runs on EVERY node, not just the leader. Balancing here is not a cluster +// decision being delegated to one node — it is each node applying the same +// replicated rule to its own forwards. Two nodes never contend for the same +// forward because a node can only release what it owns. // -// dryRun makes the balancer report what it WOULD do without publishing anything. -// That is the mode to use first on a live cluster: it answers "is my real data -// going to trip the threshold, and what would it move?" before anything on the -// ring is disturbed. -func (b *Balancer) WatchBalance(ctx context.Context, interval time.Duration, dryRun bool) { +// The first check happens after one interval rather than immediately: a node +// that has just booted adopts the ring's topology and its own load in the first +// token round, and deciding before that would act on a half-known pool. +func (b *Balancer) Watch(ctx context.Context, interval time.Duration) { tick := time.NewTicker(interval) defer tick.Stop() for { @@ -326,69 +321,46 @@ func (b *Balancer) WatchBalance(ctx context.Context, interval time.Duration, dry return case <-tick.C: } - if !b.engine.IsLeader() { - continue - } - now := time.Now().Unix() - if dryRun { - if msg := b.stepDryRun(now); msg != "" { - log.Printf("ring[%s] balance (dry-run): %s", b.engine.ID, msg) - } - continue - } - if msg := b.Step(now); msg != "" { + if msg := b.Step(time.Now().Unix()); msg != "" { log.Printf("ring[%s] balance: %s", b.engine.ID, msg) } } } -// stepDryRun reports the migration the balancer would start right now, without -// publishing a revoke. Hysteresis is still applied so the report reflects the -// same sustained-imbalance rule the real path uses. -func (b *Balancer) stepDryRun(now int64) string { - if b.InFlight() { - return "" - } - pick, dest, hot := b.engine.findBalanceCandidate(b.ratio) - if hot == "" { - b.overSince = map[string]int64{} - return "" - } - if _, seen := b.overSince[hot]; !seen { - b.overSince[hot] = now - return "" - } - if pick == nil || dest == "" { - return "" - } - if since, ok := b.overSince[pick.OwnerID]; !ok || now-since < int64(b.sustainedSec) { - return "" - } - return "would migrate " + pick.Local.Name + ":" + strconv.Itoa(pick.Link.RemotePort) + - " from " + pick.OwnerID + " to " + dest -} +// ---- read-only report for the API/UI ---- -// 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 -// log, and only when an imbalance exists AND the hysteresis window has elapsed. -// An operator asking "is this cluster balanced?" gets silence in exactly the -// healthy case, which is indistinguishable from "balancing is switched off". +// It is per-node by nature: every member has its own answer, and the answer that +// matters is the one on the node doing (or refusing to do) the releasing. The +// peer list is included because "am I over my share?" is only meaningful next to +// what everyone else holds. type BalanceReport struct { - Enabled bool `json:"enabled"` - DryRun bool `json:"dryRun"` - Ratio float64 `json:"ratio"` - SustainSec int64 `json:"sustainSec"` - // Only the leader acts; on a follower these stay false/empty. - IsLeader bool `json:"isLeader"` - // Current per-node scores. + Mode string `json:"mode"` // off | dry-run | on + Enabled bool `json:"enabled"` // mode != off + DryRun bool `json:"dryRun"` // mode != on + IsMember bool `json:"isMember"` + Self string `json:"self"` + + SelfForwards int `json:"selfForwards"` + Share int `json:"share"` + TotalForwards int `json:"totalForwards"` + AliveNodes int `json:"aliveNodes"` + SelfOverShare bool `json:"selfOverShare"` + + // WouldRelease is what this node would give up (present whenever it is over + // its share, including while blocked — see BlockedBy). + WouldRelease *BalanceReleasePlan `json:"wouldRelease,omitempty"` + // Releasing is the release this node has published and is awaiting. + Releasing *BalanceReleasePlan `json:"releasing,omitempty"` + // BlockedBy explains why an over-share node is not acting yet. + BlockedBy string `json:"blockedBy,omitempty"` + + CooldownSec int64 `json:"cooldownSec"` + LastReleaseAt int64 `json:"lastReleaseAt,omitempty"` + + // Peers: every member's replicated load, for context. Nodes []BalanceNodeReport `json:"nodes"` - // 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 { @@ -401,24 +373,47 @@ type BalanceNodeReport struct { Score float64 `json:"score"` } -type BalanceMigrationPlan struct { - Local string `json:"local"` - Remote string `json:"remote"` - Port int `json:"remotePort"` - From string `json:"from"` - To string `json:"to"` +type BalanceReleasePlan struct { + Local string `json:"local"` + Remote string `json:"remote"` + RemotePort int `json:"remotePort"` + Owner string `json:"owner"` } -// Report renders the balancer's current view. It never mutates anything — -// note that it does NOT touch overSince, so calling it cannot start or reset -// the hysteresis timer that the real Step path uses. -func (b *Balancer) Report(now int64) BalanceReport { +// Report renders the current view. It never mutates anything — in particular it +// does NOT touch lastReleaseAt or the release state, so an operator polling this +// cannot start, delay or cancel a release. (A read-only endpoint that perturbs +// the behaviour it reports is worse than no endpoint.) +func (b *Balancer) Report() BalanceReport { + now := time.Now().Unix() rep := BalanceReport{ - Enabled: !b.dryRun, - DryRun: b.dryRun, - Ratio: b.ratio, - SustainSec: int64(b.sustainedSec), - IsLeader: b.engine.IsLeader(), + Mode: b.mode, + Enabled: b.mode != BalanceOff, + DryRun: b.mode != BalanceOn, + IsMember: b.engine.IsMember(), + Self: b.engine.ID, + CooldownSec: b.cooldownSec, + } + p := b.pool() + rep.SelfForwards, rep.Share = p.mine, p.share + rep.TotalForwards, rep.AliveNodes = p.total, p.alive + rep.SelfOverShare = p.mine > p.share + + if pick, _, blocked := b.assess(now); pick != nil { + rep.WouldRelease = &BalanceReleasePlan{ + Local: pick.Local.Name, Remote: pick.Remote.Name, + RemotePort: pick.Link.RemotePort, Owner: pick.OwnerID, + } + rep.BlockedBy = blocked + } + if m := b.releasing; m != nil { + rep.Releasing = &BalanceReleasePlan{ + Local: m.local, Remote: m.remote, RemotePort: m.port, Owner: b.engine.ID, + } + rep.BlockedBy = "in-flight" + } + if b.lastReleaseAt > 0 { + rep.LastReleaseAt = b.lastReleaseAt } for i := range b.engine.state.Nodes { n := &b.engine.state.Nodes[i] @@ -428,46 +423,5 @@ func (b *Balancer) Report(now int64) BalanceReport { 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 } - -// 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 -} diff --git a/internal/cluster/ring_engine.go b/internal/cluster/ring_engine.go index 923c7b5..43542d2 100644 --- a/internal/cluster/ring_engine.go +++ b/internal/cluster/ring_engine.go @@ -24,6 +24,12 @@ type Handler interface { // creation path (idempotency guard) and the duplicate-claim guard refuse to // re-create an existing forward. 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 } @@ -406,16 +412,24 @@ func (e *Engine) runCommands(ctx context.Context, tk *Token) error { } continue // owned elsewhere — ride to the owner } - if t.Migrate { - // Addressed to one node. The whole point is that "whoever is - // lightest" is the WRONG answer here: the node we just revoked the - // forward from is the one that became lightest, and it would claim - // the task straight back. - if t.Target == e.ID { + if t.Release { + // Owner-directed, like a revoke: only the node holding the forward + // may give it up. A release executed anywhere else would be an + // attempt to move somebody else's load, which cannot work — the + // owner's own count is what has to fall, and its topology entry is + // 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 break } - continue // ride to the target + continue // owned elsewhere — ride to the owner } if t.Restart { // 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 // (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. - 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 != "" { 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, @@ -557,33 +575,46 @@ func (e *Engine) runCommands(ctx context.Context, tk *Token) error { continue } } - if claimed.Migrate { - // We are the designated destination. The forward must NOT already be - // here: the two-phase migration guarantees the old entry was gone before - // this task was published. If one somehow exists, drop the task rather - // than creating a second topology entry for the same forward. - if owner := e.state.TopologyOwner(claimed); owner != "" { - log.Printf("ring[%s] drop migrate %s: already owned by %s", e.ID, claimed.ID, owner) + if claimed.Release { + // Only reached when this node owns the forward (or it is already gone + // and we are the fallback consumer — see the selection loop). + owner := e.state.TopologyOwner(claimed) + if owner != e.ID { + // The forward vanished before we got here; nothing to give up. 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 err := e.Handler.Claim(ctx, claimed); err != nil { - // Put it back so another cycle can retry — the forward is - // 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 + if err := e.Handler.Release(ctx, claimed); err != nil { + log.Printf("ring[%s] release %s: %v", e.ID, claimed.ID, err) } } - 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 { - _, _ = e.Log.Append(e.ID, LogForwardAdd, map[string]any{ - "taskId": claimed.ID, "local": claimed.Local.Name, "remote": claimed.Remote.Name, - "migratedFrom": e.state.TopologyOwner(claimed), + _, _ = e.Log.Append(e.ID, LogForwardStop, map[string]any{ + "taskId": claimed.ID, "local": claimed.Local.Name, + "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 } if claimed.Restart { @@ -1000,27 +1031,23 @@ func (e *Engine) SubmitRestart(local store.Local, remote store.Remote, link stor 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 -// is currently lightest", and after a revoke the node we just took the forward -// from IS the lightest — so an untargeted task would be claimed straight back by -// the old owner and the migration would achieve nothing but a worker restart. -func (e *Engine) submitMigrate(local store.Local, remote store.Remote, link store.Link, target string) *Task { +// It goes through AddPending rather than writing PendingTasks directly. That is +// not cosmetic: OnToken rebuilds local-only pending commands from the ids it +// never published, and only the Add* helpers record them — a task injected by +// hand was silently dropped at the next state adoption, so it never reached the +// ring at all. +func (e *Engine) ReleaseTask(local store.Local, remote store.Remote, link store.Link) *Task { link.Disabled = false - t := &Task{ - ID: e.state.NextTaskID(), - Local: local, - Remote: remote, - Link: link, - Created: time.Now().Unix(), - Migrate: true, - Target: target, + t := e.state.AddPending(local, remote, link) + if t == nil { + return nil } - if e.state.PendingTasks == nil { - e.state.PendingTasks = map[string]*Task{} - } - e.state.PendingTasks[t.ID] = t + t.Release = true return t } diff --git a/internal/cluster/ring_engine_test.go b/internal/cluster/ring_engine_test.go index 89675ac..40d223a 100644 --- a/internal/cluster/ring_engine_test.go +++ b/internal/cluster/ring_engine_test.go @@ -12,6 +12,14 @@ type fakeHandler struct { claim func(ctx context.Context, tk *Task) error revoke 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 { diff --git a/internal/cluster/ring_handler.go b/internal/cluster/ring_handler.go index 83c92b4..82712cf 100644 --- a/internal/cluster/ring_handler.go +++ b/internal/cluster/ring_handler.go @@ -21,6 +21,12 @@ type AppHandler struct { // stopped flag + bring the worker back). Required because a stop keeps the // topology entry, so the normal claim path will not re-create it. 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. @@ -48,6 +54,22 @@ func (a *AppHandler) Revoke(ctx context.Context, tk *Task) error { 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. func (a *AppHandler) Restart(ctx context.Context, tk *Task) error { if a.RestartFn == nil { diff --git a/internal/httpapi/server.go b/internal/httpapi/server.go index ffe6a72..a24db20 100644 --- a/internal/httpapi/server.go +++ b/internal/httpapi/server.go @@ -655,7 +655,7 @@ func (h *Handler) handleClusterBalance(w http.ResponseWriter, r *http.Request) { http.Error(w, "load balancer not enabled", http.StatusNotFound) 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