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.