webui4frpc/internal/cluster/ring_remove_test.go

package cluster

import (
	"context"
	"testing"

	"webui4frpc/internal/store"
)

// TestSelfRemoveViaToken: a node-removal command reaches the target node
// through OnToken; it stops owned workers, re-queues its forwards, drops its
// ring position, and the token flows to the original successor.
func TestSelfRemoveViaToken(t *testing.T) {
	eng := newTestEngine("n1", true)
	eng.state.InsertAfter("n1", Node{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 10, NetPct: 10}})

	// n1 owns a forward (via topology)
	p := eng.state.AddPending(store.Local{Name: "web"}, store.Remote{Name: "frps1"}, store.Link{RemotePort: 18081})
	eng.state.ClaimPending(p.ID)
	eng.state.AddTopology(p, "n1")

	// n2 publishes remove-node for itself
	rm := eng.state.AddRemoveNode("n2")
	if rm.RemoveNode != "n2" {
		t.Fatalf("remove node cmd = %+v", rm)
	}

	// run OnToken at n2 (lowest-load, so it claims the remove command).
	// State order is [n2, n1] so LowestAlive picks n2.
	eng2 := newTestEngine("n2", false)
	eng2.state.UpsertNode(Node{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 10, NetPct: 10}})
	eng2.state.UpsertNode(Node{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true, Load: Load{MemPct: 90, NetPct: 90}})
	eng2.state.Topology = map[string]*TopoEntry{
		p.ID: {TaskID: p.ID, OwnerID: "n2", Active: true, Local: p.Local, Remote: p.Remote, Link: p.Link},
	}
	eng2.state.PendingTasks = map[string]*Task{rm.ID: rm}
	out, err := eng2.OnToken(context.Background(), &Token{Cycle: 1, State: eng2.state})
	if err != nil {
		t.Fatal(err)
	}
	if out == nil {
		t.Fatal("nil token after OnToken")
	}
	// n2 removed itself from ring (NOT re-added by the UpsertNode step)
	if eng2.state.Find("n2") >= 0 {
		t.Fatalf("n2 still in ring: %+v", eng2.state.Nodes)
	}
	// the forward n2 owned was re-queued as pending for another member
	if len(eng2.state.PendingList()) == 0 {
		t.Fatalf("no pending re-queue after self-remove: %+v", eng2.state.PendingList())
	}
	// removedNext captured so Forward can hand the token to the old successor
	if eng2.removedNext != "n1" {
		t.Fatalf("removedNext = %q want n1", eng2.removedNext)
	}
}

// TestRemoveCommandRidesPastLowest: a node-removal command directed at a
// NON-lowest node must NOT be claimed by the lowest-load node (the old bug
// corrupted it into a phantom forward via Handler.Claim). It stays pending
// and rides the token until it reaches the target, which then self-removes.
func TestRemoveCommandRidesPastLowest(t *testing.T) {
	// 3-node ring order [n2, n1, n3]: n2 is lowest, n3 is the remove target.
	eng2 := newTestEngine("n2", false)
	eng2.state.UpsertNode(Node{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true, Load: Load{MemPct: 50, NetPct: 50}})
	eng2.state.UpsertNode(Node{ID: "n3", Addr: "n3:7500", Alive: true, Load: Load{MemPct: 90, NetPct: 90}})
	rm := eng2.state.AddRemoveNode("n3")

	// Token reaches n2 (lowest) carrying the remove command for n3.
	if _, err := eng2.OnToken(context.Background(), &Token{Cycle: 1, State: eng2.state}); err != nil {
		t.Fatal(err)
	}
	// n2 must NOT have consumed the remove — it must still be pending, riding.
	if _, ok := eng2.state.PendingTasks[rm.ID]; !ok {
		t.Fatal("lowest node n2 consumed a remove command aimed at n3 (should ride to target)")
	}
	if i := eng2.state.Find("n3"); i < 0 {
		t.Fatal("n3 was removed by non-target n2 (should still be in ring)")
	}

	// Token now reaches the TARGET (n3): it self-removes + captures successor.
	eng3 := newTestEngine("n3", false)
	eng3.state.UpsertNode(Node{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true, Load: Load{MemPct: 50, NetPct: 50}})
	eng3.state.UpsertNode(Node{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 10, NetPct: 10}})
	eng3.state.PendingTasks = map[string]*Task{rm.ID: rm}
	if _, err := eng3.OnToken(context.Background(), &Token{Cycle: 2, State: eng3.state}); err != nil {
		t.Fatal(err)
	}
	if eng3.state.Find("n3") >= 0 {
		t.Fatalf("n3 did not self-remove: %+v", eng3.state.Nodes)
	}
	// Ring order [n3, n1, n2] → n3's original successor is n1.
	if eng3.removedNext != "n1" {
		t.Fatalf("removedNext = %q want n1", eng3.removedNext)
	}
}

