ast: fix AnnotationSet memory leak via runtime.AddCleanup cycle

AnnotationSet.MergedLabels was introduced in v1.17.0 to cache merged
label maps per rule. It used weak.Pointer[Rule] as the cache key and
registered a runtime.AddCleanup to evict the entry when the rule was
garbage-collected.

The cleanup closure captured `as` (the AnnotationSet pointer). The
AnnotationSet holds strong references to every module it was built
from (as.modules), and each module holds its rules. This meant that
once any rule had a cleanup registered:

  runtime cleanup queue → closure → AnnotationSet → modules → Rule

Rule was always reachable through that path, so the cleanup could
never fire. Nothing would ever delete the closure, so the
AnnotationSet and all of its rules were permanently retained.

In practice, OPA's dynamic bundle plugin recompiles policies on every
poll cycle. Each compilation creates a fresh AnnotationSet. With the
bug, old AnnotationSets accumulated in the heap indefinitely, causing
the OOM-kill pattern reported in #8817.

Fix: drop the cache entirely. MergedLabels now calls Chain and
mergeChainLabels on every invocation. Chain is a handful of map
lookups and MergedLabels is called at most once per evaluated rule per
request, so the recomputation cost is negligible. This removes the
mergedLabels sync.Map field, the ruleLabelsEntry type, and the
runtime/sync/weak imports.

A regression test uses weak.Pointer[AnnotationSet] to assert that an
AnnotationSet is collectable after it goes out of scope.

Fixes #8817
Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

v1/ast/annotations.go

@@ -9,11 +9,8 @@ import (
 	"fmt"
 	"maps"
 	"net/url"
-	"runtime"
 	"slices"
 	"strings"
-	"sync"
-	"weak"
 
 	"github.com/open-policy-agent/opa/internal/deepcopy"
 	astJSON "github.com/open-policy-agent/opa/v1/ast/json"
@@ -70,11 +67,10 @@ type (
 	}
 
 	AnnotationSet struct {
-		byRule       map[*Rule][]*Annotations
-		byPackage    map[int]*Annotations
-		byPath       *annotationTreeNode
-		modules      []*Module // Modules this set was constructed from
-		mergedLabels sync.Map  // map[weak.Pointer[Rule]]*ruleLabelsEntry; lazily populated, entries cleaned up via runtime.AddCleanup when rules are GC'd
+		byRule    map[*Rule][]*Annotations
+		byPackage map[int]*Annotations
+		byPath    *annotationTreeNode
+		modules   []*Module // Modules this set was constructed from
 	}
 
 	annotationTreeNode struct {
@@ -964,42 +960,19 @@ func (as *AnnotationSet) Chain(rule *Rule) AnnotationsRefSet {
 	return refs
 }
 
-// ruleLabelsEntry caches the merged labels and dedup key for a single rule.
-type ruleLabelsEntry struct {
-	labels map[string]any
-	key    string
-}
-
 // MergedLabels returns the inner-scope-wins merged labels for the given rule
-// along with a stable string suitable for content-based deduplication. The
-// result is computed once per rule and cached on the AnnotationSet; the cache
-// entry is dropped automatically when the rule is garbage-collected.
-//
-// labels is nil when the rule has no labels anywhere in its annotation chain;
-// in that case key is the empty string.
+// along with a stable JSON string suitable for content-based deduplication.
+// labels is nil when the rule has no labels anywhere in its annotation chain.
 func (as *AnnotationSet) MergedLabels(rule *Rule) (labels map[string]any, key string) {
 	if as == nil {
 		return nil, ""
 	}
-	k := weak.Make(rule)
-	if v, ok := as.mergedLabels.Load(k); ok {
-		e := v.(*ruleLabelsEntry)
-		return e.labels, e.key
+	labels = mergeChainLabels(as.Chain(rule))
+	if len(labels) > 0 {
+		b, _ := json.Marshal(labels)
+		key = string(b)
 	}
-	merged := mergeChainLabels(as.Chain(rule))
-	e := &ruleLabelsEntry{labels: merged}
-	if len(merged) > 0 {
-		b, _ := json.Marshal(merged)
-		e.key = string(b)
-	}
-	actual, loaded := as.mergedLabels.LoadOrStore(k, e)
-	if !loaded {
-		// k is a weak.Pointer (value type) — it does not keep rule alive, so
-		// the cleanup will fire once the rule becomes unreachable elsewhere.
-		runtime.AddCleanup(rule, func(k weak.Pointer[Rule]) { as.mergedLabels.Delete(k) }, k)
-	}
-	e = actual.(*ruleLabelsEntry)
-	return e.labels, e.key
+	return labels, key
 }
 
 // mergeChainLabels folds labels from a rule's annotation chain with inner-wins

v1/ast/annotations_test.go

@@ -8,7 +8,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"maps"
+	"runtime"
 	"testing"
+	"weak"
 )
 
 func TestEntrypointAnnotationScopeRequirements(t *testing.T) {
@@ -1302,3 +1304,46 @@ func schemaAnnotationFromMap(path string, def map[string]any) *SchemaAnnotation
 	var p any = def
 	return &SchemaAnnotation{Path: MustParseRef(path), Definition: &p}
 }
+
+// TestAnnotationSet_MergedLabels_Collectable verifies that an AnnotationSet
+// can be garbage-collected after it goes out of scope, even when MergedLabels
+// has been called (which populates the internal mergedLabels cache).
+//
+// Regression test for a memory leak where the cleanup closure registered by
+// runtime.AddCleanup captured the AnnotationSet itself. The AnnotationSet
+// holds strong references to all modules (and their rules) via as.modules, so
+// the cleanup could never fire: rule → cleanup closure → AnnotationSet →
+// modules → rule. Each bundle reload left the old AnnotationSet permanently
+// alive, causing unbounded heap growth.
+func TestAnnotationSet_MergedLabels_Collectable(t *testing.T) {
+	const src = `package test
+
+# METADATA
+# labels:
+#   tier: fast
+allow if true
+`
+	var watch weak.Pointer[AnnotationSet]
+
+	func() {
+		mod := MustParseModuleWithOpts(src, ParserOptions{ProcessAnnotation: true})
+		as, errs := BuildAnnotationSet([]*Module{mod})
+		if len(errs) > 0 {
+			t.Fatalf("BuildAnnotationSet: %v", errs)
+		}
+		// Populate the mergedLabels cache for all rules.
+		for _, r := range mod.Rules {
+			as.MergedLabels(r)
+		}
+		watch = weak.Make(as)
+		// mod and as go out of scope here.
+	}()
+
+	// Two GC cycles: one to discover unreachable objects, one to collect them.
+	runtime.GC()
+	runtime.GC()
+
+	if watch.Value() != nil {
+		t.Fatal("AnnotationSet was not garbage-collected: mergedLabels cache likely holds a retaining cycle")
+	}
+}