shater/protocol/group/urltest_balance_lx.go

package group

import (
	"context"
	"hash/fnv"
	"math/rand/v2"
	"strconv"
	"sync"
	"sync/atomic"

	"github.com/sagernet/sing-box/adapter"
	"github.com/sagernet/sing-box/common/urltest"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/option"
	E "github.com/sagernet/sing/common/exceptions"
	M "github.com/sagernet/sing/common/metadata"
)

// lx: SPEC 019 v2 — round_robin load-balancing for the urltest group.
//
// least_test (the default) keeps the legacy "pick lowest-delay node" behaviour and is
// handled by URLTestGroup.Select; this file drives round_robin only.
//
// Model: a FIXED-SIZE pool of slots. Slot indices [0..pool-1] never move and are never
// re-sorted — nodes flow THROUGH slots, a replacement takes the exact slot of the node it
// evicts. Two independent orderings (do not conflate — that was the v1 bug):
//   - slot order (fixed)      → sticky binding: slot[hash(key) % pool]
//   - node rank by last delay → only to find which node to evict (computed on the fly)
//
// The pool is lazily health-checked: we test no more nodes than needed to keep it full of
// live nodes. Selection runs once per new connection (DialContext/ListenPacket).

// slot is one fixed position in the pool. tag is the node currently occupying it; live is
// whether the last health-check found that node reachable. pick() routes ONLY through live
// slots (a dead occupant keeps its slot for the SPEC never-shrink invariant, but is skipped
// by selection until a health-check revives or replaces it).
type slot struct {
	tag  string
	live bool
}

// balancer holds the per-group balancing state (round_robin and random). nil for least_test.
type balancer struct {
	poolSize      int      // target slot count (config pool, before min(pool,nodes))
	poolTolerance uint16   // ms; 0 = first-live-fill, > 0 = top-N-by-delay with eviction threshold
	stickyHash    []string // sticky key components; empty → no stickiness (counter rotation)
	random        bool     // mode == random: pick a uniformly-random live slot per connection
	priority      bool     // priority == config order is the ranking (failover); enables fail-back

	// verdict, if set, reads the health-board verdict for a slot tag (wired to
	// URLTestGroup.slotVerdict). Slot liveness is the verdict when the board has fresh
	// data — alive forces live, dead forces dead — and the health-check flag otherwise
	// (untested keeps the optimistic seed usable on cold start). This is what makes a
	// death recorded by ANY prober (group checker, observatory, failed dial) take effect
	// on the next pick instead of the next tick. Health board plan §5.B.
	verdict func(tag string) urltest.HealthVerdict

	access  sync.Mutex // guards slots
	slots   []slot     // the pool; len == min(poolSize, available nodes), index = fixed slot number
	counter atomic.Uint64

	// onChange, if set, is called after the pool occupancy changes (setSlots). The
	// URLTestGroup wires it to invalidate the SPEC 020 reachable cache, since the
	// active pool just changed. Called WITHOUT the access lock held.
	onChange func()
}

