public/headscale
1
0
Fork 0
mirror of https://github.com/juanfont/headscale.git synced 2026-07-07 16:40:21 +09:00
Code Issues Packages Projects Releases Wiki Activity

state: port primary route tests, delete routes package

Browse Source 
The primary route algorithm now lives in hscontrol/state/node_store.go
(snapshotFromNodes / computePrimaries). Move the test coverage that
used to live in hscontrol/routes alongside it: primaries_test.go ports
the named scenarios (anti-flap, HA failover, recovery-no-flap, exit
routes excluded, the issue #3203 both-offline regression) and HANodes,
primaries_property_test.go ports the rapid-driven invariant check.

Both tests drive the algorithm via NodeStore.UpdateNode rather than
the deleted PrimaryRoutes API. With every consumer migrated and the
algorithm relocated, the routes package has no remaining users and
gets removed.

Updates #3203
This commit is contained in:
Kristoffer Dalby
2026-04-28 13:34:33 +00:00
parent 5dd622a31b
commit 436aa298b7
5 changed files with 650 additions and 1492 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hscontrol/routes/primary.go
-461
View File
@@ -1,461 +0,0 @@
package routes
import (
"fmt"
"net/netip"
"slices"
"sort"
"strings"
"sync"
"github.com/juanfont/headscale/hscontrol/types"
"github.com/juanfont/headscale/hscontrol/util"
"github.com/juanfont/headscale/hscontrol/util/zlog/zf"
"github.com/rs/zerolog/log"
xmaps "golang.org/x/exp/maps"
"tailscale.com/net/tsaddr"
"tailscale.com/util/set"
)
type PrimaryRoutes struct {
mu sync.Mutex
// routes is a map of prefixes that are adverties and approved and available
// in the global headscale state.
routes map[types.NodeID]set.Set[netip.Prefix]
// primaries is a map of prefixes to the node that is the primary for that prefix.
primaries map[netip.Prefix]types.NodeID
isPrimary map[types.NodeID]bool
// unhealthy tracks nodes that failed health probes. Unhealthy nodes
// are skipped during primary selection but remain in the routes map.
unhealthy set.Set[types.NodeID]
}
func New() *PrimaryRoutes {
return &PrimaryRoutes{
routes: make(map[types.NodeID]set.Set[netip.Prefix]),
primaries: make(map[netip.Prefix]types.NodeID),
isPrimary: make(map[types.NodeID]bool),
unhealthy: make(set.Set[types.NodeID]),
}
}
// updatePrimaryLocked recalculates the primary routes and updates the internal state.
// It returns true if the primary routes have changed.
// It is assumed that the caller holds the lock.
// The algorithm is as follows:
// 1. Reset the primaries map.
// 2. Iterate over the routes and count the number of times a prefix is advertised.
// 3. If a prefix is advertised by at least two nodes, it is a primary route.
// 4. If the primary routes have changed, update the internal state and return true.
// 5. Otherwise, return false.
func (pr *PrimaryRoutes) updatePrimaryLocked() bool {
log.Debug().Caller().Msg("updatePrimaryLocked starting")
// reset the primaries map, as we are going to recalculate it.
allPrimaries := make(map[netip.Prefix][]types.NodeID)
pr.isPrimary = make(map[types.NodeID]bool)
changed := false
// sort the node ids so we can iterate over them in a deterministic order.
// this is important so the same node is chosen two times in a row
// as the primary route.
ids := types.NodeIDs(xmaps.Keys(pr.routes))
sort.Sort(ids)
// Create a map of prefixes to nodes that serve them so we
// can determine the primary route for each prefix.
for _, id := range ids {
routes := pr.routes[id]
for route := range routes {
if _, ok := allPrimaries[route]; !ok {
allPrimaries[route] = []types.NodeID{id}
} else {
allPrimaries[route] = append(allPrimaries[route], id)
}
}
}
// Go through all prefixes and determine the primary route for each.
// If the number of routes is below the minimum, remove the primary.
// If the current primary is still available, continue.
// If the current primary is not available, select a new one.
for prefix, nodes := range allPrimaries {
log.Debug().
Caller().
Str(zf.Prefix, prefix.String()).
Uints64("availableNodes", func() []uint64 {
ids := make([]uint64, len(nodes))
for i, id := range nodes {
ids[i] = id.Uint64()
}
return ids
}()).
Msg("processing prefix for primary route selection")
if node, ok := pr.primaries[prefix]; ok {
// If the current primary is still available and healthy, continue.
if slices.Contains(nodes, node) && !pr.unhealthy.Contains(node) {
log.Debug().
Caller().
Str(zf.Prefix, prefix.String()).
Uint64("currentPrimary", node.Uint64()).
Msg("current primary still available and healthy, keeping it")
continue
} else {
log.Debug().
Caller().
Str(zf.Prefix, prefix.String()).
Uint64("oldPrimary", node.Uint64()).
Bool("unhealthy", pr.unhealthy.Contains(node)).
Msg("current primary no longer available or unhealthy")
}
}
// Select the first healthy node as primary.
// If all nodes are unhealthy, fall back to the first node
// (degraded service is better than no service).
var selected types.NodeID
found := false
for _, n := range nodes {
if !pr.unhealthy.Contains(n) {
selected = n
found = true
break
}
}
if !found && len(nodes) >= 1 {
selected = nodes[0]
found = true
log.Warn().
Caller().
Str(zf.Prefix, prefix.String()).
Uint64("fallbackPrimary", selected.Uint64()).
Msg("all nodes unhealthy for prefix, keeping first as degraded primary")
}
if found {
pr.primaries[prefix] = selected
changed = true
log.Debug().
Caller().
Str(zf.Prefix, prefix.String()).
Uint64("newPrimary", selected.Uint64()).
Msg("selected new primary for prefix")
}
}
// Clean up any remaining primaries that are no longer valid.
for prefix := range pr.primaries {
if _, ok := allPrimaries[prefix]; !ok {
log.Debug().
Caller().
Str(zf.Prefix, prefix.String()).
Msg("cleaning up primary route that no longer has available nodes")
delete(pr.primaries, prefix)
changed = true
}
}
// Populate the quick lookup index for primary routes
for _, nodeID := range pr.primaries {
pr.isPrimary[nodeID] = true
}
log.Debug().
Caller().
Bool(zf.Changes, changed).
Str(zf.FinalState, pr.stringLocked()).
Msg("updatePrimaryLocked completed")
return changed
}
// SetRoutes sets the routes for a given Node ID and recalculates the primary routes
// of the headscale.
// It returns true if there was a change in primary routes.
// All exit routes are ignored as they are not used in primary route context.
func (pr *PrimaryRoutes) SetRoutes(node types.NodeID, prefixes ...netip.Prefix) bool {
pr.mu.Lock()
defer pr.mu.Unlock()
nlog := log.With().Uint64(zf.NodeID, node.Uint64()).Logger()
nlog.Debug().
Caller().
Strs("prefixes", util.PrefixesToString(prefixes)).
Msg("PrimaryRoutes.SetRoutes called")
// If no routes are being set, remove the node from the routes map
// and clear any unhealthy state.
if len(prefixes) == 0 {
wasPresent := false
if _, ok := pr.routes[node]; ok {
delete(pr.routes, node)
pr.unhealthy.Delete(node)
wasPresent = true
nlog.Debug().
Caller().
Msg("removed node from primary routes (no prefixes)")
}
changed := pr.updatePrimaryLocked()
nlog.Debug().
Caller().
Bool("wasPresent", wasPresent).
Bool(zf.Changes, changed).
Str(zf.NewState, pr.stringLocked()).
Msg("SetRoutes completed (remove)")
return changed
}
rs := make(set.Set[netip.Prefix], len(prefixes))
for _, prefix := range prefixes {
if !tsaddr.IsExitRoute(prefix) {
rs.Add(prefix)
}
}
if rs.Len() != 0 {
pr.routes[node] = rs
nlog.Debug().
Caller().
Strs("routes", util.PrefixesToString(rs.Slice())).
Msg("updated node routes in primary route manager")
} else {
delete(pr.routes, node)
nlog.Debug().
Caller().
Msg("removed node from primary routes (only exit routes)")
}
changed := pr.updatePrimaryLocked()
nlog.Debug().
Caller().
Bool(zf.Changes, changed).
Str(zf.NewState, pr.stringLocked()).
Msg("SetRoutes completed (update)")
return changed
}
func (pr *PrimaryRoutes) PrimaryRoutes(id types.NodeID) []netip.Prefix {
if pr == nil {
return nil
}
pr.mu.Lock()
defer pr.mu.Unlock()
// Short circuit if the node is not a primary for any route.
if _, ok := pr.isPrimary[id]; !ok {
return nil
}
var routes []netip.Prefix
for prefix, node := range pr.primaries {
if node == id {
routes = append(routes, prefix)
}
}
slices.SortFunc(routes, netip.Prefix.Compare)
return routes
}
// SetNodeHealthy marks a node as healthy or unhealthy and recalculates
// primary routes. Returns true if primaries changed.
func (pr *PrimaryRoutes) SetNodeHealthy(
node types.NodeID,
healthy bool,
) bool {
pr.mu.Lock()
defer pr.mu.Unlock()
if healthy {
if !pr.unhealthy.Contains(node) {
return false
}
pr.unhealthy.Delete(node)
log.Debug().
Caller().
Uint64(zf.NodeID, node.Uint64()).
Msg("node marked healthy")
} else {
if pr.unhealthy.Contains(node) {
return false
}
pr.unhealthy.Add(node)
log.Info().
Caller().
Uint64(zf.NodeID, node.Uint64()).
Msg("node marked unhealthy")
}
return pr.updatePrimaryLocked()
}
// ClearUnhealthy removes a node from the unhealthy set without
// recalculating primaries. Used when a node reconnects to give
// it a fresh start.
func (pr *PrimaryRoutes) ClearUnhealthy(node types.NodeID) {
pr.mu.Lock()
defer pr.mu.Unlock()
if pr.unhealthy.Contains(node) {
pr.unhealthy.Delete(node)
log.Debug().
Caller().
Uint64(zf.NodeID, node.Uint64()).
Msg("cleared unhealthy state on reconnect")
}
}
// HANodes returns a snapshot of prefixes that have 2+ nodes
// advertising them (HA configurations), mapped to the node IDs.
// Node IDs are sorted deterministically.
func (pr *PrimaryRoutes) HANodes() map[netip.Prefix][]types.NodeID {
pr.mu.Lock()
defer pr.mu.Unlock()
// Build prefix → nodes map.
prefixNodes := make(map[netip.Prefix][]types.NodeID)
ids := types.NodeIDs(xmaps.Keys(pr.routes))
sort.Sort(ids)
for _, id := range ids {
routes := pr.routes[id]
for route := range routes {
prefixNodes[route] = append(prefixNodes[route], id)
}
}
// Keep only prefixes with 2+ advertisers.
result := make(map[netip.Prefix][]types.NodeID)
for prefix, nodes := range prefixNodes {
if len(nodes) >= 2 {
result[prefix] = nodes
}
}
return result
}
// IsNodeHealthy returns whether the node is currently considered healthy.
func (pr *PrimaryRoutes) IsNodeHealthy(node types.NodeID) bool {
pr.mu.Lock()
defer pr.mu.Unlock()
return !pr.unhealthy.Contains(node)
}
func (pr *PrimaryRoutes) String() string {
pr.mu.Lock()
defer pr.mu.Unlock()
return pr.stringLocked()
}
func (pr *PrimaryRoutes) stringLocked() string {
var sb strings.Builder
fmt.Fprintln(&sb, "Available routes:")
ids := types.NodeIDs(xmaps.Keys(pr.routes))
sort.Sort(ids)
for _, id := range ids {
prefixes := pr.routes[id]
fmt.Fprintf(&sb, "\nNode %d: %s", id, strings.Join(util.PrefixesToString(prefixes.Slice()), ", "))
}
fmt.Fprintln(&sb, "\n\nCurrent primary routes:")
for route, nodeID := range pr.primaries {
fmt.Fprintf(&sb, "\nRoute %s: %d", route, nodeID)
}
return sb.String()
}
// DebugRoutes represents the primary routes state in a structured format for JSON serialization.
type DebugRoutes struct {
// AvailableRoutes maps node IDs to their advertised routes
// In the context of primary routes, this represents the routes that are available
// for each node. A route will only be available if it is advertised by the node
// AND approved.
// Only routes by nodes currently connected to the headscale server are included.
AvailableRoutes map[types.NodeID][]netip.Prefix `json:"available_routes"`
// PrimaryRoutes maps route prefixes to the primary node serving them
PrimaryRoutes map[string]types.NodeID `json:"primary_routes"`
// UnhealthyNodes lists nodes that have failed health probes.
UnhealthyNodes []types.NodeID `json:"unhealthy_nodes,omitempty"`
}
// DebugJSON returns a structured representation of the primary routes state suitable for JSON serialization.
func (pr *PrimaryRoutes) DebugJSON() DebugRoutes {
pr.mu.Lock()
defer pr.mu.Unlock()
debug := DebugRoutes{
AvailableRoutes: make(map[types.NodeID][]netip.Prefix),
PrimaryRoutes: make(map[string]types.NodeID),
}
// Populate available routes
for nodeID, routes := range pr.routes {
prefixes := routes.Slice()
slices.SortFunc(prefixes, netip.Prefix.Compare)
debug.AvailableRoutes[nodeID] = prefixes
}
// Populate primary routes
for prefix, nodeID := range pr.primaries {
debug.PrimaryRoutes[prefix.String()] = nodeID
}
// Populate unhealthy nodes
if pr.unhealthy.Len() > 0 {
nodes := pr.unhealthy.Slice()
slices.SortFunc(nodes, func(a, b types.NodeID) int {
if a < b {
return -1
}
if a > b {
return 1
}
return 0
})
debug.UnhealthyNodes = nodes
}
return debug
}
hscontrol/routes/primary_property_test.go
-314
View File
@@ -1,314 +0,0 @@
package routes
import (
"fmt"
"net/netip"
"slices"
"testing"
"github.com/juanfont/headscale/hscontrol/types"
"pgregory.net/rapid"
)
// model mirrors the expected state of PrimaryRoutes given the operations
// applied so far. It is intentionally simple — close to what a careful
// reading of primary.go's invariants prescribes — so divergence between
// the model and the real PrimaryRoutes flags a bug in the algorithm.
type model struct {
// connected[n] is true if node n's routes are currently registered
// (SetRoutes with at least one prefix has been called and not since
// followed by SetRoutes with empty prefixes).
connected map[types.NodeID]bool
// prefixes[n] is the latest non-empty SetRoutes input for n. Cleared
// on Disconnect.
prefixes map[types.NodeID][]netip.Prefix
// unhealthy[n] mirrors PrimaryRoutes.unhealthy.
unhealthy map[types.NodeID]bool
// primary[p] is the current primary for prefix p. The implementation
// preserves the current primary across changes to avoid flapping, so
// the model has to track this across operations rather than recompute
// a fresh choice each time.
primary map[netip.Prefix]types.NodeID
}
func newModel() *model {
return &model{
connected: map[types.NodeID]bool{},
prefixes: map[types.NodeID][]netip.Prefix{},
unhealthy: map[types.NodeID]bool{},
primary: map[netip.Prefix]types.NodeID{},
}
}
// advertisersByPrefix returns the connected nodes that announce each
// prefix, sorted by NodeID (matches updatePrimaryLocked's iteration
// order).
func (m *model) advertisersByPrefix() map[netip.Prefix][]types.NodeID {
out := map[netip.Prefix][]types.NodeID{}
for n, prefs := range m.prefixes {
if !m.connected[n] {
continue
}
for _, p := range prefs {
out[p] = append(out[p], n)
}
}
for _, nodes := range out {
slices.Sort(nodes)
}
return out
}
// updatePrimaries reapplies the implementation's algorithm to recompute
// the primary for each prefix. Called after every operation.
func (m *model) updatePrimaries() {
advertisers := m.advertisersByPrefix()
// Step 1: drop primaries for prefixes that no longer have any
// advertiser at all.
for p := range m.primary {
if _, ok := advertisers[p]; !ok {
delete(m.primary, p)
}
}
// Step 2: for each prefix with advertisers, keep current primary
// if it is still in the advertiser set and not unhealthy. Otherwise
// select the first healthy advertiser; if all are unhealthy fall
// back to the lowest NodeID (degraded primary).
for p, nodes := range advertisers {
if cur, ok := m.primary[p]; ok {
if slices.Contains(nodes, cur) && !m.unhealthy[cur] {
continue
}
}
var selected types.NodeID
var found bool
for _, n := range nodes {
if !m.unhealthy[n] {
selected = n
found = true
break
}
}
if !found && len(nodes) >= 1 {
selected = nodes[0]
found = true
}
if found {
m.primary[p] = selected
}
}
}
// allPrefixes returns every prefix mentioned by any connected node.
func (m *model) allPrefixes() []netip.Prefix {
seen := map[netip.Prefix]bool{}
for n, prefs := range m.prefixes {
if !m.connected[n] {
continue
}
for _, p := range prefs {
seen[p] = true
}
}
out := make([]netip.Prefix, 0, len(seen))
for p := range seen {
out = append(out, p)
}
return out
}
func samePrefixSet(a, b []netip.Prefix) bool {
if len(a) != len(b) {
return false
}
aa := slices.Clone(a)
bb := slices.Clone(b)
slices.SortFunc(aa, netip.Prefix.Compare)
slices.SortFunc(bb, netip.Prefix.Compare)
return slices.Equal(aa, bb)
}
// checkInvariants asserts every property we expect of the real
// PrimaryRoutes given the model.
func checkInvariants(rt *rapid.T, pr *PrimaryRoutes, m *model, nodeIDs []types.NodeID) {
rt.Helper()
// Per-node primary set must match the model's expectations.
expectedByNode := map[types.NodeID][]netip.Prefix{}
for p, owner := range m.primary {
expectedByNode[owner] = append(expectedByNode[owner], p)
}
for _, id := range nodeIDs {
got := pr.PrimaryRoutes(id)
want := expectedByNode[id]
if !samePrefixSet(got, want) {
rt.Fatalf(
"PrimaryRoutes(%d) = %v, model expected %v\nstate: %s",
id, got, want, pr.String(),
)
}
if want := !m.unhealthy[id]; pr.IsNodeHealthy(id) != want {
rt.Fatalf(
"IsNodeHealthy(%d) = %v, want %v",
id, pr.IsNodeHealthy(id), want,
)
}
}
// Every prefix that has at least one connected advertiser must have a
// primary in the real PrimaryRoutes. Issue #3203 manifests as a prefix
// silently losing its primary after a disconnect/reconnect cycle.
for _, p := range m.allPrefixes() {
want, expectExists := m.primary[p]
if !expectExists {
continue
}
var foundOwner types.NodeID
var found bool
for _, id := range nodeIDs {
for _, got := range pr.PrimaryRoutes(id) {
if got == p {
foundOwner = id
found = true
break
}
}
if found {
break
}
}
if !found {
rt.Fatalf(
"prefix %s has at least one advertiser in the model but no primary in PrimaryRoutes\nstate: %s",
p, pr.String(),
)
}
if want != foundOwner {
rt.Fatalf(
"prefix %s: PrimaryRoutes assigned to node %d, model expected %d\nstate: %s",
p, foundOwner, want, pr.String(),
)
}
}
}
// TestPrimaryRoutesProperty drives PrimaryRoutes with a randomised
// sequence of high-level operations (ConnectAdvertise, Disconnect,
// ProbeUnhealthy, ProbeHealthy, ApprovedRoutesChange) and checks that
// the per-prefix primary chosen by the implementation matches a
// reference model after every step.
//
// Background: issue #3203 reports that HA tracking enters a stuck
// state after a sequence of disconnect/reconnect events. The narrow
// integration and servertest reproductions written for the bug do not
// fail on upstream/main, so this property test broadens the search by
// letting rapid generate sequences we have not enumerated by hand.
func TestPrimaryRoutesProperty(t *testing.T) {
rapid.Check(t, func(rt *rapid.T) {
// Small fixed pools so rapid can collide on the same nodes
// and prefixes — that is where the interesting interleavings
// live.
const numNodes = 4
nodeIDs := make([]types.NodeID, 0, numNodes)
for i := 1; i <= numNodes; i++ {
nodeIDs = append(nodeIDs, types.NodeID(i))
}
prefixes := []netip.Prefix{
netip.MustParsePrefix("10.0.0.0/24"),
netip.MustParsePrefix("10.0.1.0/24"),
}
pr := New()
m := newModel()
nodeGen := rapid.SampledFrom(nodeIDs)
prefixSubsetGen := rapid.SliceOfNDistinct(
rapid.SampledFrom(prefixes),
0, len(prefixes),
func(p netip.Prefix) string { return p.String() },
)
opCount := rapid.IntRange(5, 60).Draw(rt, "opCount")
for step := 0; step < opCount; step++ {
op := rapid.IntRange(0, 4).Draw(rt, fmt.Sprintf("op_%d", step))
id := nodeGen.Draw(rt, fmt.Sprintf("id_%d", step))
switch op {
case 0: // ConnectAdvertise
prefs := prefixSubsetGen.Draw(rt, fmt.Sprintf("prefs_%d", step))
// Mirror state.Connect: ClearUnhealthy then SetRoutes.
pr.ClearUnhealthy(id)
delete(m.unhealthy, id)
if len(prefs) == 0 {
pr.SetRoutes(id)
delete(m.connected, id)
delete(m.prefixes, id)
} else {
pr.SetRoutes(id, prefs...)
m.connected[id] = true
m.prefixes[id] = prefs
}
case 1: // Disconnect (state.Disconnect path).
wasConnected := m.connected[id]
pr.SetRoutes(id)
delete(m.connected, id)
delete(m.prefixes, id)
// SetRoutes(empty) only clears unhealthy if the node
// was previously in the routes map. See primary.go.
if wasConnected {
delete(m.unhealthy, id)
}
case 2: // ProbeUnhealthy
pr.SetNodeHealthy(id, false)
m.unhealthy[id] = true
case 3: // ProbeHealthy
pr.SetNodeHealthy(id, true)
delete(m.unhealthy, id)
case 4: // ApprovedRoutesChange (no ClearUnhealthy)
wasConnected := m.connected[id]
prefs := prefixSubsetGen.Draw(rt, fmt.Sprintf("prefs_%d", step))
if len(prefs) == 0 {
pr.SetRoutes(id)
delete(m.connected, id)
delete(m.prefixes, id)
if wasConnected {
delete(m.unhealthy, id)
}
} else {
pr.SetRoutes(id, prefs...)
m.connected[id] = true
m.prefixes[id] = prefs
}
}
m.updatePrimaries()
checkInvariants(rt, pr, m, nodeIDs)
}
})
}
hscontrol/routes/primary_test.go
-717
View File
@@ -1,717 +0,0 @@
package routes
import (
"net/netip"
"sync"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/juanfont/headscale/hscontrol/types"
"github.com/juanfont/headscale/hscontrol/util"
"tailscale.com/util/set"
)
// mp is a helper function that wraps netip.MustParsePrefix.
func mp(prefix string) netip.Prefix {
return netip.MustParsePrefix(prefix)
}
func TestPrimaryRoutes(t *testing.T) {
tests := []struct {
name string
operations func(pr *PrimaryRoutes) bool
expectedRoutes map[types.NodeID]set.Set[netip.Prefix]
expectedPrimaries map[netip.Prefix]types.NodeID
expectedIsPrimary map[types.NodeID]bool
expectedUnhealthy set.Set[types.NodeID]
expectedChange bool
// primaries is a map of prefixes to the node that is the primary for that prefix.
primaries map[netip.Prefix]types.NodeID
isPrimary map[types.NodeID]bool
}{
{
name: "single-node-registers-single-route",
operations: func(pr *PrimaryRoutes) bool {
return pr.SetRoutes(1, mp("192.168.1.0/24"))
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedChange: true,
},
{
name: "multiple-nodes-register-different-routes",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
return pr.SetRoutes(2, mp("192.168.2.0/24"))
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.2.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
mp("192.168.2.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
2: true,
},
expectedChange: true,
},
{
name: "multiple-nodes-register-overlapping-routes",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // true
return pr.SetRoutes(2, mp("192.168.1.0/24")) // false
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedChange: false,
},
{
name: "node-deregisters-a-route",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
return pr.SetRoutes(1) // Deregister by setting no routes
},
expectedRoutes: nil,
expectedPrimaries: nil,
expectedIsPrimary: nil,
expectedChange: true,
},
{
name: "node-deregisters-one-of-multiple-routes",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"), mp("192.168.2.0/24"))
return pr.SetRoutes(1, mp("192.168.2.0/24")) // Deregister one route by setting the remaining route
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.2.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.2.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedChange: true,
},
{
name: "node-registers-and-deregisters-routes-in-sequence",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.2.0/24"))
pr.SetRoutes(1) // Deregister by setting no routes
return pr.SetRoutes(1, mp("192.168.3.0/24"))
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.3.0/24"): {},
},
2: {
mp("192.168.2.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.2.0/24"): 2,
mp("192.168.3.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
2: true,
},
expectedChange: true,
},
{
name: "multiple-nodes-register-same-route",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true
return pr.SetRoutes(3, mp("192.168.1.0/24")) // false
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.1.0/24"): {},
},
3: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedChange: false,
},
{
name: "register-multiple-routes-shift-primary-check-primary",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 1 primary
pr.SetRoutes(3, mp("192.168.1.0/24")) // false, 1 primary
return pr.SetRoutes(1) // true, 2 primary
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
2: {
mp("192.168.1.0/24"): {},
},
3: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
2: true,
},
expectedChange: true,
},
{
name: "primary-route-map-is-cleared-up-no-primary",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 1 primary
pr.SetRoutes(3, mp("192.168.1.0/24")) // false, 1 primary
pr.SetRoutes(1) // true, 2 primary
return pr.SetRoutes(2) // true, no primary
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
3: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 3,
},
expectedIsPrimary: map[types.NodeID]bool{
3: true,
},
expectedChange: true,
},
{
name: "primary-route-map-is-cleared-up-all-no-primary",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 1 primary
pr.SetRoutes(3, mp("192.168.1.0/24")) // false, 1 primary
pr.SetRoutes(1) // true, 2 primary
pr.SetRoutes(2) // true, no primary
return pr.SetRoutes(3) // false, no primary
},
expectedChange: true,
},
{
name: "primary-route-map-is-cleared-up",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 1 primary
pr.SetRoutes(3, mp("192.168.1.0/24")) // false, 1 primary
pr.SetRoutes(1) // true, 2 primary
return pr.SetRoutes(2) // true, no primary
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
3: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 3,
},
expectedIsPrimary: map[types.NodeID]bool{
3: true,
},
expectedChange: true,
},
{
name: "primary-route-no-flake",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 1 primary
pr.SetRoutes(3, mp("192.168.1.0/24")) // false, 1 primary
pr.SetRoutes(1) // true, 2 primary
return pr.SetRoutes(1, mp("192.168.1.0/24")) // false, 2 primary
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.1.0/24"): {},
},
3: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
2: true,
},
expectedChange: false,
},
{
name: "primary-route-no-flake-check-old-primary",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 1 primary
pr.SetRoutes(3, mp("192.168.1.0/24")) // false, 1 primary
pr.SetRoutes(1) // true, 2 primary
return pr.SetRoutes(1, mp("192.168.1.0/24")) // false, 2 primary
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.1.0/24"): {},
},
3: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
2: true,
},
expectedChange: false,
},
{
name: "primary-route-no-flake-full-integration",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24")) // false
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 1 primary
pr.SetRoutes(3, mp("192.168.1.0/24")) // false, 1 primary
pr.SetRoutes(1) // true, 2 primary
pr.SetRoutes(2) // true, 3 primary
pr.SetRoutes(1, mp("192.168.1.0/24")) // true, 3 primary
pr.SetRoutes(2, mp("192.168.1.0/24")) // true, 3 primary
pr.SetRoutes(1) // true, 3 primary
return pr.SetRoutes(1, mp("192.168.1.0/24")) // false, 3 primary
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.1.0/24"): {},
},
3: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 3,
},
expectedIsPrimary: map[types.NodeID]bool{
3: true,
},
expectedChange: false,
},
{
name: "multiple-nodes-register-same-route-and-exit",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("0.0.0.0/0"), mp("192.168.1.0/24"))
return pr.SetRoutes(2, mp("192.168.1.0/24"))
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.1.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedChange: false,
},
{
name: "deregister-non-existent-route",
operations: func(pr *PrimaryRoutes) bool {
return pr.SetRoutes(1) // Deregister by setting no routes
},
expectedRoutes: nil,
expectedChange: false,
},
{
name: "register-empty-prefix-list",
operations: func(pr *PrimaryRoutes) bool {
return pr.SetRoutes(1)
},
expectedRoutes: nil,
expectedChange: false,
},
{
name: "exit-nodes",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("10.0.0.0/16"), mp("0.0.0.0/0"), mp("::/0"))
pr.SetRoutes(3, mp("0.0.0.0/0"), mp("::/0"))
return pr.SetRoutes(2, mp("0.0.0.0/0"), mp("::/0"))
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("10.0.0.0/16"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("10.0.0.0/16"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedChange: false,
},
{
name: "unhealthy-primary-fails-over",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
return pr.SetNodeHealthy(1, false)
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {mp("192.168.1.0/24"): {}},
2: {mp("192.168.1.0/24"): {}},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
2: true,
},
expectedUnhealthy: set.Set[types.NodeID]{1: {}},
expectedChange: true,
},
{
name: "unhealthy-non-primary-no-change",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
return pr.SetNodeHealthy(2, false)
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {mp("192.168.1.0/24"): {}},
2: {mp("192.168.1.0/24"): {}},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedUnhealthy: set.Set[types.NodeID]{2: {}},
expectedChange: false,
},
{
name: "all-unhealthy-keeps-current-primary",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
pr.SetNodeHealthy(1, false) // failover to 2
return pr.SetNodeHealthy(2, false) // both unhealthy, falls back to first sorted (1)
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {mp("192.168.1.0/24"): {}},
2: {mp("192.168.1.0/24"): {}},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
},
expectedUnhealthy: set.Set[types.NodeID]{1: {}, 2: {}},
expectedChange: true,
},
{
name: "recovery-marks-healthy-no-flap",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
pr.SetRoutes(3, mp("192.168.1.0/24"))
pr.SetNodeHealthy(1, false) // failover to 2
return pr.SetNodeHealthy(1, true) // recovered, but 2 stays primary
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {mp("192.168.1.0/24"): {}},
2: {mp("192.168.1.0/24"): {}},
3: {mp("192.168.1.0/24"): {}},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
2: true,
},
expectedChange: false,
},
{
name: "clear-unhealthy-on-reconnect",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
pr.SetNodeHealthy(1, false) // failover to 2
pr.ClearUnhealthy(1) // reconnect clears unhealthy
// 2 stays primary (stability), but 1 is healthy again
return pr.SetNodeHealthy(1, false) // can be marked unhealthy again
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {mp("192.168.1.0/24"): {}},
2: {mp("192.168.1.0/24"): {}},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
2: true,
},
expectedUnhealthy: set.Set[types.NodeID]{1: {}},
expectedChange: false,
},
{
name: "unhealthy-then-deregister",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
pr.SetNodeHealthy(1, false)
return pr.SetRoutes(1) // deregister clears routes and unhealthy
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
2: {mp("192.168.1.0/24"): {}},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
2: true,
},
expectedChange: false,
},
{
name: "unhealthy-primary-three-nodes-cascade",
operations: func(pr *PrimaryRoutes) bool {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
pr.SetRoutes(3, mp("192.168.1.0/24"))
pr.SetNodeHealthy(1, false) // failover to 2
return pr.SetNodeHealthy(2, false) // failover to 3
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {mp("192.168.1.0/24"): {}},
2: {mp("192.168.1.0/24"): {}},
3: {mp("192.168.1.0/24"): {}},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 3,
},
expectedIsPrimary: map[types.NodeID]bool{
3: true,
},
expectedUnhealthy: set.Set[types.NodeID]{1: {}, 2: {}},
expectedChange: true,
},
{
name: "concurrent-access",
operations: func(pr *PrimaryRoutes) bool {
var wg sync.WaitGroup
wg.Add(2)
var change1, change2 bool
go func() {
defer wg.Done()
change1 = pr.SetRoutes(1, mp("192.168.1.0/24"))
}()
go func() {
defer wg.Done()
change2 = pr.SetRoutes(2, mp("192.168.2.0/24"))
}()
wg.Wait()
return change1 || change2
},
expectedRoutes: map[types.NodeID]set.Set[netip.Prefix]{
1: {
mp("192.168.1.0/24"): {},
},
2: {
mp("192.168.2.0/24"): {},
},
},
expectedPrimaries: map[netip.Prefix]types.NodeID{
mp("192.168.1.0/24"): 1,
mp("192.168.2.0/24"): 2,
},
expectedIsPrimary: map[types.NodeID]bool{
1: true,
2: true,
},
expectedChange: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
pr := New()
change := tt.operations(pr)
if change != tt.expectedChange {
t.Errorf("change = %v, want %v", change, tt.expectedChange)
}
comps := append(util.Comparers, cmpopts.EquateEmpty())
if diff := cmp.Diff(tt.expectedRoutes, pr.routes, comps...); diff != "" {
t.Errorf("routes mismatch (-want +got):\n%s", diff)
}
if diff := cmp.Diff(tt.expectedPrimaries, pr.primaries, comps...); diff != "" {
t.Errorf("primaries mismatch (-want +got):\n%s", diff)
}
if diff := cmp.Diff(tt.expectedIsPrimary, pr.isPrimary, comps...); diff != "" {
t.Errorf("isPrimary mismatch (-want +got):\n%s", diff)
}
if diff := cmp.Diff(tt.expectedUnhealthy, pr.unhealthy, comps...); diff != "" {
t.Errorf("unhealthy mismatch (-want +got):\n%s", diff)
}
})
}
}
func TestHANodes(t *testing.T) {
tests := []struct {
name string
setup func(pr *PrimaryRoutes)
expected map[netip.Prefix][]types.NodeID
}{
{
name: "single-node-not-ha",
setup: func(pr *PrimaryRoutes) {
pr.SetRoutes(1, mp("192.168.1.0/24"))
},
expected: nil,
},
{
name: "two-nodes-same-prefix-is-ha",
setup: func(pr *PrimaryRoutes) {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
},
expected: map[netip.Prefix][]types.NodeID{
mp("192.168.1.0/24"): {1, 2},
},
},
{
name: "two-nodes-different-prefixes-not-ha",
setup: func(pr *PrimaryRoutes) {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.2.0/24"))
},
expected: nil,
},
{
name: "three-nodes-two-share-prefix",
setup: func(pr *PrimaryRoutes) {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
pr.SetRoutes(3, mp("10.0.0.0/8"))
},
expected: map[netip.Prefix][]types.NodeID{
mp("192.168.1.0/24"): {1, 2},
},
},
{
name: "three-nodes-all-share",
setup: func(pr *PrimaryRoutes) {
pr.SetRoutes(1, mp("192.168.1.0/24"))
pr.SetRoutes(2, mp("192.168.1.0/24"))
pr.SetRoutes(3, mp("192.168.1.0/24"))
},
expected: map[netip.Prefix][]types.NodeID{
mp("192.168.1.0/24"): {1, 2, 3},
},
},
{
name: "empty",
setup: func(_ *PrimaryRoutes) {
},
expected: nil,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
pr := New()
tt.setup(pr)
got := pr.HANodes()
comps := append(util.Comparers, cmpopts.EquateEmpty())
if diff := cmp.Diff(tt.expected, got, comps...); diff != "" {
t.Errorf("HANodes mismatch (-want +got):\n%s", diff)
}
})
}
}
hscontrol/state/primaries_property_test.go
+331
View File
@@ -0,0 +1,331 @@
package state
import (
"fmt"
"net/netip"
"slices"
"testing"
"github.com/juanfont/headscale/hscontrol/types"
"pgregory.net/rapid"
"tailscale.com/tailcfg"
"tailscale.com/types/key"
)
// model mirrors the expected primary-route assignment given the
// operations applied so far. It is intentionally simple — close to
// what computePrimaries (node_store.go) prescribes — so divergence
// between the model and the snapshot's primaries map flags a bug in
// the algorithm.
//
// The model lives in the test file because the algorithm under test
// is the one inside snapshotFromNodes; the model's job is to predict
// the same answer from a separate, deliberately direct implementation.
type primariesModel struct {
connected map[types.NodeID]bool
prefixes map[types.NodeID][]netip.Prefix
unhealthy map[types.NodeID]bool
// primary[p] is the current primary for prefix p. The
// implementation preserves the current primary across changes to
// avoid flapping, so the model has to track this across
// operations rather than recompute a fresh choice each time.
primary map[netip.Prefix]types.NodeID
}
func newPrimariesModel() *primariesModel {
return &primariesModel{
connected: map[types.NodeID]bool{},
prefixes: map[types.NodeID][]netip.Prefix{},
unhealthy: map[types.NodeID]bool{},
primary: map[netip.Prefix]types.NodeID{},
}
}
// advertisersByPrefix returns the connected nodes that announce each
// prefix, sorted by NodeID (matches computePrimaries' iteration).
func (m *primariesModel) advertisersByPrefix() map[netip.Prefix][]types.NodeID {
out := map[netip.Prefix][]types.NodeID{}
for n, prefs := range m.prefixes {
if !m.connected[n] {
continue
}
for _, p := range prefs {
out[p] = append(out[p], n)
}
}
for _, nodes := range out {
slices.Sort(nodes)
}
return out
}
// updatePrimaries reapplies the algorithm to recompute the primary
// for each prefix. Called after every operation.
func (m *primariesModel) updatePrimaries() {
advertisers := m.advertisersByPrefix()
// Drop primaries for prefixes that no longer have any advertiser.
for p := range m.primary {
if _, ok := advertisers[p]; !ok {
delete(m.primary, p)
}
}
for p, nodes := range advertisers {
if cur, ok := m.primary[p]; ok {
if slices.Contains(nodes, cur) && !m.unhealthy[cur] {
continue
}
}
var (
selected types.NodeID
found bool
)
for _, n := range nodes {
if !m.unhealthy[n] {
selected = n
found = true
break
}
}
if !found && len(nodes) >= 1 {
selected = nodes[0]
found = true
}
if found {
m.primary[p] = selected
}
}
}
// allPrefixes returns every prefix mentioned by any connected node.
func (m *primariesModel) allPrefixes() []netip.Prefix {
seen := map[netip.Prefix]bool{}
for n, prefs := range m.prefixes {
if !m.connected[n] {
continue
}
for _, p := range prefs {
seen[p] = true
}
}
out := make([]netip.Prefix, 0, len(seen))
for p := range seen {
out = append(out, p)
}
return out
}
func samePrefixSet(a, b []netip.Prefix) bool {
if len(a) != len(b) {
return false
}
aa := slices.Clone(a)
bb := slices.Clone(b)
slices.SortFunc(aa, netip.Prefix.Compare)
slices.SortFunc(bb, netip.Prefix.Compare)
return slices.Equal(aa, bb)
}
// checkInvariants asserts every property we expect of the snapshot's
// primaries map given the model.
func checkPrimariesInvariants(rt *rapid.T, ns *NodeStore, m *primariesModel, nodeIDs []types.NodeID) {
rt.Helper()
expectedByNode := map[types.NodeID][]netip.Prefix{}
for p, owner := range m.primary {
expectedByNode[owner] = append(expectedByNode[owner], p)
}
for _, id := range nodeIDs {
got := ns.PrimaryRoutesForNode(id)
want := expectedByNode[id]
if !samePrefixSet(got, want) {
rt.Fatalf(
"PrimaryRoutesForNode(%d) = %v, model expected %v",
id, got, want,
)
}
if want := !m.unhealthy[id]; ns.IsNodeHealthy(id) != want {
rt.Fatalf(
"IsNodeHealthy(%d) = %v, want %v",
id, ns.IsNodeHealthy(id), want,
)
}
}
// Every prefix that has at least one connected advertiser must
// have a primary in the snapshot. Issue #3203 manifests as a
// prefix silently losing its primary after a disconnect/reconnect
// cycle.
for _, p := range m.allPrefixes() {
want, expectExists := m.primary[p]
if !expectExists {
continue
}
got, ok := ns.PrimaryRouteFor(p)
if !ok {
rt.Fatalf(
"prefix %s has at least one advertiser in the model but no primary in NodeStore",
p,
)
}
if want != got {
rt.Fatalf(
"prefix %s: snapshot primary = %d, model expected %d",
p, got, want,
)
}
}
}
// nodeForRapid builds a minimal types.Node for use in property
// tests. Tests drive (IsOnline, Hostinfo.RoutableIPs, ApprovedRoutes,
// Unhealthy) via UpdateNode; the rest stays fixed.
func nodeForRapid(id types.NodeID) types.Node {
mk := key.NewMachine()
nk := key.NewNode()
return types.Node{
ID: id,
Hostname: fmt.Sprintf("rapid-%d", id),
MachineKey: mk.Public(),
NodeKey: nk.Public(),
UserID: new(uint(1)),
User: &types.User{Name: "rapid"},
IsOnline: new(false),
Hostinfo: &tailcfg.Hostinfo{},
}
}
// TestPrimaryRoutesProperty drives NodeStore with a randomised
// sequence of high-level operations and checks that the snapshot's
// primaries map matches a reference model after every step.
//
// Background: issue #3203 reports that HA tracking enters a stuck
// state after a sequence of disconnect/reconnect events. The narrow
// integration and servertest reproductions written for the bug do
// not fail on upstream/main, so this property test broadens the
// search by letting rapid generate sequences we have not enumerated
// by hand.
func TestPrimaryRoutesProperty(t *testing.T) {
rapid.Check(t, func(rt *rapid.T) {
const numNodes = 4
nodeIDs := make([]types.NodeID, 0, numNodes)
for i := 1; i <= numNodes; i++ {
nodeIDs = append(nodeIDs, types.NodeID(i))
}
prefixes := []netip.Prefix{
netip.MustParsePrefix("10.0.0.0/24"),
netip.MustParsePrefix("10.0.1.0/24"),
}
ns := NewNodeStore(nil, allowAllPeersFunc, TestBatchSize, TestBatchTimeout)
ns.Start()
defer ns.Stop()
for _, id := range nodeIDs {
ns.PutNode(nodeForRapid(id))
}
m := newPrimariesModel()
nodeGen := rapid.SampledFrom(nodeIDs)
prefixSubsetGen := rapid.SliceOfNDistinct(
rapid.SampledFrom(prefixes),
0, len(prefixes),
func(p netip.Prefix) string { return p.String() },
)
opCount := rapid.IntRange(5, 60).Draw(rt, "opCount")
for step := range opCount {
op := rapid.IntRange(0, 4).Draw(rt, fmt.Sprintf("op_%d", step))
id := nodeGen.Draw(rt, fmt.Sprintf("id_%d", step))
switch op {
case 0: // ConnectAdvertise — Connect path clears Unhealthy.
prefs := prefixSubsetGen.Draw(rt, fmt.Sprintf("prefs_%d", step))
ns.UpdateNode(id, func(n *types.Node) {
n.IsOnline = new(true)
n.Unhealthy = false
n.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: prefs}
n.ApprovedRoutes = prefs
})
if len(prefs) == 0 {
delete(m.connected, id)
delete(m.prefixes, id)
} else {
m.connected[id] = true
m.prefixes[id] = prefs
}
delete(m.unhealthy, id)
case 1: // Disconnect — IsOnline=false; ApprovedRoutes persists.
ns.UpdateNode(id, func(n *types.Node) {
n.IsOnline = new(false)
})
delete(m.connected, id)
case 2: // ProbeUnhealthy — HA prober marks node bad.
ns.UpdateNode(id, func(n *types.Node) {
n.Unhealthy = true
})
m.unhealthy[id] = true
case 3: // ProbeHealthy — HA prober marks node good.
ns.UpdateNode(id, func(n *types.Node) {
n.Unhealthy = false
})
delete(m.unhealthy, id)
case 4: // ApprovedRoutesChange — change advertised prefs without touching health.
prefs := prefixSubsetGen.Draw(rt, fmt.Sprintf("prefs_%d", step))
ns.UpdateNode(id, func(n *types.Node) {
n.IsOnline = new(true)
n.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: prefs}
n.ApprovedRoutes = prefs
})
if len(prefs) == 0 {
delete(m.connected, id)
delete(m.prefixes, id)
} else {
m.connected[id] = true
m.prefixes[id] = prefs
}
}
m.updatePrimaries()
checkPrimariesInvariants(rt, ns, m, nodeIDs)
}
})
}
hscontrol/state/primaries_test.go
+319
View File
@@ -0,0 +1,319 @@
package state
import (
"net/netip"
"slices"
"testing"
"github.com/juanfont/headscale/hscontrol/types"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"tailscale.com/tailcfg"
)
// mp wraps netip.MustParsePrefix.
func mp(prefix string) netip.Prefix {
return netip.MustParsePrefix(prefix)
}
// primariesFixture builds a NodeStore with the requested node IDs
// pre-registered (offline, no routes) and provides terse helpers for
// driving the kinds of state transitions the algorithm cares about.
type primariesFixture struct {
t *testing.T
ns *NodeStore
}
func newPrimariesFixture(t *testing.T, ids ...types.NodeID) *primariesFixture {
t.Helper()
ns := NewNodeStore(nil, allowAllPeersFunc, TestBatchSize, TestBatchTimeout)
ns.Start()
t.Cleanup(ns.Stop)
for _, id := range ids {
ns.PutNode(nodeForRapid(id))
}
return &primariesFixture{t: t, ns: ns}
}
// advertise mirrors State.Connect: marks the node online, clears
// Unhealthy, and sets approved + announced routes to prefs. An empty
// prefs argument leaves the node online but advertising nothing.
func (f *primariesFixture) advertise(id types.NodeID, prefs ...netip.Prefix) {
f.t.Helper()
f.ns.UpdateNode(id, func(n *types.Node) {
n.IsOnline = new(true)
n.Unhealthy = false
n.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: prefs}
n.ApprovedRoutes = prefs
})
}
// approveRoutes mirrors State.SetApprovedRoutes / Hostinfo updates:
// it changes the node's announced + approved set without touching
// Unhealthy.
func (f *primariesFixture) approveRoutes(id types.NodeID, prefs ...netip.Prefix) {
f.t.Helper()
f.ns.UpdateNode(id, func(n *types.Node) {
n.IsOnline = new(true)
n.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: prefs}
n.ApprovedRoutes = prefs
})
}
// disconnect mirrors State.Disconnect: marks the node offline. The
// snapshot rebuild treats an offline node as a non-advertiser.
func (f *primariesFixture) disconnect(id types.NodeID) {
f.t.Helper()
f.ns.UpdateNode(id, func(n *types.Node) {
n.IsOnline = new(false)
})
}
// unhealthy mirrors State.SetNodeUnhealthy(id, true).
func (f *primariesFixture) unhealthy(id types.NodeID) {
f.t.Helper()
f.ns.UpdateNode(id, func(n *types.Node) {
n.Unhealthy = true
})
}
// healthy mirrors State.SetNodeUnhealthy(id, false).
func (f *primariesFixture) healthy(id types.NodeID) {
f.t.Helper()
f.ns.UpdateNode(id, func(n *types.Node) {
n.Unhealthy = false
})
}
// requirePrimary asserts that prefix has node id as its primary.
func (f *primariesFixture) requirePrimary(prefix netip.Prefix, id types.NodeID) {
f.t.Helper()
got, ok := f.ns.PrimaryRouteFor(prefix)
require.True(f.t, ok, "expected a primary for %s, got none", prefix)
require.Equal(f.t, id, got, "primary for %s", prefix)
}
// requireNoPrimary asserts that prefix has no primary at all.
func (f *primariesFixture) requireNoPrimary(prefix netip.Prefix) {
f.t.Helper()
_, ok := f.ns.PrimaryRouteFor(prefix)
require.False(f.t, ok, "expected no primary for %s", prefix)
}
// requireNodeRoutes asserts the set of prefixes for which id is the
// primary, regardless of order.
func (f *primariesFixture) requireNodeRoutes(id types.NodeID, want ...netip.Prefix) {
f.t.Helper()
got := f.ns.PrimaryRoutesForNode(id)
gotSorted := slices.Clone(got)
wantSorted := slices.Clone(want)
slices.SortFunc(gotSorted, netip.Prefix.Compare)
slices.SortFunc(wantSorted, netip.Prefix.Compare)
require.Equal(f.t, wantSorted, gotSorted, "primary routes for node %d", id)
}
func TestPrimaries_SingleNodeSingleRoute(t *testing.T) {
f := newPrimariesFixture(t, 1)
f.advertise(1, mp("192.168.1.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 1)
f.requireNodeRoutes(1, mp("192.168.1.0/24"))
}
func TestPrimaries_TwoNodesDifferentRoutes(t *testing.T) {
f := newPrimariesFixture(t, 1, 2)
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.2.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 1)
f.requirePrimary(mp("192.168.2.0/24"), 2)
}
func TestPrimaries_OverlappingRoutesLowerIDWins(t *testing.T) {
f := newPrimariesFixture(t, 1, 2)
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 1)
f.requireNodeRoutes(1, mp("192.168.1.0/24"))
f.requireNodeRoutes(2)
}
func TestPrimaries_AntiFlapPreservesCurrentPrimary(t *testing.T) {
// A primary that disappears (advertiser leaves the set) should
// trigger failover. When the original primary returns, the new
// primary keeps the assignment — anti-flap.
f := newPrimariesFixture(t, 1, 2)
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 1)
f.disconnect(1)
f.requirePrimary(mp("192.168.1.0/24"), 2)
f.advertise(1, mp("192.168.1.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 2)
}
func TestPrimaries_ClearRoutesDropsPrimary(t *testing.T) {
f := newPrimariesFixture(t, 1)
f.advertise(1, mp("192.168.1.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 1)
f.approveRoutes(1)
f.requireNoPrimary(mp("192.168.1.0/24"))
}
func TestPrimaries_DisconnectDropsLastAdvertiserPrimary(t *testing.T) {
f := newPrimariesFixture(t, 1)
f.advertise(1, mp("192.168.1.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 1)
f.disconnect(1)
f.requireNoPrimary(mp("192.168.1.0/24"))
}
func TestPrimaries_UnhealthyTriggersFailover(t *testing.T) {
f := newPrimariesFixture(t, 1, 2)
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
f.requirePrimary(mp("192.168.1.0/24"), 1)
f.unhealthy(1)
f.requirePrimary(mp("192.168.1.0/24"), 2)
}
func TestPrimaries_RecoveryFromUnhealthyNoFlap(t *testing.T) {
f := newPrimariesFixture(t, 1, 2)
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
f.unhealthy(1)
f.requirePrimary(mp("192.168.1.0/24"), 2)
f.healthy(1)
f.requirePrimary(mp("192.168.1.0/24"), 2)
}
func TestPrimaries_AllUnhealthyKeepsLowestIDPrimary(t *testing.T) {
// When every advertiser is unhealthy the algorithm degrades to
// the lowest-ID advertiser rather than going dark — peers should
// still see *some* primary so connectivity can recover when one
// of them flips healthy.
f := newPrimariesFixture(t, 1, 2)
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
f.unhealthy(1)
f.unhealthy(2)
f.requirePrimary(mp("192.168.1.0/24"), 1)
}
func TestPrimaries_ExitRouteNotElected(t *testing.T) {
// Exit routes (0.0.0.0/0, ::/0) are not subject to HA primary
// election — every approved exit-route advertiser keeps it.
f := newPrimariesFixture(t, 1)
exitV4 := mp("0.0.0.0/0")
f.advertise(1, exitV4)
f.requireNoPrimary(exitV4)
}
func TestPrimaries_RegressionIssue3203_BothOfflineThenOneReturns(t *testing.T) {
// Issue #3203: with two HA advertisers, dropping both then
// bringing one back used to leave the prefix without any
// primary. After the refactor the snapshot recomputes primaries
// on every NodeStore write, so the returning advertiser must
// be elected.
prefix := mp("10.0.0.0/24")
f := newPrimariesFixture(t, 1, 2)
f.advertise(1, prefix)
f.advertise(2, prefix)
f.requirePrimary(prefix, 1)
f.disconnect(1)
f.requirePrimary(prefix, 2)
f.disconnect(2)
f.requireNoPrimary(prefix)
f.advertise(2, prefix)
f.requirePrimary(prefix, 2)
}
func TestPrimaries_HANodes(t *testing.T) {
tests := []struct {
name string
setup func(*primariesFixture)
want map[netip.Prefix][]types.NodeID
}{
{
name: "single-node-not-ha",
setup: func(f *primariesFixture) {
f.advertise(1, mp("192.168.1.0/24"))
},
want: map[netip.Prefix][]types.NodeID{},
},
{
name: "two-nodes-same-prefix-is-ha",
setup: func(f *primariesFixture) {
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
},
want: map[netip.Prefix][]types.NodeID{
mp("192.168.1.0/24"): {1, 2},
},
},
{
name: "two-nodes-different-prefixes-not-ha",
setup: func(f *primariesFixture) {
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.2.0/24"))
},
want: map[netip.Prefix][]types.NodeID{},
},
{
name: "three-nodes-two-share-prefix",
setup: func(f *primariesFixture) {
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
f.advertise(3, mp("10.0.0.0/8"))
},
want: map[netip.Prefix][]types.NodeID{
mp("192.168.1.0/24"): {1, 2},
},
},
{
name: "three-nodes-all-share",
setup: func(f *primariesFixture) {
f.advertise(1, mp("192.168.1.0/24"))
f.advertise(2, mp("192.168.1.0/24"))
f.advertise(3, mp("192.168.1.0/24"))
},
want: map[netip.Prefix][]types.NodeID{
mp("192.168.1.0/24"): {1, 2, 3},
},
},
{
name: "empty",
setup: func(*primariesFixture) {
},
want: map[netip.Prefix][]types.NodeID{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
f := newPrimariesFixture(t, 1, 2, 3)
tt.setup(f)
got := f.ns.HANodes()
assert.Equal(t, tt.want, got)
})
}
}