// TestRemoveCommandFulfilledNoPhantom: when a remove command's target has
// already left the ring (or was never a member), the lowest-load node that
// re-encounters it must DROP it — NOT fall through to Handler.Claim and spawn
// a phantom empty forward (the remove task carries no local/remote/link).
func TestRemoveCommandFulfilledNoPhantom(t *testing.T) {
	var claimed []*Task
	eng := NewEngine("n1", "n1:7500", "u", "p", "0.1.0", nil,
		&fakeHandler{load: Load{MemPct: 5, NetPct: 5},
			claim: func(ctx context.Context, tk *Task) error { claimed = append(claimed, tk); return nil }},
		sendNull, "n1:7500", true, "")
	rm := eng.state.AddRemoveNode("node-x:7500") // target not in the ring
	if _, err := eng.OnToken(context.Background(), &Token{Cycle: 1, State: eng.state}); err != nil {
		t.Fatal(err)
	}
	if _, ok := eng.state.PendingTasks[rm.ID]; ok {
		t.Fatal("fulfilled remove command was not dropped from pending")
	}
	if len(claimed) != 0 {
		t.Fatalf("Handler.Claim called for a remove command: %+v", claimed)
	}
	if len(eng.state.TopologyList()) != 0 {
		t.Fatalf("phantom topology entry created: %+v", eng.state.TopologyList())
	}
}

// TestLeaderSelfRemoveDesignatesSuccessor: when the LEADER self-removes, it
// must designate its successor as the new leader before publishing the token.
// Without this, LeaderID would be "" (SelfRemove clears it), no node would
// call Send (cycle never advances), and WatchLeader can't find
// AlivePredecessor("") to promote anyone — the ring runs leaderless and dies.
func TestLeaderSelfRemoveDesignatesSuccessor(t *testing.T) {
	eng := newTestEngine("n1", true)
	eng.state.InsertAfter("n1", Node{
		ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 90, NetPct: 90},
	})
	// n1 (leader) publishes a remove-node command for itself.
	rm := eng.state.AddRemoveNode("n1")
	eng.state.PendingTasks = map[string]*Task{rm.ID: rm}

	out, err := eng.OnToken(context.Background(), &Token{Cycle: 1, State: eng.state})
	if err != nil {
		t.Fatal(err)
	}
	if out == nil {
		t.Fatal("nil token after OnToken")
	}
	// n1 removed itself from the ring.
	if eng.state.Find("n1") >= 0 {
		t.Fatalf("n1 still in ring: %+v", eng.state.Nodes)
	}
	// n2 designated as the new leader (not "" — the old bug).
	if eng.state.LeaderID != "n2" {
		t.Fatalf("LeaderID = %q want n2 (successor should be designated as leader)", eng.state.LeaderID)
	}
	// n2 marked IsLeader in Nodes.
	if i := eng.state.Find("n2"); i >= 0 && !eng.state.Nodes[i].IsLeader {
		t.Fatalf("n2 not marked IsLeader: %+v", eng.state.Nodes[i])
	}
	// removedNext captured so Forward can hand the token to the old successor.
	if eng.removedNext != "n2" {
		t.Fatalf("removedNext = %q want n2", eng.removedNext)
	}
}

// TestDetachAfterForward: after self-leave + Forward (token handed to the old
// successor), the engine resets to a fresh standalone state so Snapshot() no
// longer serves the old cluster picture (members, topology, pending, log).
func TestDetachAfterForward(t *testing.T) {
	eng := newTestEngine("n1", true)
	eng.state.InsertAfter("n1", Node{
		ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 90, NetPct: 90},
	})
	// Give n1 an owned forward so the old state has a non-empty topology.
	p := eng.state.AddPending(store.Local{Name: "web"}, store.Remote{Name: "frps1"}, store.Link{RemotePort: 18081})
	eng.state.ClaimPending(p.ID)
	eng.state.AddTopology(p, "n1")

	rm := eng.state.AddRemoveNode("n1")
	eng.state.PendingTasks = map[string]*Task{rm.ID: rm}

	out, err := eng.OnToken(context.Background(), &Token{Cycle: 1, State: eng.state})
	if err != nil {
		t.Fatal(err)
	}
	if out == nil {
		t.Fatal("nil token after OnToken")
	}
	// Before Forward: self-removed, removedNext captured, old state retained
	// (n2 still in Nodes, n1's forward re-queued as pending by SelfRemove).
	if !eng.selfRemoved {
		t.Fatal("selfRemoved not set before Forward")
	}
	if eng.removedNext != "n2" {
		t.Fatalf("removedNext = %q want n2", eng.removedNext)
	}
	if eng.state.Find("n2") < 0 {
		t.Fatal("n2 (old member) missing before Forward — test setup wrong")
	}

	// Forward hands off the token (sendNull no-op) then detaches to standalone.
	if err := eng.Forward(context.Background(), out); err != nil {
		t.Fatal(err)
	}
	// After Forward: fresh standalone state.
	if eng.state.LeaderID != "n1" {
		t.Fatalf("LeaderID = %q want n1 (standalone leader after detach)", eng.state.LeaderID)
	}
	if len(eng.state.Nodes) != 1 || eng.state.Nodes[0].ID != "n1" {
		t.Fatalf("state not standalone (want just n1): %+v", eng.state.Nodes)
	}
	if len(eng.state.Topology) != 0 {
		t.Fatalf("topology not cleared after detach: %+v", eng.state.Topology)
	}
	if len(eng.state.PendingTasks) != 0 {
		t.Fatalf("pending not cleared after detach: %+v", eng.state.PendingList())
	}
	if eng.selfRemoved {
		t.Fatal("selfRemoved should be cleared after detach")
	}
	if eng.removedNext != "" {
		t.Fatalf("removedNext should be empty after detach, got %q", eng.removedNext)
	}
}