webui4frpc/internal/cluster/ring_engine_test.go

package cluster

import (
	"context"
	"testing"

	"webui4frpc/internal/store"
)

type fakeHandler struct {
	load    Load
	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 {
	if h.revoke != nil {
		return h.revoke(ctx, tk)
	}
	return nil
}

func (h *fakeHandler) Restart(ctx context.Context, tk *Task) error {
	if h.restart != nil {
		return h.restart(ctx, tk)
	}
	return nil
}

func (h *fakeHandler) RuntimeLoad() Load { return h.load }
func (h *fakeHandler) Claim(ctx context.Context, tk *Task) error {
	if h.claim != nil {
		return h.claim(ctx, tk)
	}
	return nil
}

// sendNull is a no-op sender (single-node ring).
func sendNull(ctx context.Context, next string, tk *Token) error { return nil }

// TestSingleNodeCycle: leader starts ring, processes phases locally, and the
// task gets claimed (lowest load = only node).
func TestSingleNodeCycle(t *testing.T) {
	eng := NewEngine("n1", "n1:7500", "u", "p", "0.71.0", []string{"0.71.0"},
		&fakeHandler{load: Load{MemPct: 30, NetPct: 30}}, sendNull, "n1:7500", true, "")
	eng.StartRing(context.Background())

	// single node: token stays local; no successor so nothing travels.
	if eng.State().Nodes[0].ID != "n1" {
		t.Fatalf("nodes = %+v", eng.State().Nodes)
	}
}

// TestTwoNodesPhaseCollect: leader n1 sends to n2; n1 collects both after n2
// stamps; then phase flips and second round syncs.
func TestTwoNodesSingleRound(t *testing.T) {
	eng2 := NewEngine("n2", "n2:7500", "u", "p", "0.71.0", nil,
		&fakeHandler{load: Load{MemPct: 10, NetPct: 10}}, sendNull, "n2:7500", false, "")
	eng2.state.UpsertNode(Node{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true, Load: Load{MemPct: 50, NetPct: 50}})

	// n1 sends a token carrying the cluster picture; single-round processing:
	// the receiving node adopts it and appends its own info in one pass.
	tk := &Token{Cycle: 1, State: State{
		Nodes: []Node{
			{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true},
			{ID: "n2", Addr: "n2:7500", Alive: true},
		},
	}}

	out, err := eng2.OnToken(context.Background(), tk)
	if err != nil {
		t.Fatal(err)
	}
	if eng2.State().Find("n2") < 0 {
		t.Fatal("n2 should be present in adopted state after single-round")
	}
	if out == nil {
		t.Fatal("nil token after processing")
	}
}

// TestLeaderRoundTripNoResurrection: when the leader submits a task and the
// token completes a full round (leader→n2 claims→n3→leader), the consumed task
// MUST NOT be resurrected on the leader's next OnToken, and the topology must
// still attribute the forward to the single claimer (n2).
//
// This holds NOT via a published-set mark on StartRing, but via Go map
// reference semantics: StartRing sets tk.State = e.state, so the token's
// PendingTasks map IS the leader's own map. When n2 calls ClaimPending it
// deletes from that shared map — the deletion is visible to the leader too.
// By the time the token returns, the claimed task is already gone from the
// leader's e.state.PendingTasks, so the localPending re-merge has nothing to
// re-inject. Single token + remove-on-claim = single claim (plan §M6).
func TestLeaderRoundTripNoResurrection(t *testing.T) {
	// 3-node ring: n1 (leader+submitter), n2 (lowest, claims), n3 (idle).
	var sent *Token
	sendCap := func(ctx context.Context, next string, tk *Token) error {
		sent = tk
		return nil
	}
	n1 := NewEngine("n1", "n1:7500", "u", "p", "v", nil,
		&fakeHandler{load: Load{MemPct: 50, NetPct: 50}}, sendCap, "n1:7500", true, "")
	// n2/n3 carry a lower stored load than n1's NewEngine default ({10,10})
	// so LowestAlive picks n2 (first among the tied low nodes) as the claimer.
	n1.state.UpsertNode(Node{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 1, NetPct: 1}})
	n1.state.UpsertNode(Node{ID: "n3", Addr: "n3:7500", Alive: true, Load: Load{MemPct: 1, NetPct: 1}})

	task := n1.SubmitTask(store.Local{Name: "l1"}, store.Remote{Name: "r1"}, store.Link{RemotePort: 100})
	if task == nil {
		t.Fatal("submit returned nil")
	}
	n1.StartRing(context.Background())
	if sent == nil {
		t.Fatal("StartRing did not send a token")
	}

	// n2 receives, claims t1 (lowest load), establishes topology.
	n2 := NewEngine("n2", "n2:7500", "u", "p", "v", nil,
		&fakeHandler{load: Load{MemPct: 10, NetPct: 10}}, sendCap, "n2:7500", false, "")
	out2, err := n2.OnToken(context.Background(), sent)
	if err != nil {
		t.Fatal(err)
	}
	if got := len(n2.state.TopologyList()); got != 1 {
		t.Fatalf("n2 should own 1 forward, topo=%+v", n2.state.TopologyList())
	}
	if len(n2.state.PendingList()) != 0 {
		t.Fatalf("pending should be empty after n2 claim, got %+v", n2.state.PendingList())
	}

	// n3 receives, nothing to claim, forwards.
	n3 := NewEngine("n3", "n3:7500", "u", "p", "v", nil,
		&fakeHandler{load: Load{MemPct: 10, NetPct: 10}}, sendCap, "n3:7500", false, "")
	out3, err := n3.OnToken(context.Background(), out2)
	if err != nil {
		t.Fatal(err)
	}

	// Token returns to n1. The consumed task t1 MUST NOT be resurrected.
	if _, err := n1.OnToken(context.Background(), out3); err != nil {
		t.Fatal(err)
	}
	if got := len(n1.state.PendingList()); got != 0 {
		t.Fatalf("n1 resurrected a consumed task: pending=%+v (single-token + shared-map must prevent this)",
			n1.state.PendingList())
	}
	// Topology still attributes t1 to n2 (not overwritten by a re-claim).
	topo := n1.state.TopologyList()
	if len(topo) != 1 || topo[0].OwnerID != "n2" {
		t.Fatalf("topology should be t1@n2, got %+v", topo)
	}
}

