From c122df1868778b422df41eabfc7e8b7c581f3efd Mon Sep 17 00:00:00 2001 From: Anders Eknert Date: Wed, 8 Oct 2025 14:16:40 +0200 Subject: [PATCH] Performance improvements in formatter (#7967) - Use new `util.SlicePool` to avoid cost of temporary slices in `formatTerm` - Add `SkipDefensiveCopying` option and enable it for all `Source`* functions ``` 676179 ns/op 995204 B/op 8850 allocs/op // regression in 0fb7526 513570 ns/op 378787 B/op 8775 allocs/op // addressed regression with sync.Pool 481681 ns/op 352130 B/op 7954 allocs/op // new util.NewSlicePool using only pointers 365116 ns/op 160528 B/op 2098 allocs/op // new SkipDefensiveCopying option ``` Signed-off-by: Anders Eknert --- v1/format/format.go | 96 +++++++++++++++++++++++++------------ v1/format/format_test.go | 24 +++++----- v1/util/performance.go | 47 ++++++++++++++++++ v1/util/performance_test.go | 9 ++++ 4 files changed, 135 insertions(+), 41 deletions(-) diff --git a/v1/format/format.go b/v1/format/format.go index 4867594905..75514d39c0 100644 --- a/v1/format/format.go +++ b/v1/format/format.go @@ -18,6 +18,20 @@ import ( "github.com/open-policy-agent/opa/internal/future" "github.com/open-policy-agent/opa/v1/ast" "github.com/open-policy-agent/opa/v1/types" + "github.com/open-policy-agent/opa/v1/util" +) + +// defaultLocationFile is the file name used in `Ast()` for terms +// without a location, as could happen when pretty-printing the +// results of partial eval. +const defaultLocationFile = "__format_default__" + +var ( + elseVar ast.Value = ast.Var("else") + + expandedConst = ast.NewBody(ast.NewExpr(ast.InternedTerm(true))) + commentsSlicePool = util.NewSlicePool[*ast.Comment](50) + varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$") ) // Opts lets you control the code formatting via `AstWithOpts()`. @@ -38,6 +52,11 @@ type Opts struct { // Imports are only removed if [Opts.RegoVersion] makes them redundant. DropV0Imports bool + // SkipDefensiveCopying, if true, will avoid deep-copying the AST before formatting it. + // This is true by default for all Source* functions, but false by default for Ast* functions, + // as some formatting operations may otherwise mutate the AST. + SkipDefensiveCopying bool + Capabilities *ast.Capabilities } @@ -48,16 +67,11 @@ func (o Opts) effectiveRegoVersion() ast.RegoVersion { return o.RegoVersion } -// defaultLocationFile is the file name used in `Ast()` for terms -// without a location, as could happen when pretty-printing the -// results of partial eval. -const defaultLocationFile = "__format_default__" - // Source formats a Rego source file. The bytes provided must describe a complete // Rego module. If they don't, Source will return an error resulting from the attempt // to parse the bytes. func Source(filename string, src []byte) ([]byte, error) { - return SourceWithOpts(filename, src, Opts{}) + return SourceWithOpts(filename, src, Opts{SkipDefensiveCopying: true}) } func SourceWithOpts(filename string, src []byte, opts Opts) ([]byte, error) { @@ -72,6 +86,9 @@ func SourceWithOpts(filename string, src []byte, opts Opts) ([]byte, error) { parserOpts.RegoVersion = ast.RegoV1 } + // Copying the node does not make sense when both input and output are bytes. + opts.SkipDefensiveCopying = true + if parserOpts.RegoVersion == ast.RegoUndefined { parserOpts.RegoVersion = ast.DefaultRegoVersion } @@ -166,7 +183,9 @@ func AstWithOpts(x any, opts Opts) ([]byte, error) { // The node has to be deep copied because it may be mutated below. Alternatively, // we could avoid the copy by checking if mutation will occur first. For now, // since format is not latency sensitive, just deep copy in all cases. - x = ast.Copy(x) + if !opts.SkipDefensiveCopying { + x = ast.Copy(x) + } wildcards := map[ast.Var]*ast.Term{} @@ -233,10 +252,11 @@ func AstWithOpts(x any, opts Opts) ([]byte, error) { } case *ast.Rule: - if len(n.Head.Ref()) > 2 { + headLen := len(n.Head.Ref()) + if headLen > 2 { o.refHeads = true } - if len(n.Head.Ref()) == 2 && n.Head.Key != nil && n.Head.Value == nil { // p.q contains "x" + if headLen == 2 && n.Head.Key != nil && n.Head.Value == nil { // p.q contains "x" o.refHeads = true } } @@ -339,6 +359,7 @@ func AstWithOpts(x any, opts Opts) ([]byte, error) { if len(w.errs) > 0 { return nil, w.errs } + return squashTrailingNewlines(w.buf.Bytes()), nil } @@ -545,8 +566,6 @@ func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) ([]*ast. return comments, nil } -var expandedConst = ast.NewBody(ast.NewExpr(ast.InternedTerm(true))) - func (w *writer) groupableOneLiner(rule *ast.Rule) bool { // Location required to determine if two rules are adjacent in the policy. // If not, we respect line breaks between rules. @@ -667,8 +686,6 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) return comments, nil } -var elseVar ast.Value = ast.Var("else") - func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) ([]*ast.Comment, error) { // If there was nothing else on the line before the "else" starts // then preserve this style of else block, otherwise it will be @@ -1127,18 +1144,33 @@ func (w *writer) writeWith(with *ast.With, comments []*ast.Comment, indented boo return comments, nil } +// saveComments saves a copy of the comments slice in a pooled slice to and returns it. +// This is to avoid having to create a new slice every time we need to save comments. +// The caller is responsible for putting the slice back in the pool when done. +func saveComments(comments []*ast.Comment) *[]*ast.Comment { + cmlen := len(comments) + saved := commentsSlicePool.Get(cmlen) + + copy(*saved, comments) + + return saved +} + func (w *writer) writeTerm(term *ast.Term, comments []*ast.Comment) ([]*ast.Comment, error) { - currentComments := make([]*ast.Comment, len(comments)) - copy(currentComments, comments) + if len(comments) == 0 { + return w.writeTermParens(false, term, comments) + } currentLen := w.buf.Len() + currentComments := saveComments(comments) + defer commentsSlicePool.Put(currentComments) comments, err := w.writeTermParens(false, term, comments) if err != nil { if errors.As(err, &unexpectedCommentError{}) { w.buf.Truncate(currentLen) - comments, uErr := w.writeUnformatted(term.Location, currentComments) + comments, uErr := w.writeUnformatted(term.Location, *currentComments) if uErr != nil { return nil, uErr } @@ -1156,16 +1188,16 @@ func (w *writer) writeUnformatted(location *ast.Location, currentComments []*ast return nil, errors.New("original unformatted text is empty") } - rawRule := string(location.Text) - rowNum := len(strings.Split(rawRule, "\n")) + rowNum := bytes.Count(location.Text, []byte{'\n'}) + 1 - w.write(string(location.Text)) + w.writeBytes(location.Text) comments := make([]*ast.Comment, 0, len(currentComments)) for _, c := range currentComments { // if there is a body then wait to write the last comment if w.writeCommentOnFinalLine && c.Location.Row == location.Row+rowNum-1 { - w.write(" " + string(c.Location.Text)) + w.write(" ") + w.writeBytes(c.Location.Text) continue } @@ -1227,19 +1259,19 @@ func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Co case ast.String: if term.Location.Text[0] == '`' { // To preserve raw strings, we need to output the original text, - w.write(string(term.Location.Text)) + w.writeBytes(term.Location.Text) } else { // x.String() cannot be used by default because it can change the input string "\u0000" to "\x00" - var after, quote string + var after, quote []byte var found bool // term.Location.Text could contain the prefix `else :=`, remove it switch term.Location.Text[len(term.Location.Text)-1] { case '"': - quote = "\"" - _, after, found = strings.Cut(string(term.Location.Text), quote) + quote = []byte{'"'} + _, after, found = bytes.Cut(term.Location.Text, quote) case '`': - quote = "`" - _, after, found = strings.Cut(string(term.Location.Text), quote) + quote = []byte{'`'} + _, after, found = bytes.Cut(term.Location.Text, quote) } if !found { @@ -1247,7 +1279,8 @@ func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Co // e.g. partial_set.y to partial_set["y"] w.write(x.String()) } else { - w.write(quote + after) + w.writeBytes(quote) + w.writeBytes(after) } } @@ -1310,8 +1343,6 @@ func (w *writer) writeBracketed(str string) { w.write("[" + str + "]") } -var varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$") - func (w *writer) writeRefStringPath(s ast.String, l *ast.Location) { str := string(s) if w.shouldBracketRefTerm(str, l) { @@ -2130,11 +2161,16 @@ func (w *writer) blankLine() { w.write("\n") } -// write the input string and writes it to the buffer. +// write writes string s to the buffer. func (w *writer) write(s string) { w.buf.WriteString(s) } +// writeBytes writes []byte b to the buffer. +func (w *writer) writeBytes(b []byte) { + w.buf.Write(b) +} + // writeLine writes the string on a newly started line, then terminate the line. func (w *writer) writeLine(s string) { if !w.inline { diff --git a/v1/format/format_test.go b/v1/format/format_test.go index 16d8a030f9..97994b2325 100644 --- a/v1/format/format_test.go +++ b/v1/format/format_test.go @@ -1004,24 +1004,26 @@ func TestFormatKeywordsInRefs(t *testing.T) { } } -// 382 3064960 ns/op 4573131 B/op 26266 allocs/op // no optimizations -// 685 1737719 ns/op 1972193 B/op 14160 allocs/op // pre-allocate partitionComments -// 708 1674343 ns/op 1916700 B/op 11556 allocs/op // static memberRef & memberWithKeyRef -// 746 1594546 ns/op 1882652 B/op 10644 allocs/op // various minor fixes -// 1250 853508 ns/op 441730 B/op 8895 allocs/op // partitionComments early return if unchanged -// 1396 812859 ns/op 362651 B/op 8811 allocs/op // partitionComments reuse backing array +// 3064960 ns/op 4573131 B/op 26266 allocs/op // no optimizations +// 1737719 ns/op 1972193 B/op 14160 allocs/op // pre-allocate partitionComments +// 1674343 ns/op 1916700 B/op 11556 allocs/op // static memberRef & memberWithKeyRef +// 1594546 ns/op 1882652 B/op 10644 allocs/op // various minor fixes +// _853508 ns/op 441730 B/op 8895 allocs/op // partitionComments early return if unchanged +// _812859 ns/op 362651 B/op 8811 allocs/op // partitionComments reuse backing array +// _676179 ns/op 995204 B/op 8850 allocs/op // regression in 0fb7526 +// _513570 ns/op 378787 B/op 8775 allocs/op // addressed regression with sync.Pool +// _481681 ns/op 352130 B/op 7954 allocs/op // new util.NewSlicePool using only pointers +// _365116 ns/op 160528 B/op 2098 allocs/op // new SkipDefensiveCopying option func BenchmarkFormatLargePolicy(b *testing.B) { contents, err := os.ReadFile("testdata/bench.rego") if err != nil { b.Fatalf("Failed to read rego source: %v", err) } module := ast.MustParseModule(string(contents)) + opts := Opts{RegoVersion: ast.RegoV1, SkipDefensiveCopying: true} - b.ResetTimer() - - for range b.N { - _, err := AstWithOpts(module, Opts{RegoVersion: ast.RegoV1}) - if err != nil { + for b.Loop() { + if _, err := AstWithOpts(module, opts); err != nil { b.Fatal(err) } } diff --git a/v1/util/performance.go b/v1/util/performance.go index 7be2e69c64..09a8e74011 100644 --- a/v1/util/performance.go +++ b/v1/util/performance.go @@ -4,6 +4,7 @@ import ( "math" "slices" "strings" + "sync" "unsafe" ) @@ -91,3 +92,49 @@ func SplitMap[T any](text string, delim string, fn func(string) T) []T { return sl } + +// SlicePool is a pool for (pointers to) slices of type T. +// It uses sync.Pool to pool the slices, and grows them as needed. +type SlicePool[T any] struct { + pool sync.Pool +} + +// NewSlicePool creates a new SlicePool for slices of type T with the given initial length. +// This number is only a hint, as the slices will grow as needed. For best performance, store +// slices of similar lengths in the same pool. +func NewSlicePool[T any](length int) *SlicePool[T] { + return &SlicePool[T]{ + pool: sync.Pool{ + New: func() any { + s := make([]T, length) + return &s + }, + }, + } +} + +// Get returns a pointer to a slice of type T with the given length +// from the pool. The slice capacity will grow as needed to accommodate +// the requested length. The returned slice will have all its elements +// set to the zero value of T. Returns a pointer to avoid allocating. +func (sp *SlicePool[T]) Get(length int) *[]T { + s := sp.pool.Get().(*[]T) + d := *s + + if cap(d) < length { + d = slices.Grow(d, length) + } + + d = d[:length] // reslice to requested length, while keeping capacity + + clear(d) + + *s = d + + return s +} + +// Put returns a pointer to a slice of type T to the pool. +func (sp *SlicePool[T]) Put(s *[]T) { + sp.pool.Put(s) +} diff --git a/v1/util/performance_test.go b/v1/util/performance_test.go index b3a70eea3a..9bac4609bb 100644 --- a/v1/util/performance_test.go +++ b/v1/util/performance_test.go @@ -59,6 +59,15 @@ func BenchmarkSplitMap(b *testing.B) { }) } +// Zero allocations +func BenchmarkSlicePoolGetPut(b *testing.B) { + sp := NewSlicePool[int](4) + for b.Loop() { + s := sp.Get(4) + sp.Put(s) + } +} + func mustAtoi(s string) int { v, _ := strconv.Atoi(s) return v