webui4frpc/internal/cluster/ring.go

// Package cluster implements the token-ring cooperative network.
// Authoritative design: plan.md §M6 (令牌环网拓扑).
//
// Token payload semantics (per the design):
//   - PendingTasks: forward requests (中间转发意图) riding in the token,
//     waiting to be claimed by the lowest-load node. Once claimed, the task
//     disappears from PendingTasks.
//   - Topology: the cluster's ACTIVE forward map — every established forward
//     with its owning node. The claiming node appends it back into the token
//     so every member converges on a complete picture (who runs what).
package cluster

import (
	"fmt"
	"sort"
	"time"

	"webui4frpc/internal/store"
)

// ---- 数据面 ---- //

// Node is one cluster member in the ring.
type Node struct {
	ID       string   `json:"id"`
	Addr     string   `json:"addr"`
	IsLeader bool     `json:"isLeader,omitempty"`
	Alive    bool     `json:"alive"`
	Load     Load     `json:"load"`
	Version  string   `json:"version,omitempty"`
	Cache    []string `json:"cache,omitempty"`
	LastSeen int64    `json:"lastSeen,omitempty"`
	// NodeKey is this member's cluster admission key. It travels in the
	// token so every node knows every peer's key — a crashed node can
	// rejoin via ANY cached peer by presenting that peer's key. Without
	// this, a rejoiner would only know its original sponsor's key (from
	// the -join-key flag) and couldn't rejoin through a different peer.
	NodeKey string `json:"nodeKey,omitempty"`
}

// Load is the combined load metric used to pick the task claimer.
type Load struct {
	MemPct   float64 `json:"memPct"`
	NetPct   float64 `json:"netPct"`
	Forwards int     `json:"forwards,omitempty"` // active forwards owned by this node (primary signal)
}

// Score weights owned forwards heavily so the node with the fewest active
// forwards is picked first; mem/net only break ties at equal forward count.
func (l Load) Score() float64 { return float64(l.Forwards)*100 + l.MemPct + l.NetPct }

// Task is a PENDING forward request circulated in the token. It carries the
// intermediate forwarding intent (local/remote/link) — NOT a rendered frpc
// config. The claiming node renders its local frpc config and spawns the worker.
type Task struct {
	ID      string       `json:"id"`
	Local   store.Local  `json:"local"`
	Remote  store.Remote `json:"remote"`
	Link    store.Link   `json:"link"`
	Created int64        `json:"created"`
	// Revoke marks a REVOCATION task: instead of creating the forward, the
	// owning node cancels it (stop worker, drop from topology). Reuses the
	// same publish channel as creation (round-1 inject, round-2 apply).
	Revoke bool `json:"revoke,omitempty"`
	// Restart marks a RE-ENABLE task: the forward is already claimed and its
	// topology entry still exists (a stop keeps the entry as the flag's
	// carrier), so the creation path would dedupe the submission and the
	// duplicate-claim guard would discard it. The owner applies this one
	// unconditionally: re-mark the entry enabled and (re)spawn the worker.
	Restart bool `json:"restart,omitempty"`
	// RemoveNode: node ID to remove from the ring; the node self-removes when
	// the command reaches it via the token.
	RemoveNode string `json:"removeNode,omitempty"`
}

// TopoEntry is an ESTABLISHED forward in the cluster topology: who (OwnerID)
// runs which forward (rendered from the same intermediate intent).
type TopoEntry struct {
	TaskID  string       `json:"taskId"`
	OwnerID string       `json:"ownerId"`
	Local   store.Local  `json:"local"`
	Remote  store.Remote `json:"remote"`
	Link    store.Link   `json:"link"`
	Active  bool         `json:"active"`
}

// State is the full cluster picture, replicated on every node:
// ring membership + pending tasks + active forward topology.
type State struct {
	LeaderID string `json:"leaderId"`
	Nodes    []Node `json:"nodes"`
	Cycle    int64  `json:"cycle"`
	// RoundDelay rides the token so every node converges on the same ring
	// cadence (serialized as nanoseconds; the leader refreshes it each round
	// via tkDelaySince). MUST be serialized — otherwise a token round-trip
	// zeroed it and LossTimeout collapsed to the floor.
	RoundDelay time.Duration `json:"roundDelay,omitempty"`
	// PendingTasks: not yet claimed (disappear when claimed).
	PendingTasks map[string]*Task `json:"pendingTasks,omitempty"`
	// Topology: active forwards owned by members (full cluster view).
	Topology map[string]*TopoEntry `json:"topology,omitempty"`
}

