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