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

policy,state: key the peer map by node ID

Browse Source 
Adjacency becomes immutable, so a snapshot can resolve peers through its
own fresh views instead of storing them.

Updates #3417

(cherry picked from commit 95ba787417)
This commit is contained in:
Kristoffer Dalby
2026-09-09 14:41:27 +00:00
parent 9fbf7b9b60
commit 938c2bd753
8 changed files with 96 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hscontrol/state/node_store.go
+28 -9
View File
@@ -148,13 +148,18 @@ func NewNodeStore(allNodes types.Nodes, peersFunc PeersFunc, batchSize int, batc
type Snapshot struct {
// nodesByID is the main source of truth for nodes.
nodesByID map[types.NodeID]types.Node
// nodeViewsByID resolves peer adjacency IDs to immutable views from this
// snapshot without rebuilding every peer slice on each write.
nodeViewsByID map[types.NodeID]types.NodeView
// calculated from nodesByID
nodesByNodeKey map[key.NodePublic]types.NodeView
nodesByMachineKey map[key.MachinePublic]map[types.UserID]types.NodeView
peersByNode map[types.NodeID][]types.NodeView
nodesByUser map[types.UserID][]types.NodeView
allNodes []types.NodeView
// peersByNode stores immutable adjacency as IDs; ListPeers resolves the
// corresponding views through nodeViewsByID.
peersByNode map[types.NodeID][]types.NodeID
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
@@ -167,7 +172,7 @@ type Snapshot struct {
// with the relationships between nodes and their peers.
// This will typically be used to calculate which nodes can see each other
// based on the current policy.
type PeersFunc func(nodes []types.NodeView) map[types.NodeID][]types.NodeView
type PeersFunc func(nodes []types.NodeView) map[types.NodeID][]types.NodeID
// work represents a single operation to be performed on the [NodeStore].
type work struct {
@@ -617,14 +622,19 @@ func snapshotFromNodes(
defer timer.ObserveDuration()
allNodes := make([]types.NodeView, 0, len(nodes))
nodeViewsByID := make(map[types.NodeID]types.NodeView, len(nodes))
for _, n := range nodes {
allNodes = append(allNodes, n.View())
nv := n.View()
allNodes = append(allNodes, nv)
nodeViewsByID[n.ID] = nv
}
routes, isPrimaryRoute := electPrimaryRoutes(nodes, prevRoutes)
newSnap := Snapshot{
nodesByID: nodes,
nodeViewsByID: nodeViewsByID,
allNodes: allNodes,
nodesByNodeKey: make(map[key.NodePublic]types.NodeView),
nodesByMachineKey: make(map[key.MachinePublic]map[types.UserID]types.NodeView),
@@ -633,7 +643,7 @@ func snapshotFromNodes(
// it will use the list of all nodes, combined with the
// current policy to precalculate which nodes are peers and
// can see each other.
peersByNode: func() map[types.NodeID][]types.NodeView {
peersByNode: func() map[types.NodeID][]types.NodeID {
peersTimer := prometheus.NewTimer(nodeStorePeersCalculationDuration)
defer peersTimer.ObserveDuration()
@@ -647,7 +657,7 @@ func snapshotFromNodes(
// Build nodesByUser, nodesByNodeKey, and nodesByMachineKey maps
for _, n := range nodes {
nodeView := n.View()
nodeView := nodeViewsByID[n.ID]
userID := n.TypedUserID()
// Tagged nodes are owned by their tags, not a user,
@@ -658,7 +668,6 @@ func snapshotFromNodes(
newSnap.nodesByNodeKey[n.NodeKey] = nodeView
// Build machine key index
if newSnap.nodesByMachineKey[n.MachineKey] == nil {
newSnap.nodesByMachineKey[n.MachineKey] = make(map[types.UserID]types.NodeView)
}
@@ -886,7 +895,17 @@ func (s *NodeStore) ListPeers(id types.NodeID) views.Slice[types.NodeView] {
nodeStoreOperations.WithLabelValues("list_peers").Inc()
return views.SliceOf(s.data.Load().peersByNode[id])
snapshot := s.data.Load()
peerIDs := snapshot.peersByNode[id]
peers := make([]types.NodeView, 0, len(peerIDs))
for _, peerID := range peerIDs {
if peer, ok := snapshot.nodeViewsByID[peerID]; ok {
peers = append(peers, peer)
}
}
return views.SliceOf(peers)
}
// PrimaryRouteFor returns the current primary advertiser for prefix.
hscontrol/state/node_store_test.go
+33 -33
View File
@@ -26,8 +26,8 @@ func TestSnapshotFromNodes(t *testing.T) {
name: "empty nodes",
setupFunc: func() (map[types.NodeID]types.Node, PeersFunc) {
nodes := make(map[types.NodeID]types.Node)
peersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeView {
return make(map[types.NodeID][]types.NodeView)
peersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeID {
return make(map[types.NodeID][]types.NodeID)
}
return nodes, peersFunc
@@ -79,9 +79,9 @@ func TestSnapshotFromNodes(t *testing.T) {
// Each node sees the other as peer (but not itself)
assert.Len(t, snapshot.peersByNode[1], 1)
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[1][0].ID())
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[1][0])
assert.Len(t, snapshot.peersByNode[2], 1)
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[2][0].ID())
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[2][0])
assert.Len(t, snapshot.nodesByUser[1], 2)
},
},
@@ -133,17 +133,17 @@ func TestSnapshotFromNodes(t *testing.T) {
// Odd nodes should only see other odd nodes as peers
require.Len(t, snapshot.peersByNode[1], 1)
assert.Equal(t, types.NodeID(3), snapshot.peersByNode[1][0].ID())
assert.Equal(t, types.NodeID(3), snapshot.peersByNode[1][0])
require.Len(t, snapshot.peersByNode[3], 1)
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[3][0].ID())
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[3][0])
// Even nodes should only see other even nodes as peers
require.Len(t, snapshot.peersByNode[2], 1)
assert.Equal(t, types.NodeID(4), snapshot.peersByNode[2][0].ID())
assert.Equal(t, types.NodeID(4), snapshot.peersByNode[2][0])
require.Len(t, snapshot.peersByNode[4], 1)
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[4][0].ID())
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[4][0])
},
},
}
@@ -190,14 +190,14 @@ func createTestNode(nodeID types.NodeID, userID uint, username, hostname string)
// Peer functions
func allowAllPeersFunc(nodes []types.NodeView) map[types.NodeID][]types.NodeView {
ret := make(map[types.NodeID][]types.NodeView, len(nodes))
func allowAllPeersFunc(nodes []types.NodeView) map[types.NodeID][]types.NodeID {
ret := make(map[types.NodeID][]types.NodeID, len(nodes))
for _, node := range nodes {
var peers []types.NodeView
var peers []types.NodeID
for _, n := range nodes {
if n.ID() != node.ID() {
peers = append(peers, n)
peers = append(peers, n.ID())
}
}
@@ -207,10 +207,10 @@ func allowAllPeersFunc(nodes []types.NodeView) map[types.NodeID][]types.NodeView
return ret
}
func oddEvenPeersFunc(nodes []types.NodeView) map[types.NodeID][]types.NodeView {
ret := make(map[types.NodeID][]types.NodeView, len(nodes))
func oddEvenPeersFunc(nodes []types.NodeView) map[types.NodeID][]types.NodeID {
ret := make(map[types.NodeID][]types.NodeID, len(nodes))
for _, node := range nodes {
var peers []types.NodeView
var peers []types.NodeID
nodeIsOdd := node.ID()%2 == 1
@@ -223,7 +223,7 @@ func oddEvenPeersFunc(nodes []types.NodeView) map[types.NodeID][]types.NodeView
// Only add peer if both are odd or both are even
if nodeIsOdd == peerIsOdd {
peers = append(peers, n)
peers = append(peers, n.ID())
}
}
@@ -313,9 +313,9 @@ func TestNodeStoreOperations(t *testing.T) {
// Now both nodes should see each other as peers
assert.Len(t, snapshot.peersByNode[1], 1)
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[1][0].ID())
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[1][0])
assert.Len(t, snapshot.peersByNode[2], 1)
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[2][0].ID())
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[2][0])
assert.Len(t, snapshot.nodesByUser[1], 2)
},
},
@@ -382,9 +382,9 @@ func TestNodeStoreOperations(t *testing.T) {
// Remaining nodes should see each other as peers
assert.Len(t, snapshot.peersByNode[1], 1)
assert.Equal(t, types.NodeID(3), snapshot.peersByNode[1][0].ID())
assert.Equal(t, types.NodeID(3), snapshot.peersByNode[1][0])
assert.Len(t, snapshot.peersByNode[3], 1)
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[3][0].ID())
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[3][0])
// User groupings updated
assert.Len(t, snapshot.nodesByUser[1], 1) // user1 now has only node 1
@@ -475,16 +475,16 @@ func TestNodeStoreOperations(t *testing.T) {
// Verify odd-even peer relationships
require.Len(t, snapshot.peersByNode[1], 1)
assert.Equal(t, types.NodeID(3), snapshot.peersByNode[1][0].ID())
assert.Equal(t, types.NodeID(3), snapshot.peersByNode[1][0])
require.Len(t, snapshot.peersByNode[2], 1)
assert.Equal(t, types.NodeID(4), snapshot.peersByNode[2][0].ID())
assert.Equal(t, types.NodeID(4), snapshot.peersByNode[2][0])
require.Len(t, snapshot.peersByNode[3], 1)
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[3][0].ID())
assert.Equal(t, types.NodeID(1), snapshot.peersByNode[3][0])
require.Len(t, snapshot.peersByNode[4], 1)
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[4][0].ID())
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[4][0])
},
},
{
@@ -500,9 +500,9 @@ func TestNodeStoreOperations(t *testing.T) {
// Even nodes should still see each other
require.Len(t, snapshot.peersByNode[2], 1)
assert.Equal(t, types.NodeID(4), snapshot.peersByNode[2][0].ID())
assert.Equal(t, types.NodeID(4), snapshot.peersByNode[2][0])
require.Len(t, snapshot.peersByNode[4], 1)
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[4][0].ID())
assert.Equal(t, types.NodeID(2), snapshot.peersByNode[4][0])
},
},
},
@@ -1245,16 +1245,16 @@ func TestRebuildPeerMapsWithChangedPeersFunc(t *testing.T) {
// This simulates how PolicyManager.BuildPeerMap works - it reads state
// that can change between calls
dynamicPeersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeView {
ret := make(map[types.NodeID][]types.NodeView, len(nodes))
dynamicPeersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeID {
ret := make(map[types.NodeID][]types.NodeID, len(nodes))
if allowPeers {
// Allow all peers
for _, node := range nodes {
var peers []types.NodeView
var peers []types.NodeID
for _, n := range nodes {
if n.ID() != node.ID() {
peers = append(peers, n)
peers = append(peers, n.ID())
}
}
@@ -1263,7 +1263,7 @@ func TestRebuildPeerMapsWithChangedPeersFunc(t *testing.T) {
} else {
// Allow no peers
for _, node := range nodes {
ret[node.ID()] = []types.NodeView{}
ret[node.ID()] = []types.NodeID{}
}
}
@@ -1285,8 +1285,8 @@ func TestRebuildPeerMapsWithChangedPeersFunc(t *testing.T) {
snapshot := store.data.Load()
require.Len(t, snapshot.peersByNode[1], 1, "node1 should have 1 peer initially")
require.Len(t, snapshot.peersByNode[2], 1, "node2 should have 1 peer initially")
require.Equal(t, types.NodeID(2), snapshot.peersByNode[1][0].ID())
require.Equal(t, types.NodeID(1), snapshot.peersByNode[2][0].ID())
require.Equal(t, types.NodeID(2), snapshot.peersByNode[1][0])
require.Equal(t, types.NodeID(1), snapshot.peersByNode[2][0])
// Now "change the policy" by disabling peers
allowPeers = false
hscontrol/state/state.go
+1 -1
View File
@@ -281,7 +281,7 @@ func NewState(cfg *types.Config) (*State, error) {
// This moves the complex peer relationship logic into the policy package where it belongs.
nodeStore := NewNodeStore(
nodes,
func(nodes []types.NodeView) map[types.NodeID][]types.NodeView {
func(nodes []types.NodeView) map[types.NodeID][]types.NodeID {
return polMan.BuildPeerMap(views.SliceOf(nodes))
},
batchSize,