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

@@ -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")
+}

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.