webui4frpc/internal/cluster/balance_test.go

package cluster

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

	"webui4frpc/internal/store"
)

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

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

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

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

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

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

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