// newBalancer builds the balancer from validated options. Returns nil for least_test (and
// empty mode), so callers branch cheaply on a nil balancer.
func newBalancer(options option.URLTestOutboundOptions) (*balancer, error) {
	mode := options.Mode
	if mode == "" || mode == C.URLTestModeLeastTest {
		return nil, nil
	}
	if mode != C.URLTestModeRoundRobin && mode != C.URLTestModeRandom {
		return nil, E.New("unknown urltest mode: ", mode)
	}
	isRandom := mode == C.URLTestModeRandom
	// stickyDefault: round_robin binds a flow to a slot by default (session stability); random is
	// a per-connection independent draw, so stickiness is meaningless there — default it OFF.
	var stickyDefault []string
	if !isRandom {
		stickyDefault = []string{C.URLTestStickyProcess, C.URLTestStickyDomain}
	}
	bo := options.Balancer
	pool := C.DefaultURLTestPool
	var stickyHash []string
	if bo != nil {
		// pool: a Go int with omitempty can't tell "0" from "absent", so 0 means default.
		// Only an explicitly negative value is rejected.
		if bo.Pool < 0 {
			return nil, E.New("urltest balancer.pool must be >= 1")
		}
		if bo.Pool > 0 {
			pool = bo.Pool
		}
		// Omitted (nil) or empty → mode default. To DISABLE stickiness use the explicit ["none"]
		// sentinel: the config decoder collapses a bare [] to nil (see URLTestStickyNone), so an
		// empty list cannot mean "off". ["none"] → stickiness off (no components).
		if len(bo.StickyHash) == 0 {
			stickyHash = stickyDefault
		} else if len(bo.StickyHash) == 1 && bo.StickyHash[0] == C.URLTestStickyNone {
			stickyHash = nil // sticky off → pure counter rotation
		} else {
			stickyHash = bo.StickyHash
		}
	} else {
		// no balancer block → mode defaults (round_robin: sticky on; random: sticky off).
		stickyHash = stickyDefault
	}
	for _, component := range stickyHash {
		switch component {
		case C.URLTestStickyProcess, C.URLTestStickyDomain, C.URLTestStickySourceIP,
			C.URLTestStickyDestIP, C.URLTestStickyDestPort:
		case C.URLTestStickyNone:
			// "none" is only valid as the sole component (handled above); mixed with real
			// components it is ambiguous (disable AND key by something?) — reject.
			return nil, E.New("urltest sticky_hash: \"none\" must be the only component (it disables stickiness)")
		default:
			return nil, E.New("unknown urltest sticky_hash component: ", component)
		}
	}
	var poolTolerance uint16
	var priority bool
	if bo != nil {
		poolTolerance = bo.PoolTolerance
		priority = bo.Priority
	}
	return &balancer{poolSize: pool, poolTolerance: poolTolerance, stickyHash: stickyHash, random: isRandom, priority: priority}, nil
}

// pick selects one node for this connection from the current pool. Selection runs ONLY over
// LIVE slots (health-checked reachable) — a dead slot is never returned while any live slot
// exists, for every balancing mode. fallback is returned when the pool is empty (start, before
// the first health-check fills it) or when every slot is currently dead.
func (b *balancer) pick(ctx context.Context, destination M.Socksaddr, fallback adapter.Outbound, resolve func(tag string) adapter.Outbound) adapter.Outbound {
	b.access.Lock()
	n := len(b.slots)
	if n == 0 {
		b.access.Unlock()
		return fallback
	}
	liveCount := 0
	for i := range b.slots {
		if b.slotIsLive(i) {
			liveCount++
		}
	}
	var tag string
	switch {
	case liveCount == 0:
		// No live slot: nothing safe to route to. fallback (the group's Select) covers it —
		// which itself only returns a node with a working history, else config order.
		b.access.Unlock()
		return fallback
	case b.random:
		// Uniform random draw over the LIVE slots. rand.IntN (math/rand/v2) is safe for
		// concurrent use (per-P generator state, no shared global lock) and auto-seeded, so no
		// mutex or manual seeding is needed — the lock here only guards the slots slice.
		tag = b.nthLiveTag(rand.IntN(liveCount))
	case len(b.stickyHash) > 0:
		// slot-hash: fixed slot index from the key over ALL slots, so a living occupant keeps its
		// keys (the strict-zero-reconnect invariant). If that slot is dead, degrade to a LIVE slot
		// picked deterministically from the same key — the flow still lands somewhere working.
		h := hashKey(b.stickyKey(ctx, destination))
		idx := int(h % uint64(n))
		if b.slotIsLive(idx) {
			tag = b.slots[idx].tag
		} else {
			tag = b.nthLiveTag(int(h % uint64(liveCount)))
		}
	default:
		// plain round-robin over the LIVE slots only.
		idx := int(b.counter.Add(1)-1) % liveCount
		tag = b.nthLiveTag(idx)
	}
	b.access.Unlock()
	if node := resolve(tag); node != nil {
		return node
	}
	return fallback
}