// Token is the circulating message: one physical token per cycle (single
// round — adopt + append + execute in one pass, no phase split).
type Token struct {
	Cycle  int64      `json:"cycle"`
	State  State      `json:"state"`
	Log    []LogEntry `json:"logDelta,omitempty"`
	SentAt int64      `json:"sentAt,omitempty"`
}

// ---- 环图操作 ---- //

func (s *State) Find(id string) int {
	for i, n := range s.Nodes {
		if n.ID == id {
			return i
		}
	}
	return -1
}

func (s *State) AliveSuccessor(id string) (string, bool) {
	i := s.Find(id)
	if i < 0 {
		return "", false
	}
	n := len(s.Nodes)
	for k := 1; k <= n; k++ {
		j := (i + k) % n
		if s.Nodes[j].Alive {
			return s.Nodes[j].ID, true
		}
	}
	return "", false
}

func (s *State) AlivePredecessor(id string) (string, bool) {
	i := s.Find(id)
	if i < 0 {
		return "", false
	}
	n := len(s.Nodes)
	for k := 1; k <= n; k++ {
		j := ((i-k)%n + n) % n
		if s.Nodes[j].Alive {
			return s.Nodes[j].ID, true
		}
	}
	return "", false
}

func (s *State) UpsertNode(n Node) {
	i := s.Find(n.ID)
	if i < 0 {
		if n.IsLeader {
			s.LeaderID = n.ID
		}
		n.Alive = true
		n.LastSeen = time.Now().Unix()
		s.Nodes = append(s.Nodes, n)
		return
	}
	old := s.Nodes[i]
	n.IsLeader = old.IsLeader || n.IsLeader
	if n.IsLeader {
		s.LeaderID = n.ID
	}
	n.Alive = true
	n.LastSeen = time.Now().Unix()
	s.Nodes[i] = n
}

func (s *State) MarkOffline(id string) int {
	i := s.Find(id)
	if i < 0 {
		return -1
	}
	s.Nodes[i].Alive = false
	s.Nodes[i].LastSeen = time.Now().Unix()
	return i
}

// InsertAfter inserts newNode right after anchor (join: newNode lands behind
// its sponsor; the former successor moves after the newcomer).
func (s *State) InsertAfter(anchor string, n Node) {
	i := s.Find(n.ID)
	if i >= 0 {
		s.Nodes[i].Alive = true
		s.Nodes[i].LastSeen = time.Now().Unix()
		return
	}
	n.Alive = true
	n.LastSeen = time.Now().Unix()
	j := s.Find(anchor)
	if j < 0 {
		s.Nodes = append(s.Nodes, n)
		return
	}
	out := make([]Node, 0, len(s.Nodes)+1)
	out = append(out, s.Nodes[:j+1]...)
	out = append(out, n)
	out = append(out, s.Nodes[j+1:]...)
	s.Nodes = out
}

// LowestAlive returns alive node with lowest combined load (nil if none).
func (s *State) LowestAlive() *Node {
	var best *Node
	for i := range s.Nodes {
		if !s.Nodes[i].Alive {
			continue
		}
		if best == nil || s.Nodes[i].Load.Score() < best.Load.Score() {
			b := s.Nodes[i]
			best = &b
		}
	}
	return best
}

func (s *State) NextTaskID() string {
	// Collision-free id allocation by scanning the ids actually in flight
	// (pending + topology) and taking max+1. The ring is mutually exclusive
	// (one token holder at a time), so when a node mints an id it has the
	// authoritative full view — scanning existing ids guarantees a fresh id.
	// n is tiny (handful of forwards). No Seq counter is needed: a cross-node
	// Seq was previously adopted wholesale on every OnToken (e.state = tk.State),
	// which dropped local increments and could regress below an id still in use,
	// recycling it and overwriting an active topology entry.
	max := int64(0)
	for id := range s.PendingTasks {
		if n := taskIDNum(id); n > max {
			max = n
		}
	}
	for _, e := range s.Topology {
		if n := taskIDNum(e.TaskID); n > max {
			max = n
		}
	}
	return fmt.Sprintf("t%d", max+1)
}

