webui4frpc/internal/cluster/balance_test.go

package cluster

import (
	"context"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"webui4frpc/internal/store"
)

// mustState round-trips the seed through JSON so the engine under test owns its
// own copy of Nodes/Topology. AdoptState assigns the state by value, but Nodes is
// a slice: handing one seed to several engines would alias one backing array and
// make a mutation on one node visible on the others — a test artefact that
// previously produced a false green here.
func mustState(t *testing.T, s State) *State {
	t.Helper()
	b, err := json.Marshal(s)
	if err != nil {
		t.Fatal(err)
	}
	var out State
	if err := json.Unmarshal(b, &out); err != nil {
		t.Fatal(err)
	}
	e := newTestEngine("probe", true)
	e.AdoptState(out)
	return e.State()
}

// TestPlanBalanceMovesOneForwardOffTheOverloadedNode covers the happy path and
// the "one per pass" rule.
func TestPlanBalanceMovesOneForwardOffTheOverloadedNode(t *testing.T) {
	s := State{
		LeaderID: "n1",
		Nodes: []Node{
			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 3, MemPct: 5}},
			{ID: "n2", Alive: true, Load: Load{Forwards: 0, MemPct: 5}},
			{ID: "n3", Alive: true, Load: Load{Forwards: 0, MemPct: 5}},
		},
		Topology: map[string]*TopoEntry{},
	}
	for i, p := range []int{100, 101, 102} {
		s.Topology[string(rune('a'+i))] = &TopoEntry{
			TaskID: string(rune('a' + i)), OwnerID: "n1", Active: true,
			Local:  store.Local{Name: "web"},
			Remote: store.Remote{Name: "frps1"},
			Link:   store.Link{RemotePort: p},
		}
	}
	st := mustState(t, s)

	pick, dest, ok := planBalance(st, 2.0)
	if !ok {
		t.Fatal("expected a migration plan for a 3-vs-0 imbalance")
	}
	if dest != "n2" && dest != "n3" {
		t.Fatalf("destination must be an idle node, got %q", dest)
	}
	if pick.OwnerID != "n1" {
		t.Fatalf("must move a forward off the overloaded owner, got owner %q", pick.OwnerID)
	}
	// Exactly one forward, and the cheapest (lowest port) is a fine choice.
	if pick.Link.RemotePort != 100 {
		t.Logf("chose port %d (any single forward is acceptable)", pick.Link.RemotePort)
	}
}

// TestPlanBalanceNeverMovesDisabledForwards: moving a stopped forward would
// resurrect work the user turned off.
func TestPlanBalanceNeverMovesDisabledForwards(t *testing.T) {
	s := State{
		LeaderID: "n1",
		Nodes: []Node{
			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 2, MemPct: 5}},
			{ID: "n2", Alive: true, Load: Load{Forwards: 0, MemPct: 5}},
		},
		Topology: map[string]*TopoEntry{},
	}
	// n1's only forwards are stopped.
	for i, p := range []int{200, 201} {
		s.Topology[string(rune('a'+i))] = &TopoEntry{
			TaskID: string(rune('a' + i)), OwnerID: "n1", Active: false,
			Local:  store.Local{Name: "web"},
			Remote: store.Remote{Name: "frps1"},
			Link:   store.Link{RemotePort: p, Disabled: true},
		}
	}
	st := mustState(t, s)
	if _, _, ok := planBalance(st, 2.0); ok {
		t.Fatal("must not plan a migration when every forward on the hot node is stopped")
	}
}

// TestPlanBalanceIsQuietWhenBalanced: the whole point of thresholding is that a
// healthy ring is left alone.
func TestPlanBalanceIsQuietWhenBalanced(t *testing.T) {
	s := State{
		LeaderID: "n1",
		Nodes: []Node{
			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 1, MemPct: 5}},
			{ID: "n2", Alive: true, Load: Load{Forwards: 1, MemPct: 5}},
			{ID: "n3", Alive: true, Load: Load{Forwards: 1, MemPct: 5}},
		},
		Topology: map[string]*TopoEntry{},
	}
	for i, id := range []string{"n1", "n2", "n3"} {
		s.Topology[id] = &TopoEntry{
			TaskID: id, OwnerID: id, Active: true,
			Local:  store.Local{Name: "web"},
			Remote: store.Remote{Name: "frps1"},
			Link:   store.Link{RemotePort: 300 + i},
		}
	}
	st := mustState(t, s)
	if _, _, ok := planBalance(st, 2.0); ok {
		t.Fatal("a perfectly balanced ring must not trigger a migration")
	}
}

