From 8f3ba2d8cd9e266bbbc90a76810547f40c599411 Mon Sep 17 00:00:00 2001
From: Serhii Matsyshyn <serhii.matsyshyn@ucu.edu.ua>
Date: Thu, 15 Oct 2026 19:25:19 +0300
Subject: [PATCH] state: rebuild peer maps in ReloadPolicy only when the policy
 changed

The SetPolicy API handler calls State.SetPolicy, which already rebuilds the peer map when the policy changed, and then ReloadPolicy, which rebuilt it again even when nothing changed. So every policy write built the peer map twice.

ReloadPolicy now rebuilds only when the policy really changed. A test checks one build per changed write and none for an unchanged reload.

Part of #3549
---
 hscontrol/state/reload_policy_test.go | 88 +++++++++++++++++++++++++++
 hscontrol/state/state.go              | 12 +++-
 2 files changed, 99 insertions(+), 1 deletion(-)
 create mode 100644 hscontrol/state/reload_policy_test.go

diff --git a/hscontrol/state/reload_policy_test.go b/hscontrol/state/reload_policy_test.go
new file mode 100644
index 000000000..6bb5dff5c
--- /dev/null
+++ b/hscontrol/state/reload_policy_test.go
@@ -0,0 +1,88 @@
+package state
+
+import (
+	"sync/atomic"
+	"testing"
+
+	"github.com/juanfont/headscale/hscontrol/db"
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/stretchr/testify/require"
+)
+
+// TestReloadPolicyRebuildsPeerMapOnChangeOnly covers the peer-map rebuilds
+// of a policy write. The API handlers apply a policy with [State.SetPolicy],
+// store it and then call [State.ReloadPolicy]; the O(n^2) rebuild must run
+// once for that, not once per step. A reload that is the first to see a new
+// policy (SIGHUP, a policy stored by another path) must still rebuild.
+func TestReloadPolicyRebuildsPeerMapOnChangeOnly(t *testing.T) {
+	const (
+		denyAll  = `{"grants": []}`
+		allowAll = `{"grants": [{"src": ["*"], "dst": ["*"], "ip": ["*"]}]}`
+	)
+
+	cfg := persistTestConfig(t.TempDir() + "/headscale.db")
+
+	database, err := db.NewHeadscaleDatabase(cfg)
+	require.NoError(t, err)
+
+	nodes := database.CreateRegisteredNodesForTest(database.CreateUserForTest("reload-user"), 2, "reload-node")
+	require.NoError(t, database.Close())
+
+	s, err := NewState(cfg)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = s.Close() })
+
+	// Count the peer-map builds. The writer goroutine reads peersFunc only
+	// while applying a batch, and none has been queued yet.
+	var builds atomic.Int32
+
+	peersFunc := s.nodeStore.peersFunc
+	s.nodeStore.peersFunc = func(n []types.NodeView) map[types.NodeID][]types.NodeID {
+		builds.Add(1)
+
+		return peersFunc(n)
+	}
+
+	// write is the sequence of the SetPolicy API handlers.
+	write := func(pol string) {
+		t.Helper()
+
+		_, err := s.SetPolicy([]byte(pol))
+		require.NoError(t, err)
+
+		_, err = s.SetPolicyInDB(pol)
+		require.NoError(t, err)
+
+		cs, err := s.ReloadPolicy()
+		require.NoError(t, err)
+		require.NotEmpty(t, cs, "a reload always reports a policy change to distribute")
+	}
+	peers := func() int { return s.ListPeers(nodes[0].ID).Len() }
+
+	require.Equal(t, 1, peers(), "no policy: allow all")
+
+	write(denyAll)
+	require.EqualValues(t, 1, builds.Swap(0), "a changed policy is rebuilt once per write")
+	require.Zero(t, peers())
+
+	write(denyAll)
+	require.Zero(t, builds.Swap(0), "the same policy again needs no rebuild")
+	require.Zero(t, peers())
+
+	write(allowAll)
+	require.EqualValues(t, 1, builds.Swap(0))
+	require.Equal(t, 1, peers())
+
+	// A reload that finds a policy the manager has not seen yet.
+	_, err = s.SetPolicyInDB(denyAll)
+	require.NoError(t, err)
+
+	_, err = s.ReloadPolicy()
+	require.NoError(t, err)
+	require.EqualValues(t, 1, builds.Swap(0), "a reload that changes the policy rebuilds")
+	require.Zero(t, peers())
+
+	_, err = s.ReloadPolicy()
+	require.NoError(t, err)
+	require.Zero(t, builds.Swap(0), "a reload of the current policy does not")
+}
diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go
index 04f2c8cde..fcf79db41 100644
--- a/hscontrol/state/state.go
+++ b/hscontrol/state/state.go
@@ -368,7 +368,17 @@ func (s *State) ReloadPolicy() ([]change.Change, error) {
 	// uses the [policy.PolicyManager]'s filters. Without this, nodes won't see
 	// newly allowed peers until a node is added/removed, causing autogroup:self
 	// policies to not propagate correctly when switching between policy types.
-	s.nodeStore.RebuildPeerMaps()
+	//
+	// An unchanged policy leaves the matchers as they were, and the other
+	// inputs of the peer map rebuild it themselves (node writes in the
+	// [NodeStore], [State.updatePolicyManagerNodes] and
+	// [State.updatePolicyManagerUsers]), so the O(n^2) rebuild is skipped
+	// then. The policy API handlers always get here that way: they apply
+	// the policy with [State.SetPolicy], which rebuilds, store it, and
+	// reload the very same bytes.
+	if policyChanged {
+		s.nodeStore.RebuildPeerMaps()
+	}
 
 	// Nodes whose CapMap shifted get their self refresh from
 	// [State.DrainSelfRefreshes] when these changes are dispatched.