webui4frpc/internal/cluster/balance_test.go

package cluster

import (
	"encoding/json"
	"testing"

	"webui4frpc/internal/store"
)

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

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

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

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

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

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