diff --git a/hscontrol/policy/v2/policy.go b/hscontrol/policy/v2/policy.go index d5ff21043..518247aad 100644 --- a/hscontrol/policy/v2/policy.go +++ b/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) } diff --git a/hscontrol/policy/v2/policy_test.go b/hscontrol/policy/v2/policy_test.go index 18eca4d92..65692d37b 100644 --- a/hscontrol/policy/v2/policy_test.go +++ b/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. diff --git a/hscontrol/policy/v2/types.go b/hscontrol/policy/v2/types.go index 0235a72eb..c37225624 100644 --- a/hscontrol/policy/v2/types.go +++ b/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]. diff --git a/hscontrol/policy/v2/types_test.go b/hscontrol/policy/v2/types_test.go index eeac6f8cc..624769cdb 100644 --- a/hscontrol/policy/v2/types_test.go +++ b/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"},