perf: improved annotations parsing (#8210)

This wasn't really work I planned to do, and not driven by performance requirements. Just stumbled upon the metadata parsing code and thought it could be made better looking. And I think it is now, while also performing a bit better. Almost all the remaining cost now is unmarshalling YAML though, so I think this is about as good as it gets. Also added two new benchmarks for this. Signed-off-by: Anders Eknert <anders.eknert@apple.com>
2026-08-12 19:32:48 -06:00 · 2026-01-13 10:40:44 +01:00
parent bcd57a207c
commit a3c28af2b6
3 changed files with 147 additions and 79 deletions
@@ -71,6 +71,10 @@ var (
 	// copy them to  the call term only when needed
 	memberWithKeyRef = MemberWithKey.Ref()
 	memberRef        = Member.Ref()
+
+	newlineBytes       = []byte{'\n'}
+	metadataBytes      = []byte("METADATA")
+	metadataParserPool = util.NewSyncPool[metadataParser]()
 )

 func (v RegoVersion) Int() int {
@@ -540,44 +544,46 @@ func (p *Parser) parseAnnotations(stmts []Statement) []Statement {
 	return stmts
 }

-func parseAnnotations(comments []*Comment) ([]*Annotations, Errors) {
+func parseAnnotations(comments []*Comment) (stmts []*Annotations, errs Errors) {
+	numBlocks := CountFunc(comments, isMetadataComment)
+	if numBlocks == 0 {
+		return nil, nil
+	}

-	var hint = []byte("METADATA")
-	var curr *metadataParser
-	var blocks []*metadataParser
+	stmts = make([]*Annotations, 0, numBlocks)
+	mdp := metadataParserPool.Get()
+	if mdp.buf == nil {
+		mdp.buf = &bytes.Buffer{}
+	}

 	for i := range comments {
-		if curr != nil {
-			if comments[i].Location.Row == comments[i-1].Location.Row+1 && comments[i].Location.Col == 1 {
-				curr.Append(comments[i])
-				continue
+		if isMetadataComment(comments[i]) { // scan until end of block
+			mdp.Reset(comments[i].Location)
+			for i++; i < len(comments) && !blockBuster(comments[i], comments[i-1]); i++ {
+				mdp.Append(comments[i])
+			}
+
+			if a, err := mdp.Parse(); err != nil {
+				errs = append(errs, &Error{Code: ParseErr, Message: err.Error(), Location: mdp.loc})
+			} else {
+				stmts = append(stmts, a)
 			}
-			curr = nil
-		}
-		if bytes.HasPrefix(bytes.TrimSpace(comments[i].Text), hint) {
-			curr = newMetadataParser(comments[i].Location)
-			blocks = append(blocks, curr)
 		}
 	}

-	stmts := make([]*Annotations, 0, len(blocks))
-
-	var errs Errors
-	for _, b := range blocks {
-		if a, err := b.Parse(); err != nil {
-			errs = append(errs, &Error{
-				Code:     ParseErr,
-				Message:  err.Error(),
-				Location: b.loc,
-			})
-		} else {
-			stmts = append(stmts, a)
-		}
-	}
+	metadataParserPool.Put(mdp)

 	return stmts, errs
 }

+func isMetadataComment(c *Comment) bool {
+	return c.Location.Col == 1 && bytes.HasPrefix(bytes.TrimSpace(c.Text), metadataBytes)
+}
+
+func blockBuster(curr, prev *Comment) bool { // or endOfBlock, but the name was too good to pass up
+	return curr.Location.Col != 1 || curr.Location.Row-1 != prev.Location.Row
+}
+
 func (p *Parser) parsePackage() *Package {
 	if p.s.tok != tokens.Package {
 		return nil
@@ -2745,13 +2751,17 @@ type rawAnnotation struct {
 }

 type metadataParser struct {
-	buf      *bytes.Buffer
 	comments []*Comment
+	buf      *bytes.Buffer
 	loc      *location.Location
 }

-func newMetadataParser(loc *Location) *metadataParser {
-	return &metadataParser{loc: loc, buf: bytes.NewBuffer(nil)}
+func (b *metadataParser) Reset(loc *location.Location) {
+	b.comments = b.comments[:0]
+	b.loc = loc
+	if b.buf != nil {
+		b.buf.Reset()
+	}
 }

 func (b *metadataParser) Append(c *Comment) {
@@ -2762,14 +2772,12 @@ func (b *metadataParser) Append(c *Comment) {

 var yamlLineErrRegex = regexp.MustCompile(`^yaml:(?: unmarshal errors:[\n\s]*)? line ([[:digit:]]+):`)

-func (b *metadataParser) Parse() (*Annotations, error) {
-
-	var raw rawAnnotation
-
+func (b *metadataParser) Parse() (result *Annotations, err error) {
 	if len(bytes.TrimSpace(b.buf.Bytes())) == 0 {
 		return nil, errors.New("expected METADATA block, found whitespace")
 	}

+	var raw rawAnnotation
 	if err := yaml.Unmarshal(b.buf.Bytes(), &raw); err != nil {
 		var comment *Comment
 		match := yamlLineErrRegex.FindStringSubmatch(err.Error())
@@ -2792,13 +2800,14 @@ func (b *metadataParser) Parse() (*Annotations, error) {
 		return nil, augmentYamlError(err, b.comments)
 	}

-	var result Annotations
-	result.comments = b.comments
-	result.Scope = raw.Scope
-	result.Entrypoint = raw.Entrypoint
-	result.Title = raw.Title
-	result.Description = raw.Description
-	result.Organizations = raw.Organizations
+	result = &Annotations{
+		comments:      b.comments,
+		Scope:         raw.Scope,
+		Entrypoint:    raw.Entrypoint,
+		Title:         raw.Title,
+		Description:   raw.Description,
+		Organizations: raw.Organizations,
+	}

 	for _, v := range raw.RelatedResources {
 		rr, err := parseRelatedResource(v)
@@ -2880,32 +2889,30 @@ func (b *metadataParser) Parse() (*Annotations, error) {
 		result.Authors = append(result.Authors, author)
 	}

-	result.Custom = make(map[string]any)
-	for k, v := range raw.Custom {
-		val, err := convertYAMLMapKeyTypes(v, nil)
-		if err != nil {
-			return nil, err
+	if raw.Custom != nil {
+		result.Custom = make(map[string]any, len(raw.Custom))
+		for k, v := range raw.Custom {
+			if result.Custom[k], err = convertYAMLMapKeyTypes(v, nil); err != nil {
+				return nil, err
+			}
 		}
-		result.Custom[k] = val
 	}

 	result.Location = b.loc

 	// recreate original text of entire metadata block for location text attribute
-	sb := strings.Builder{}
-	sb.WriteString("# METADATA\n")
+	original := bytes.TrimSuffix(b.buf.Bytes(), newlineBytes)
+	numLines := bytes.Count(original, newlineBytes) + 1
+	preAlloc := len("# METADATA\n") + len(original) + numLines*2 // '# ' prefix added per line

-	lines := bytes.Split(b.buf.Bytes(), []byte{'\n'})
+	result.Location.Text = append(make([]byte, 0, preAlloc), "# METADATA\n"...)

-	for _, line := range lines[:len(lines)-1] {
-		sb.WriteString("# ")
-		sb.Write(line)
-		sb.WriteByte('\n')
+	for line := range bytes.SplitAfterSeq(original, newlineBytes) {
+		result.Location.Text = append(result.Location.Text, "# "...)
+		result.Location.Text = append(result.Location.Text, line...)
 	}

-	result.Location.Text = []byte(strings.TrimSuffix(sb.String(), "\n"))
-
-	return &result, nil
+	return result, err
 }

 // augmentYamlError augments a YAML error with hints intended to help the user figure out the cause of an otherwise
@@ -2914,30 +2921,29 @@ func (b *metadataParser) Parse() (*Annotations, error) {
 func augmentYamlError(err error, comments []*Comment) error {
 	// Adding hints for when key/value ':' separator isn't suffixed with a legal YAML space symbol
 	for _, comment := range comments {
-		txt := string(comment.Text)
-		parts := strings.Split(txt, ":")
-		if len(parts) > 1 {
-			parts = parts[1:]
-			var invalidSpaces []string
-			for partIndex, part := range parts {
-				if len(part) == 0 && partIndex == len(parts)-1 {
-					invalidSpaces = []string{}
-					break
-				}
+		if bytes.IndexByte(comment.Text, ':') == -1 {
+			continue
+		}
+		parts := bytes.Split(comment.Text, []byte{':'})[1:]

-				r, _ := utf8.DecodeRuneInString(part)
-				if r == ' ' || r == '\t' {
-					invalidSpaces = []string{}
-					break
-				}
+		var invalidSpaces []string
+		for partIndex, part := range parts {
+			if len(part) == 0 && partIndex == len(parts)-1 {
+				break
+			}

-				invalidSpaces = append(invalidSpaces, fmt.Sprintf("%+q", r))
-			}
-			if len(invalidSpaces) > 0 {
-				err = fmt.Errorf(
-					"%s\n  Hint: on line %d, symbol(s) %v immediately following a key/value separator ':' is not a legal yaml space character",
-					err.Error(), comment.Location.Row, invalidSpaces)
+			r, _ := utf8.DecodeRune(part)
+			if r == ' ' || r == '\t' {
+				break
 			}
+
+			invalidSpaces = append(invalidSpaces, fmt.Sprintf("%+q", r))
+		}
+		if len(invalidSpaces) > 0 {
+			err = fmt.Errorf(
+				"%s\n  Hint: on line %d, symbol(s) %v immediately following a"+
+					" key/value separator ':' is not a legal yaml space character",
+				err.Error(), comment.Location.Row, invalidSpaces)
 		}
 	}
 	return err
@@ -3055,7 +3061,7 @@ func parseAuthorString(s string) (*AuthorAnnotation, error) {
 	if len(trailing) >= len(emailPrefix)+len(emailSuffix) && strings.HasPrefix(trailing, emailPrefix) &&
 		strings.HasSuffix(trailing, emailSuffix) {
 		email = trailing[len(emailPrefix):]
-		email = email[0 : len(email)-len(emailSuffix)]
+		email = email[:len(email)-len(emailSuffix)]
 		namePartCount -= 1
 	}

@@ -259,3 +259,50 @@ func generateObjectOrSetStatement(depth int) string {
 	}
 	return s.String()
 }
+
+// _7136 ns/op     5744 B/op       40 allocs/op // parsing only "package p"
+// 34255 ns/op    40760 B/op      506 allocs/op // with annotations
+// 33261 ns/op    38841 B/op      499 allocs/op // pre-alloc location text buffer
+// 32817 ns/op    37319 B/op      487 allocs/op // use single metadataParser instance with reset
+func BenchmarkParseAnnotations(b *testing.B) {
+	policy := `
+# METADATA
+# title: Example Policy
+# description: Annotations are fun
+# organizations:
+#   - Open Policy Agent
+#   - Cloud Native Computing Foundation
+# related_resources:
+#   - https://www.openpolicyagent.org
+#   - https://www.cncf.io
+# authors:
+#   - Alice
+#   - Bob
+# custom:
+#   tags:
+#     - example
+#     - demo
+# scope: subpackages
+# schemas:
+#   - input: {"type": "object", "properties": {"user": {"type": "string"}}}
+package p
+`
+	for b.Loop() {
+		MustParseModuleWithOpts(policy, ParserOptions{ProcessAnnotation: true})
+	}
+}
+
+// 296108 ns/op	  882355 B/op	    7230 allocs/op
+// 279484 ns/op	  842892 B/op	    6315 allocs/op // pre-alloc location text buffer and reuse metadataParser
+func BenchmarkParseManyAnnotations(b *testing.B) {
+	sb := &strings.Builder{}
+	sb.WriteString("package p\n\n")
+	for i := range 100 {
+		fmt.Fprintf(sb, "# METADATA\n# title: annotation %d\n\n", i)
+	}
+	policy := strings.TrimSpace(sb.String()) + "\nrule if true\n"
+
+	for b.Loop() {
+		MustParseModuleWithOpts(policy, ParserOptions{ProcessAnnotation: true})
+	}
+}
@@ -0,0 +1,15 @@
+// Copyright 2026 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package ast
+
+// CountFunc counts the number of items in a slice S that satisfy predicate function f.
+func CountFunc[T any, S ~[]T](items S, f func(T) bool) (n int) {
+	for i := range items {
+		if f(items[i]) {
+			n++
+		}
+	}
+	return n
+}