// TestPlanBalanceIgnoresDeadNodes: a node marked offline must never be chosen as
// a destination, or we would migrate onto a machine that is not there.
func TestPlanBalanceIgnoresDeadNodes(t *testing.T) {
	s := State{
		LeaderID: "n1",
		Nodes: []Node{
			{ID: "n1", Alive: true, IsLeader: true, Load: Load{Forwards: 3, MemPct: 5}},
			{ID: "n2", Alive: false, Load: Load{Forwards: 0, MemPct: 0}},
			{ID: "n3", Alive: true, Load: Load{Forwards: 1, MemPct: 5}},
		},
		Topology: map[string]*TopoEntry{},
	}
	s.Topology["a"] = &TopoEntry{
		TaskID: "a", OwnerID: "n1", Active: true,
		Local:  store.Local{Name: "web"},
		Remote: store.Remote{Name: "frps1"},
		Link:   store.Link{RemotePort: 100},
	}
	st := mustState(t, s)
	_, dest, ok := planBalance(st, 2.0)
	if !ok {
		t.Skip("this arrangement is within threshold; the dead-node case is covered by the alive filter above")
	}
	if dest == "n2" {
		t.Fatal("chose an OFFLINE node as the migration destination")
	}
}

// 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()
	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, 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: ownCount, MemPct: 5}},
			{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 5}},
			{ID: "n3", Addr: "n3:7500", Alive: true, Load: Load{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)
	}
	// n1 claims `ownCount` forwards.
	for i := 0; i < ownCount; i++ {
		n1.state.AddTopology(&Task{
			ID:    "f" + string(rune('a'+i)),
			Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
			Link: store.Link{RemotePort: 100 + i},
		}, "n1")
	}
	var snap State
	b, _ := json.Marshal(n1.state)
	_ = json.Unmarshal(b, &snap)
	for _, e := range []*Engine{n1, n2, n3} {
		e.AdoptState(snap)
	}
	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

	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
			}}
	}
	n1.Handler = mk("n1", 50)
	n2.Handler = mk("n2", 40) // the TARGET — heavier than n3
	n3.Handler = mk("n3", 1)  // strictly the lightest
	var snap 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)
			}
		}
	}
	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 }}

	// 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)
	}
	if claimed != 0 {
		t.Errorf("Claim ran %d time(s) while the forward was still owned; "+
			"this would create a duplicate proxy", claimed)
	}
	if n := len(n1.State().TopologyList()); n != 3 {
		t.Errorf("topology has %d entries, want 3 (no duplicate created)", n)
	}
}

// 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)

	// 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")
	}
	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")
	}
}

// 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)
	for i := 0; i < 5; i++ {
		if msg := b.Step(int64(1000 + i*600)); msg != "" {
			t.Fatalf("a balanced ring produced action: %q", msg)
		}
	}
	if b.InFlight() {
		t.Error("balancer started work on a balanced ring")
	}
}

// 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)
		}
	}
	if b.InFlight() {
		t.Error("dry-run started an in-flight migration")
	}
	if n := len(n1.State().TopologyList()); n != 3 {
		t.Errorf("topology changed under dry-run: %d entries, want 3", n)
	}
}

// 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")
	}

	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()

	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))
		}
	}
	if got := b.overSince["n1"]; got != started {
		t.Errorf("Report reset the overload timer: %d -> %d", started, got)
	}

	// 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")
	}
}

// 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 {
		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")
		}
	}
}

// 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
	if !b.InFlight() {
		t.Fatal("precondition: a migration should be in flight")
	}
	rep := b.Report(1001)
	if rep.InFlight == nil {
		t.Fatal("report did not surface the in-flight migration")
	}
	if rep.InFlight.Local == "" || rep.InFlight.To == "" {
		t.Fatalf("in-flight plan is incomplete: %+v", rep.InFlight)
	}
}