policy/v2: index nodes by tag once per compile

Each tag alias scanned every node, so with many tags a compile, run on
every NodeStore batch, was mostly HasTag.

hscontrol/policy/v2/policy.go

@@ -246,9 +246,12 @@ func NewPolicyManager(b []byte, users []types.User, nodes views.Slice[types.Node
 // pm is written, so a caller that restores its own input on error leaves pm
 // as it was, and a retry with the same input recompiles.
 func (pm *PolicyManager) updateLocked() (bool, error) {
+	// One tag index serves every resolution in this compile.
+	pol := pm.pol.withNodesByTag(pm.nodes)
+
 	// Compile all grants once. Both global and per-node filter
 	// rules are derived from these compiled grants.
-	grants := pm.pol.compileGrants(pm.users, pm.nodes)
+	grants := pol.compileGrants(pm.users, pm.nodes)
 
 	relayTargetIPs, err := collectRelayTargetIPs(grants)
 	if err != nil {
@@ -258,12 +261,12 @@ func (pm *PolicyManager) updateLocked() (bool, error) {
 	// Order matters, tags might be used in autoapprovers, so we need to ensure
 	// that the map for tag owners is resolved before resolving autoapprovers.
 	// TODO(kradalby): Order might not matter after #2417
-	tagMap, err := resolveTagOwners(pm.pol, pm.users, pm.nodes)
+	tagMap, err := resolveTagOwners(pol, pm.users, pm.nodes)
 	if err != nil {
 		return false, fmt.Errorf("resolving tag owners map: %w", err)
 	}
 
-	autoMap, exitSet, err := resolveAutoApprovers(pm.pol, pm.users, pm.nodes)
+	autoMap, exitSet, err := resolveAutoApprovers(pol, pm.users, pm.nodes)
 	if err != nil {
 		return false, fmt.Errorf("resolving auto approvers map: %w", err)
 	}
@@ -279,7 +282,7 @@ func (pm *PolicyManager) updateLocked() (bool, error) {
 
 	var nodeAttrs map[types.NodeID]tailcfg.NodeCapMap
 	if refreshNodeAttrs {
-		nodeAttrs, err = pm.pol.compileNodeAttrs(pm.users, pm.nodes)
+		nodeAttrs, err = pol.compileNodeAttrs(pm.users, pm.nodes)
 		if err != nil {
 			return false, fmt.Errorf("compiling nodeAttrs: %w", err)
 		}

hscontrol/policy/v2/policy_test.go

@@ -3431,6 +3431,33 @@ func BenchmarkSetNodes(b *testing.B) {
 	}
 }
 
+// BenchmarkSetNodesManyTags measures the recompile a join or delete costs
+// when every node carries its own tag and grants reference them, a CI fleet
+// shape where resolving tags dominated the compile.
+func BenchmarkSetNodesManyTags(b *testing.B) {
+	for _, n := range []int{1000, 4000} {
+		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
+			nodes := scanManyRulesNodes(n)
+			pm, err := NewPolicyManager([]byte(scanManyRulesPolicy(n, 400)), nil, nodes.ViewSlice())
+			require.NoError(b, err)
+
+			sets := []views.Slice[types.NodeView]{nodes[:n-1].ViewSlice(), nodes.ViewSlice()}
+
+			b.ReportAllocs()
+
+			i := 0
+			for b.Loop() {
+				i++
+
+				_, err := pm.SetNodes(sets[i%2])
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
+	}
+}
+
 // TestNodesGenerationCountsChangingSetNodes pins that NodesGeneration
 // moves exactly when SetNodes reports a change, so a caller that did not
 // run the SetNodes itself can still tell its write moved the policy.

hscontrol/policy/v2/types.go

@@ -575,6 +575,14 @@ func (t *Tag) Resolve(p *Policy, users types.Users, nodes views.Slice[types.Node
 func (t *Tag) resolve(p *Policy, users types.Users, nodes views.Slice[types.NodeView]) (*netipx.IPSet, error) {
 	var ips netipx.IPSetBuilder
 
+	if p != nil && p.nodesByTag != nil {
+		for _, node := range p.nodesByTag[string(*t)] {
+			node.AppendToIPSet(&ips)
+		}
+
+		return ips.IPSet()
+	}
+
 	for _, node := range nodes.All() {
 		// Check if node has this tag
 		if node.HasTag(string(*t)) {
@@ -1961,6 +1969,33 @@ type Policy struct {
 	Tests               []PolicyTest       `json:"tests,omitempty"`
 	SSHTests            []SSHPolicyTest    `json:"sshTests,omitempty"`
 	RandomizeClientPort bool               `json:"randomizeClientPort,omitempty"`
+
+	// nodesByTag is set only on the copy [Policy.withNodesByTag] makes for
+	// one compile, and indexes the nodes that compile resolves against.
+	nodesByTag map[string][]types.NodeView
+}
+
+// withNodesByTag returns a copy of pol to resolve against nodes, and only
+// nodes. Its tag aliases look their members up instead of scanning every
+// node: a compile resolves every tag alias occurrence, so with many tags the
+// scans were most of the compile.
+func (pol *Policy) withNodesByTag(nodes views.Slice[types.NodeView]) *Policy {
+	if pol == nil {
+		return nil
+	}
+
+	byTag := make(map[string][]types.NodeView)
+
+	for _, node := range nodes.All() {
+		for _, tag := range node.Tags().All() {
+			byTag[tag] = append(byTag[tag], node)
+		}
+	}
+
+	indexed := *pol
+	indexed.nodesByTag = byTag
+
+	return &indexed
 }
 
 // MarshalJSON is deliberately not implemented for [Policy].

hscontrol/policy/v2/types_test.go

@@ -2823,6 +2823,33 @@ func TestResolvePolicy(t *testing.T) {
 	}
 }
 
+// TestTagResolveIndexedMatchesScan resolves tags with and without the
+// per-compile tag index: the index must name exactly the nodes HasTag does.
+func TestTagResolveIndexedMatchesScan(t *testing.T) {
+	v4 := func(s string) *netip.Addr { a := netip.MustParseAddr(s); return &a }
+
+	nodes := types.Nodes{
+		{ID: 1, IPv4: v4("100.64.0.1"), IPv6: v4("fd7a:115c:a1e0::1"), Tags: []string{"tag:a", "tag:b"}},
+		{ID: 2, IPv4: v4("100.64.0.2"), Tags: []string{"tag:b"}},
+		{ID: 3, IPv4: v4("100.64.0.3")},
+		{ID: 4, IPv4: v4("100.64.0.4"), Tags: []string{"tag:a", "tag:a"}},
+		{ID: 5, IPv6: v4("fd7a:115c:a1e0::5"), Tags: []string{"tag:c"}},
+	}
+
+	pol := &Policy{}
+	indexed := pol.withNodesByTag(nodes.ViewSlice())
+
+	for _, tag := range []Tag{"tag:a", "tag:b", "tag:c", "tag:missing"} {
+		want, err := tag.resolve(pol, nil, nodes.ViewSlice())
+		require.NoError(t, err)
+
+		got, err := tag.resolve(indexed, nil, nodes.ViewSlice())
+		require.NoError(t, err)
+
+		assert.Equal(t, want.Prefixes(), got.Prefixes(), "%s", tag)
+	}
+}
+
 func TestResolveAutoApprovers(t *testing.T) {
 	users := types.Users{
 		{ID: 1, Name: "user1"},