From 3feb307a12a79badf669c3f5f46811100efab02b Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Tue, 28 Apr 2026 12:44:58 +0000 Subject: [PATCH] state: compute primary routes inside NodeStore snapshot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add primaries + isPrimary fields to Snapshot. snapshotFromNodes now elects per-prefix primary advertisers from the current node set, mirroring routes.PrimaryRoutes.updatePrimaryLocked: skip exit routes, prefer the previous primary if still a valid healthy advertiser, fall back to the lowest healthy NodeID, then to the lowest NodeID overall when all are unhealthy. Anti-flap memory rides the previous snapshot through applyBatch. Caller-side coordination is unnecessary; the writer goroutine serialises every mutation that could change the primary set. Add NodeStore reader methods (PrimaryFor, PrimaryRoutesForNode, HANodes, IsNodeHealthy) that mirror the legacy routes.PrimaryRoutes API on the snapshot. Consumers will switch in the next commit. No behaviour change yet — primaries are computed but not read. Updates #3203 --- hscontrol/state/node_store.go | 180 +++++++++++++++++++++++++++-- hscontrol/state/node_store_test.go | 2 +- 2 files changed, 173 insertions(+), 9 deletions(-) diff --git a/hscontrol/state/node_store.go b/hscontrol/state/node_store.go index 0453ad2d..8f9d30e7 100644 --- a/hscontrol/state/node_store.go +++ b/hscontrol/state/node_store.go @@ -4,6 +4,8 @@ import ( "errors" "fmt" "maps" + "net/netip" + "slices" "strconv" "strings" "sync/atomic" @@ -12,6 +14,7 @@ import ( "github.com/juanfont/headscale/hscontrol/types" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" + "tailscale.com/net/tsaddr" "tailscale.com/types/key" "tailscale.com/types/views" "tailscale.com/util/dnsname" @@ -114,7 +117,7 @@ func NewNodeStore(allNodes types.Nodes, peersFunc PeersFunc, batchSize int, batc nodes[n.ID] = *n } - snap := snapshotFromNodes(nodes, peersFunc) + snap := snapshotFromNodes(nodes, peersFunc, nil) store := &NodeStore{ peersFunc: peersFunc, @@ -144,6 +147,12 @@ type Snapshot struct { peersByNode map[types.NodeID][]types.NodeView nodesByUser map[types.UserID][]types.NodeView allNodes []types.NodeView + + // routes maps each prefix to its current primary advertiser. The + // previous assignment is carried over when still valid so the + // primary does not flap on every unrelated batch. + routes map[netip.Prefix]types.NodeID + isPrimaryRoute map[types.NodeID]bool } // PeersFunc is a function that takes a list of nodes and returns a map @@ -459,7 +468,8 @@ func (s *NodeStore) applyBatch(batch []work) { } } - newSnap := snapshotFromNodes(nodes, s.peersFunc) + prev := s.data.Load() + newSnap := snapshotFromNodes(nodes, s.peersFunc, prev.routes) s.data.Store(&newSnap) // Update node count gauge @@ -543,12 +553,14 @@ func resolveGivenName(nodes map[types.NodeID]types.Node, self types.NodeID, base } } -// snapshotFromNodes creates a new Snapshot from the provided nodes. -// It builds a lot of "indexes" to make lookups fast for datasets we -// that is used frequently, like nodesByNodeKey, peersByNode, and nodesByUser. -// This is not a fast operation, it is the "slow" part of our copy-on-write -// structure, but it allows us to have fast reads and efficient lookups. -func snapshotFromNodes(nodes map[types.NodeID]types.Node, peersFunc PeersFunc) Snapshot { +// snapshotFromNodes builds the index maps and primary-route table for +// a new Snapshot. prevRoutes carries forward the previous primary +// assignment so a still-valid choice survives unrelated batches. +func snapshotFromNodes( + nodes map[types.NodeID]types.Node, + peersFunc PeersFunc, + prevRoutes map[netip.Prefix]types.NodeID, +) Snapshot { timer := prometheus.NewTimer(nodeStoreSnapshotBuildDuration) defer timer.ObserveDuration() @@ -557,6 +569,8 @@ func snapshotFromNodes(nodes map[types.NodeID]types.Node, peersFunc PeersFunc) S allNodes = append(allNodes, n.View()) } + routes, isPrimaryRoute := electPrimaryRoutes(nodes, prevRoutes) + newSnap := Snapshot{ nodesByID: nodes, allNodes: allNodes, @@ -574,6 +588,9 @@ func snapshotFromNodes(nodes map[types.NodeID]types.Node, peersFunc PeersFunc) S return peersFunc(allNodes) }(), nodesByUser: make(map[types.UserID][]types.NodeView), + + routes: routes, + isPrimaryRoute: isPrimaryRoute, } // Build nodesByUser, nodesByNodeKey, and nodesByMachineKey maps @@ -600,6 +617,81 @@ func snapshotFromNodes(nodes map[types.NodeID]types.Node, peersFunc PeersFunc) S return newSnap } +// electPrimaryRoutes picks the primary advertiser for each non-exit +// prefix. The previous primary is preserved when it is still online +// and healthy (anti-flap); otherwise the lowest-NodeID healthy +// advertiser wins, falling back to the lowest-NodeID candidate when +// every advertiser is unhealthy so peers see *some* primary instead +// of none. +func electPrimaryRoutes( + nodes map[types.NodeID]types.Node, + prev map[netip.Prefix]types.NodeID, +) (map[netip.Prefix]types.NodeID, map[types.NodeID]bool) { + ids := make([]types.NodeID, 0, len(nodes)) + for id := range nodes { + ids = append(ids, id) + } + + slices.Sort(ids) + + advertisers := make(map[netip.Prefix][]types.NodeID) + + for _, id := range ids { + n := nodes[id] + if n.IsOnline == nil || !*n.IsOnline { + continue + } + + for _, p := range n.AllApprovedRoutes() { + if tsaddr.IsExitRoute(p) { + continue + } + + advertisers[p] = append(advertisers[p], id) + } + } + + routes := make(map[netip.Prefix]types.NodeID, len(advertisers)) + for prefix, candidates := range advertisers { + if cur, ok := prev[prefix]; ok && + slices.Contains(candidates, cur) && + !nodes[cur].Unhealthy { + routes[prefix] = cur + continue + } + + var ( + selected types.NodeID + found bool + ) + + for _, c := range candidates { + if !nodes[c].Unhealthy { + selected = c + found = true + + break + } + } + + if !found && len(candidates) >= 1 { + selected = candidates[0] + found = true + } + + if found { + routes[prefix] = selected + } + } + + isPrimaryRoute := make(map[types.NodeID]bool, len(routes)) + for _, id := range routes { + isPrimaryRoute[id] = true + } + + return routes, isPrimaryRoute +} + // GetNode retrieves a node by its ID. // The bool indicates if the node exists or is available (like "err not found"). // The NodeView might be invalid, so it must be checked with .Valid(), which must be used to ensure @@ -752,6 +844,78 @@ func (s *NodeStore) ListPeers(id types.NodeID) views.Slice[types.NodeView] { return views.SliceOf(s.data.Load().peersByNode[id]) } +// PrimaryRouteFor returns the current primary advertiser for prefix. +func (s *NodeStore) PrimaryRouteFor(prefix netip.Prefix) (types.NodeID, bool) { + id, ok := s.data.Load().routes[prefix] + return id, ok +} + +// PrimaryRoutesForNode returns the prefixes for which id is the current +// primary advertiser. +func (s *NodeStore) PrimaryRoutesForNode(id types.NodeID) []netip.Prefix { + snap := s.data.Load() + if !snap.isPrimaryRoute[id] { + return nil + } + + out := make([]netip.Prefix, 0) + + for prefix, nodeID := range snap.routes { + if nodeID == id { + out = append(out, prefix) + } + } + + return out +} + +// HANodes returns the prefixes with two or more online advertisers, the +// candidate set the HA prober needs to monitor. +func (s *NodeStore) HANodes() map[netip.Prefix][]types.NodeID { + snap := s.data.Load() + + advertisers := make(map[netip.Prefix][]types.NodeID) + + for id, n := range snap.nodesByID { + if n.IsOnline == nil || !*n.IsOnline { + continue + } + + for _, p := range n.AllApprovedRoutes() { + if tsaddr.IsExitRoute(p) { + continue + } + + advertisers[p] = append(advertisers[p], id) + } + } + + out := make(map[netip.Prefix][]types.NodeID) + + for p, ids := range advertisers { + if len(ids) < 2 { + continue + } + + slices.Sort(ids) + out[p] = ids + } + + return out +} + +// IsNodeHealthy reports whether the HA prober considers id healthy. +// Unknown nodes report healthy so absence does not exclude them from +// election. +func (s *NodeStore) IsNodeHealthy(id types.NodeID) bool { + n, ok := s.data.Load().nodesByID[id] + if !ok { + return true + } + + return !n.Unhealthy +} + // RebuildPeerMaps rebuilds the peer relationship map using the current peersFunc. // This must be called after policy changes because peersFunc uses PolicyManager's // filters to determine which nodes can see each other. Without rebuilding, the diff --git a/hscontrol/state/node_store_test.go b/hscontrol/state/node_store_test.go index a8b02d8d..6b421754 100644 --- a/hscontrol/state/node_store_test.go +++ b/hscontrol/state/node_store_test.go @@ -150,7 +150,7 @@ func TestSnapshotFromNodes(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { nodes, peersFunc := tt.setupFunc() - snapshot := snapshotFromNodes(nodes, peersFunc) + snapshot := snapshotFromNodes(nodes, peersFunc, nil) tt.validate(t, nodes, snapshot) }) }