// slotIsLive reports whether slot i is selectable: the board verdict wins when fresh
// (alive → live, dead → dead), the last health-check flag decides for untested slots.
// Caller must hold access. Health board plan §5.B.
func (b *balancer) slotIsLive(i int) bool {
	s := b.slots[i]
	if s.tag == "" {
		return false
	}
	if b.verdict != nil {
		switch b.verdict(s.tag) {
		case urltest.VerdictAlive:
			return true
		case urltest.VerdictDead:
			return false
		}
	}
	return s.live
}

// nthLiveTag returns the tag of the k-th live slot (0-based, in slot order). Caller must hold
// access and pass k in [0, liveCount). Walks without allocating (pools may be large for random).
func (b *balancer) nthLiveTag(k int) string {
	for i := range b.slots {
		if !b.slotIsLive(i) {
			continue
		}
		if k == 0 {
			return b.slots[i].tag
		}
		k--
	}
	return ""
}

// poolTags returns the current slot tags (snapshot under lock). Used by Pool()/GetPool.
func (b *balancer) poolTags() []string {
	b.access.Lock()
	defer b.access.Unlock()
	tags := make([]string, len(b.slots))
	for i, s := range b.slots {
		tags[i] = s.tag
	}
	return tags
}

// setSlots replaces the pool atomically. The health-check (urltest.go) computes the new
// occupancy — which tag sits in which slot — and hands the ordered tag list here, plus the set
// of tags found LIVE this round. A slot whose tag is absent from live is retained (never-shrink)
// but marked dead, so pick() skips it until it is revived or replaced.
func (b *balancer) setSlots(tags []string, live map[string]bool) {
	b.access.Lock()
	b.slots = make([]slot, len(tags))
	for i, tag := range tags {
		b.slots[i] = slot{tag: tag, live: tag != "" && live[tag]}
	}
	b.access.Unlock()
	// lx: SPEC 020 — the active pool changed; invalidate the reachable cache.
	// Called outside the access lock so the callback never nests under it.
	if b.onChange != nil {
		b.onChange()
	}
}

// --- sticky key ---------------------------------------------------------------------

// stickyKey builds the binding key from the configured components, in order. An absent
// component contributes "" (SPEC: all-empty → key "" → one fixed slot).
func (b *balancer) stickyKey(ctx context.Context, destination M.Socksaddr) string {
	metadata := adapter.ContextFrom(ctx)
	var buf []byte
	for i, component := range b.stickyHash {
		if i > 0 {
			buf = append(buf, 0) // NUL separator
		}
		buf = append(buf, stickyComponent(component, metadata, destination)...)
	}
	return string(buf)
}

func stickyComponent(component string, metadata *adapter.InboundContext, destination M.Socksaddr) string {
	switch component {
	case C.URLTestStickyProcess:
		if metadata == nil || metadata.ProcessInfo == nil {
			return ""
		}
		if len(metadata.ProcessInfo.AndroidPackageNames) > 0 {
			return metadata.ProcessInfo.AndroidPackageNames[0]
		}
		return metadata.ProcessInfo.ProcessPath
	case C.URLTestStickyDomain:
		// The router resolves a domain destination to an IP and overwrites metadata.Destination
		// (route.go: `metadata.Destination = M.Socksaddr{Addr: ...}`) BEFORE the group's
		// DialContext runs — so destination.Fqdn is empty here for domain traffic. The original
		// domain survives in metadata.Domain (set by sniffing / reverse mapping). Read that first;
		// fall back to destination.Fqdn only if it is still a domain (no metadata, direct dial).
		if metadata != nil && metadata.Domain != "" {
			return metadata.Domain
		}
		return destination.Fqdn
	case C.URLTestStickySourceIP:
		if metadata == nil || !metadata.Source.Addr.IsValid() {
			return ""
		}
		return metadata.Source.Addr.String()
	case C.URLTestStickyDestIP:
		if !destination.Addr.IsValid() {
			return ""
		}
		return destination.Addr.String()
	case C.URLTestStickyDestPort:
		if destination.Port == 0 {
			return ""
		}
		return strconv.Itoa(int(destination.Port))
	default:
		return ""
	}
}