// taskIDNum extracts the numeric suffix of a task id "t12" -> 12 (0 if it does
// not parse). Used only to keep NextTaskID collision-free.
func taskIDNum(id string) int64 {
	if len(id) < 2 || id[0] != 't' {
		return 0
	}
	var n int64
	for _, c := range id[1:] {
		if c < '0' || c > '9' {
			return 0
		}
		n = n*10 + int64(c-'0')
	}
	return n
}

// AddRemoveNode publishes a node-removal command via the token; the target
// node self-removes when it receives the command (re-queue forwards, drop
// ring position, pass token to former successor).
func (s *State) AddRemoveNode(nodeID string) *Task {
	t := &Task{
		ID: s.NextTaskID(), Created: time.Now().Unix(), RemoveNode: nodeID,
	}
	if s.PendingTasks == nil {
		s.PendingTasks = map[string]*Task{}
	}
	s.PendingTasks[t.ID] = t
	return t
}

// AddRevoke publishes a REVOCATION task (same channel as creation): when a
// node wants to cancel an established forward, it injects a Revoke task;
// the owning node cancels the worker and drops it from topology.
func (s *State) AddRevoke(local store.Local, remote store.Remote, link store.Link) *Task {
	t := &Task{
		ID:      s.NextTaskID(),
		Local:   local,
		Remote:  remote,
		Link:    link,
		Created: time.Now().Unix(),
		Revoke:  true,
	}
	if s.PendingTasks == nil {
		s.PendingTasks = map[string]*Task{}
	}
	s.PendingTasks[t.ID] = t
	return t
}

// AddPending attaches a NEW forward request to the token (round-1 inject).
func (s *State) AddPending(local store.Local, remote store.Remote, link store.Link) *Task {
	t := &Task{
		ID:      s.NextTaskID(),
		Local:   local,
		Remote:  remote,
		Link:    link,
		Created: time.Now().Unix(),
	}
	if s.PendingTasks == nil {
		s.PendingTasks = map[string]*Task{}
	}
	s.PendingTasks[t.ID] = t
	return t
}

// PendingList returns pending tasks, stable by id.
func (s *State) PendingList() []*Task {
	out := make([]*Task, 0, 4)
	for _, t := range s.PendingTasks {
		out = append(out, t)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out
}

// ClaimPending removes a pending task and returns it (the claimer then renders
// its local frpc config and spawns the worker). Returns nil if absent.
func (s *State) ClaimPending(id string) *Task {
	if s.PendingTasks == nil {
		return nil
	}
	t := s.PendingTasks[id]
	if t == nil {
		return nil
	}
	delete(s.PendingTasks, id)
	return t
}

// AddTopology records an ESTABLISHED forward into the cluster topology.
// Called by the claiming node right after it spawns the worker: the task
// disappears from pending and lives on as active topology.
func (s *State) AddTopology(t *Task, ownerID string) *TopoEntry {
	if s.Topology == nil {
		s.Topology = map[string]*TopoEntry{}
	}
	// Active mirrors the task's disabled flag rather than being unconditionally
	// true. A claim can legitimately carry a stopped forward (e.g. a node
	// re-claiming from a departed peer), and forcing Active=true there would
	// resurrect it: OfflineReassign only re-queues Active entries, and every
	// reader that filters on Active would treat it as running again.
	e := &TopoEntry{
		TaskID: t.ID, OwnerID: ownerID,
		Local: t.Local, Remote: t.Remote, Link: t.Link,
		Active: !t.Link.Disabled,
	}
	s.Topology[t.ID] = e
	return e
}

// SelfRemove removes this node from the ring: re-queues its own forwards
// as pending (so another member claims them), drops its topology entries
// and removes its ring position. Returns the re-queued task ids.
func (s *State) SelfRemove(nodeID string) []string {
	reattached := s.OfflineReassign(nodeID)
	i := s.Find(nodeID)
	if i >= 0 {
		s.Nodes = append(s.Nodes[:i], s.Nodes[i+1:]...)
	}
	if s.LeaderID == nodeID {
		s.LeaderID = ""
	}
	return reattached
}

// RemoveTopology drops the active forward for the given local/remote/port
// (returns true if removed — the owning node revokes it).
func (s *State) RemoveTopology(local, remote string, port int) bool {
	for id, e := range s.Topology {
		if e.Local.Name == local && e.Remote.Name == remote && e.Link.RemotePort == port {
			delete(s.Topology, id)
			return true
		}
	}
	return false
}

// TopologyOwner returns the node that owns the active forward matching the
// given task (by local/remote/port), or "" if no such entry exists. Used to
// route a revocation to the OWNING node (plan §任务撤销: "持有该转发的节点
// 收到撤销任务后取消") instead of letting any lowest-load node claim it.
func (s *State) TopologyOwner(t *Task) string {
	for _, e := range s.Topology {
		if e.Local.Name == t.Local.Name && e.Remote.Name == t.Remote.Name && e.Link.RemotePort == t.Link.RemotePort {
			return e.OwnerID
		}
	}
	return ""
}

// TopologyList returns active forwards, stable by task id.
func (s *State) TopologyList() []*TopoEntry {
	out := make([]*TopoEntry, 0, len(s.Topology))
	for _, e := range s.Topology {
		out = append(out, e)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].TaskID < out[j].TaskID })
	return out
}

