diff --git a/hscontrol/routes/primary.go b/hscontrol/routes/primary.go deleted file mode 100644 index dd1694ac..00000000 --- a/hscontrol/routes/primary.go +++ /dev/null @@ -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 -} diff --git a/hscontrol/routes/primary_property_test.go b/hscontrol/routes/primary_property_test.go deleted file mode 100644 index 7047e596..00000000 --- a/hscontrol/routes/primary_property_test.go +++ /dev/null @@ -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) - } - }) -} diff --git a/hscontrol/routes/primary_test.go b/hscontrol/routes/primary_test.go deleted file mode 100644 index bd0c615b..00000000 --- a/hscontrol/routes/primary_test.go +++ /dev/null @@ -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) - } - }) - } -} diff --git a/hscontrol/state/primaries_property_test.go b/hscontrol/state/primaries_property_test.go new file mode 100644 index 00000000..2b7a7562 --- /dev/null +++ b/hscontrol/state/primaries_property_test.go @@ -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) + } + }) +} diff --git a/hscontrol/state/primaries_test.go b/hscontrol/state/primaries_test.go new file mode 100644 index 00000000..38f9bb53 --- /dev/null +++ b/hscontrol/state/primaries_test.go @@ -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) + }) + } +}