func hashKey(key string) uint64 {
	h := fnv.New64a()
	_, _ = h.Write([]byte(key))
	return h.Sum64()
}

// --- pool maintenance (pure planning, no I/O) ---------------------------------------
//
// These compute the NEW slot occupancy from current slots + fresh test results, honouring
// the SPEC invariants: fixed slot count, replace-in-slot, never shrink, dead node keeps its
// slot until a live replacement exists. The caller (urltest.go health-check) runs the actual
// URL tests and applies the result via setSlots.

// candidate is a node and its just-measured delay (0 == not measured / dead).
type candidate struct {
	tag   string
	delay uint16
	alive bool
}

// planFirstLivePool (pool_tolerance == 0) computes slot occupancy WITHOUT ranking by delay:
// every live slot member keeps its exact index, dead/empty slots are refilled from fillOrder
// (live nodes not already pooled, in caller order), and any leftover hole keeps a dead member
// so the pool never shrinks. This is the replace-in-slot twin of balancePoolFirstLive, used by
// the manual-test rebuild path where re-ranking would needlessly relocate living nodes and
// break their sticky bindings.
//
// current = present slot tags (slot order). live = set of tags that tested alive now.
// fillOrder = candidate tags to drop into holes (already filtered to non-pool, in order).
func planFirstLivePool(current []string, live map[string]bool, fillOrder []string, size int) []string {
	slotCount := size
	if len(current) > slotCount {
		slotCount = len(current)
	}
	next := make([]string, slotCount)
	for i, tag := range current {
		if tag != "" && live[tag] {
			next[i] = tag
		}
	}
	placed := make(map[string]bool, len(next))
	for _, tag := range next {
		if tag != "" {
			placed[tag] = true
		}
	}
	// Refill holes with fresh live nodes, writing each into the hole's own index.
	fi := 0
	for i := range next {
		if next[i] != "" {
			continue
		}
		for fi < len(fillOrder) {
			tag := fillOrder[fi]
			fi++
			if tag != "" && live[tag] && !placed[tag] {
				next[i] = tag
				placed[tag] = true
				break
			}
		}
	}
	// Any hole still open: keep a dead member there (never shrink). A dead occupant first keeps
	// the slot it already held; surplus dead members fall into leftover holes.
	for i, tag := range current {
		if i < len(next) && next[i] == "" && tag != "" && !live[tag] && !placed[tag] {
			next[i] = tag
			placed[tag] = true
		}
	}
	for i := range next {
		if next[i] != "" {
			continue
		}
		for _, tag := range current {
			if tag != "" && !live[tag] && !placed[tag] {
				next[i] = tag
				placed[tag] = true
				break
			}
		}
	}
	return next
}

// planPriorityPool (priority balancer, used by strategy=failover) makes the CONFIG ORDER the
// ranking: the pool holds the first `size` LIVE nodes in configuration order. Unlike
// planFirstLivePool / planTolerantPool it deliberately does NOT honour the replace-in-slot
// invariant — a live lower-priority node is never allowed to keep a slot ahead of a
// higher-priority live one. That is exactly what gives fail-back: a revived node #1 re-takes
// slot 0 on the next tick even while #2 is still live. Relocating a node across slots is safe
// here ONLY because the priority balancer always runs with stickiness off (failoverBalancer
// forces StickyHash ["none"]), so no per-flow key is bound to a slot index and moving a node
// between slots breaks nothing.
//
// When fewer than `size` nodes are live the remaining slots are filled with DEAD nodes in
// config order (SPEC never-shrink: the pool keeps its size so it can never collapse to empty).
// A dead occupant is only ever reached when EVERY slot is dead, and pick() returns the group
// fallback in that state regardless, so which dead node holds a hole is immaterial.
//
// configOrder = every member tag in configuration order. live = set of tags that tested alive.
func planPriorityPool(configOrder []string, live map[string]bool, size int) []string {
	next := make([]string, 0, size)
	placed := make(map[string]bool, size)
	// 1. First `size` LIVE nodes in config order — the priority ranking.
	for _, tag := range configOrder {
		if len(next) >= size {
			break
		}
		if tag != "" && live[tag] && !placed[tag] {
			next = append(next, tag)
			placed[tag] = true
		}
	}
	// 2. Not enough live nodes: fill the remaining slots with dead nodes in config order so the
	// pool never shrinks below `size`. A dead node never displaces a live one (live went first).
	for _, tag := range configOrder {
		if len(next) >= size {
			break
		}
		if tag != "" && !placed[tag] {
			next = append(next, tag)
			placed[tag] = true
		}
	}
	return next
}