// TestClaimGuardDropsDuplicateForward: when a node (lowest load) receives a
// pending task whose forward is ALREADY in the topology owned by another node
// (a resurrected/stale copy), the claim path must drop it WITHOUT spawning a
// worker or overwriting the topology. Without the guard a second worker for
// the same forward would be spawned and orphaned (the topology entry is keyed
// by task id, so AddTopology would silently overwrite the owner).
func TestClaimGuardDropsDuplicateForward(t *testing.T) {
	// n1 is lowest load and would otherwise claim; the fake Claim hook fails
	// the test if ever called.
	n1 := NewEngine("n1", "n1:7500", "u", "p", "v", nil,
		&fakeHandler{
			load: Load{MemPct: 10, NetPct: 10},
			claim: func(ctx context.Context, tk *Task) error {
				t.Fatalf("Claim must not be called for an already-owned forward: %s", tk.ID)
				return nil
			},
		}, sendNull, "n1:7500", true, "")

	dup := &Task{ID: "t1", Local: store.Local{Name: "l1"}, Remote: store.Remote{Name: "r1"}, Link: store.Link{RemotePort: 100}}
	tk := &Token{Cycle: 1, State: State{
		Nodes: []Node{
			{ID: "n1", Addr: "n1:7500", Alive: true, Load: Load{MemPct: 10, NetPct: 10}},
			{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 50, NetPct: 50}},
		},
		PendingTasks: map[string]*Task{"t1": dup},
		Topology: map[string]*TopoEntry{"t1": {
			TaskID: "t1", OwnerID: "n2", Local: store.Local{Name: "l1"},
			Remote: store.Remote{Name: "r1"}, Link: store.Link{RemotePort: 100}, Active: true,
		}},
	}}
	if _, err := n1.OnToken(context.Background(), tk); err != nil {
		t.Fatal(err)
	}
	// Pending drained (the stale task was consumed/dropped, not left to ride).
	if got := len(n1.state.PendingList()); got != 0 {
		t.Fatalf("stale task should be dropped, pending=%+v", n1.state.PendingList())
	}
	// Topology untouched: still owned by n2.
	topo := n1.state.TopologyList()
	if len(topo) != 1 || topo[0].OwnerID != "n2" {
		t.Fatalf("topology should remain t1@n2, got %+v", topo)
	}
}