// ForwardsOwnedBy returns topology entries owned by a node (for offline
// reassignment and UI per-node views).
func (s *State) ForwardsOwnedBy(owner string) []*TopoEntry {
	var out []*TopoEntry
	for _, e := range s.Topology {
		if e.OwnerID == owner {
			out = append(out, e)
		}
	}
	return out
}

// UpdateTopologyGroup sets the group label on the topology entry matching the
// (local, remote, port) triple. Returns true if found. The updated entry
// propagates to all nodes via the next token cycle — group changes sync
// through the ring without a dedicated command.
// UpdateTopologyDisabled sets the disabled flag on the topology entry matching
// the forward's natural key. Mirror of UpdateTopologyGroup: it makes a locally
// decided stop/start part of the topology so the next token cycle carries it to
// every other member (the alternative — a one-shot revoke task addressed at the
// owner — only converges if that owner happens to be online right then).
//
// Returns false when the forward is not in the topology, which is not an error:
// a stopped forward may legitimately have no entry yet.
func (s *State) UpdateTopologyDisabled(local, remote string, port int, disabled bool) bool {
	for _, e := range s.Topology {
		if e.Local.Name == local && e.Remote.Name == remote && e.Link.RemotePort == port {
			e.Link.Disabled = disabled
			// Keep Active consistent with the flag so readers that key off it
			// (status page, load accounting) agree with Link.Disabled.
			e.Active = !disabled
			return true
		}
	}
	return false
}

// TopologyDisabled returns the disabled flag the cluster currently holds for a
// forward, and whether such an entry exists. Used by the adoption path to learn
// a peer's decision into the local store.
func (s *State) TopologyDisabled(local, remote string, port int) (bool, bool) {
	for _, e := range s.Topology {
		if e.Local.Name == local && e.Remote.Name == remote && e.Link.RemotePort == port {
			return e.Link.Disabled, true
		}
	}
	return false, false
}

func (s *State) UpdateTopologyGroup(local, remote string, port int, group string) bool {
	for _, e := range s.Topology {
		if e.Local.Name == local && e.Remote.Name == remote && e.Link.RemotePort == port {
			e.Link.Group = group
			return true
		}
	}
	return false
}

// OfflineReassign moves all active forwards owned by an offline node back into
// PendingTasks (they become new tasks for the next lowest-load member). Returns
// the reassigned task ids.
func (s *State) OfflineReassign(offlineID string) []string {
	var reattached []string
	for id, e := range s.Topology {
		if e.OwnerID != offlineID || !e.Active {
			continue
		}
		delete(s.Topology, id)
		t := &Task{ID: id, Local: e.Local, Remote: e.Remote, Link: e.Link,
			Created: time.Now().Unix()}
		if s.PendingTasks == nil {
			s.PendingTasks = map[string]*Task{}
		}
		s.PendingTasks[id] = t
		reattached = append(reattached, id)
	}
	return reattached
}