// planTolerantPool (pool_tolerance > 0): pick the top-`size` live candidates by delay, but
// keep a node in its current slot when it survives (replace-in-slot, no needless churn). A
// living slot is only displaced when some out-of-pool node is faster than it by > tolerance.
//
// current = present slot tags (in slot order). results = delay for every node that was
// tested (alive ones); size = min(poolSize, len(all nodes)).
func planTolerantPool(current []string, results map[string]candidate, size int, tolerance uint16) []string {
	// Rank all alive candidates by delay ascending (tag breaks ties — deterministic).
	alive := make([]candidate, 0, len(results))
	for _, c := range results {
		if c.alive {
			alive = append(alive, c)
		}
	}
	sortCandidatesByDelay(alive)

	next := make([]string, len(current))
	copy(next, current)
	// Grow to size if we have spare slots and alive nodes not yet placed.
	inPool := make(map[string]bool, len(next))
	for _, tag := range next {
		if tag != "" {
			inPool[tag] = true
		}
	}
	for len(next) < size {
		next = append(next, "")
	}
	// Fill empty slots first with the fastest unused alive nodes.
	for i := range next {
		if next[i] != "" {
			continue
		}
		for _, c := range alive {
			if !inPool[c.tag] {
				next[i] = c.tag
				inPool[c.tag] = true
				break
			}
		}
	}
	// Eviction: for each slot, if a faster-than-by-tolerance unused alive node exists, swap it
	// IN PLACE. We do NOT return the evicted occupant to the candidate set: re-inserting it at a
	// later slot would relocate a node across slots and move every sticky key bound to it. A
	// surviving occupant therefore always keeps its slot index; only the evicted-and-replaced
	// slot changes its node (the SPEC replace-in-slot invariant).
	for i := range next {
		occupant := next[i]
		occDelay, occAlive := slotDelay(occupant, results)
		for _, c := range alive {
			if inPool[c.tag] {
				continue
			}
			// Replace if the slot is dead, or the candidate beats it by > tolerance.
			if !occAlive || (occAlive && c.delay+tolerance < occDelay) {
				next[i] = c.tag
				inPool[c.tag] = true
				break
			}
		}
	}
	return next
}

func slotDelay(tag string, results map[string]candidate) (uint16, bool) {
	if tag == "" {
		return 0, false
	}
	if c, ok := results[tag]; ok {
		return c.delay, c.alive
	}
	return 0, false
}

func sortCandidatesByDelay(cs []candidate) {
	// insertion sort — pool/result sets are tiny; avoids pulling sort.Slice closure cost.
	for i := 1; i < len(cs); i++ {
		for j := i; j > 0; j-- {
			a, b := cs[j-1], cs[j]
			if a.delay < b.delay || (a.delay == b.delay && a.tag <= b.tag) {
				break
			}
			cs[j-1], cs[j] = cs[j], cs[j-1]
		}
	}
}

// ActiveTags returns the tags this urltest group is CURRENTLY routing through —
// the whole rotation pool for a round_robin (balanced) group, or the single
// current node for a legacy least_test group. Used by the SPEC 020 reachability
// walk: every node a balanced group could dial right now is reachable (and must
// not be idle-suspended), not just Now(). Returns nil for an empty/cold group.
func (s *URLTest) ActiveTags() []string {
	if s.balancer != nil {
		if tags := s.balancer.poolTags(); len(tags) > 0 {
			return tags
		}
		// Cold start: pool not yet filled — fall back to the single current pick.
	}
	if now := s.Now(); now != "" {
		return []string{now}
	}
	return nil
}