Use sync.Pool for eval func objects (#8054)

These did not show up in escape analysis as they weren't pointers.
I guess you could say they... escaped the analysis. But some pprof
exploration revealed these were heavy allocation sources, and just
like that, we have almost 2 million allocations less running the
`regal lint bundle` benchmark. And a nice reduction in `B/op` too!

Also in this PR a simple change to the regexp built-in functions,
where a `RWMutex` is now used to reduce blocking, which also showed
up in the pprof data (although with marginal impact).

**Before**
```
756176916 ns/op	2614292780 B/op	53440051 allocs/op
```
**After**
```
737606792 ns/op	2469976856 B/op	51680605 allocs/op
```

Signed-off-by: Anders Eknert <[email protected]>

Author: anderseknert

v1/topdown/eval.go

@@ -119,9 +119,17 @@ type eval struct {
 	defined                     bool
 }
 
-type evp struct {
-	pool sync.Pool
-}
+type (
+	evp struct {
+		pool sync.Pool
+	}
+	evfp struct {
+		pool sync.Pool
+	}
+	evbp struct {
+		pool sync.Pool
+	}
+)
 
 func (ep *evp) Put(e *eval) {
 	ep.pool.Put(e)
@@ -131,14 +139,52 @@ func (ep *evp) Get() *eval {
 	return ep.pool.Get().(*eval)
 }
 
-var evalPool = evp{
-	pool: sync.Pool{
-		New: func() any {
-			return &eval{}
-		},
-	},
+func (ep *evfp) Put(e *evalFunc) {
+	if e != nil {
+		e.e, e.terms, e.ir = nil, nil, nil
+		ep.pool.Put(e)
+	}
+}
+
+func (ep *evfp) Get() *evalFunc {
+	return ep.pool.Get().(*evalFunc)
+}
+
+func (ep *evbp) Put(e *evalBuiltin) {
+	if e != nil {
+		e.e, e.bi, e.bctx, e.f, e.terms = nil, nil, nil, nil, nil
+		ep.pool.Put(e)
+	}
+}
+
+func (ep *evbp) Get() *evalBuiltin {
+	return ep.pool.Get().(*evalBuiltin)
 }
 
+var (
+	evalPool = &evp{
+		pool: sync.Pool{
+			New: func() any {
+				return &eval{}
+			},
+		},
+	}
+	evalFuncPool = &evfp{
+		pool: sync.Pool{
+			New: func() any {
+				return &evalFunc{}
+			},
+		},
+	}
+	evalBuiltinPool = &evbp{
+		pool: sync.Pool{
+			New: func() any {
+				return &evalBuiltin{}
+			},
+		},
+	}
+)
+
 func (e *eval) Run(iter evalIterator) error {
 	if !e.traceEnabled {
 		// avoid function literal escaping to heap if we don't need the trace
@@ -401,9 +447,11 @@ func (e *eval) evalExpr(iter evalIterator) error {
 		}
 		return nil
 	}
-	expr := e.query[e.index]
 
-	e.traceEval(expr)
+	expr := e.query[e.index]
+	if e.traceEnabled {
+		e.traceEval(expr)
+	}
 
 	if len(expr.With) > 0 {
 		return e.evalWith(iter)
@@ -521,7 +569,7 @@ func (e *eval) evalStep(iter evalIterator) error {
 		// generateVar inlined here to avoid extra allocations in hot path
 		rterm := ast.VarTerm(e.fmtVarTerm())
 		err = e.unify(terms, rterm, func() error {
-			if e.saveSet.Contains(rterm, e.bindings) {
+			if e.saveSet != nil && e.saveSet.Contains(rterm, e.bindings) {
 				return e.saveExpr(ast.NewExpr(rterm), e.bindings, func() error {
 					return iter(e)
 				})
@@ -888,7 +936,6 @@ func (e *eval) evalNotPartialSupport(negationID uint64, expr *ast.Expr, unknowns
 }
 
 func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
-
 	ref := terms[0].Value.(ast.Ref)
 
 	mock, mocked := e.functionMocks.Get(ref)
@@ -912,8 +959,8 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
 
 	if ref[0].Equal(ast.DefaultRootDocument) {
 		if mocked {
-			f := e.compiler.TypeEnv.Get(ref).(*types.Function)
-			return e.evalCallValue(f.Arity(), terms, mock, iter)
+			arity := e.compiler.TypeEnv.GetByRef(ref).(*types.Function).Arity()
+			return e.evalCallValue(arity, terms, mock, iter)
 		}
 
 		var ir *ast.IndexResult
@@ -928,11 +975,13 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
 			return err
 		}
 
-		eval := evalFunc{
-			e:     e,
-			terms: terms,
-			ir:    ir,
-		}
+		eval := evalFuncPool.Get()
+		defer evalFuncPool.Put(eval)
+
+		eval.e = e
+		eval.terms = terms
+		eval.ir = ir
+
 		return eval.eval(iter)
 	}
 
@@ -991,13 +1040,14 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
 		}
 	}
 
-	eval := evalBuiltin{
-		e:     e,
-		bi:    bi,
-		bctx:  bctx,
-		f:     f,
-		terms: terms[1:],
-	}
+	eval := evalBuiltinPool.Get()
+	defer evalBuiltinPool.Put(eval)
+
+	eval.e = e
+	eval.bi = bi
+	eval.bctx = bctx
+	eval.f = f
+	eval.terms = terms[1:]
 
 	return eval.eval(iter)
 }
@@ -1054,7 +1104,9 @@ func (e *eval) biunify(a, b *ast.Term, b1, b2 *bindings, iter unifyIterator) err
 		case ast.Var, ast.Ref, *ast.ArrayComprehension:
 			return e.biunifyValues(a, b, b1, b2, iter)
 		case *ast.Array:
-			return e.biunifyArrays(vA, vB, b1, b2, iter)
+			if vA.Len() == vB.Len() {
+				return e.biunifyArraysRec(vA, vB, b1, b2, iter, 0)
+			}
 		}
 	case ast.Object:
 		switch vB := b.Value.(type) {
@@ -1069,13 +1121,6 @@ func (e *eval) biunify(a, b *ast.Term, b1, b2 *bindings, iter unifyIterator) err
 	return nil
 }
 
-func (e *eval) biunifyArrays(a, b *ast.Array, b1, b2 *bindings, iter unifyIterator) error {
-	if a.Len() != b.Len() {
-		return nil
-	}
-	return e.biunifyArraysRec(a, b, b1, b2, iter, 0)
-}
-
 func (e *eval) biunifyArraysRec(a, b *ast.Array, b1, b2 *bindings, iter unifyIterator, idx int) error {
 	if idx == a.Len() {
 		return iter()
@@ -2052,8 +2097,7 @@ type evalFunc struct {
 	terms []*ast.Term
 }
 
-func (e evalFunc) eval(iter unifyIterator) error {
-
+func (e *evalFunc) eval(iter unifyIterator) error {
 	if e.ir.Empty() {
 		return nil
 	}
@@ -2065,13 +2109,13 @@ func (e evalFunc) eval(iter unifyIterator) error {
 		argCount = len(e.ir.Default.Head.Args)
 	}
 
-	if len(e.ir.Else) > 0 && e.e.unknown(e.e.query[e.e.index], e.e.bindings) {
-		// Partial evaluation of ordered rules is not supported currently. Save the
-		// expression and continue. This could be revisited in the future.
-		return e.e.saveCall(argCount, e.terms, iter)
-	}
-
 	if e.e.partial() {
+		if len(e.ir.Else) > 0 && e.e.unknown(e.e.query[e.e.index], e.e.bindings) {
+			// Partial evaluation of ordered rules is not supported currently. Save the
+			// expression and continue. This could be revisited in the future.
+			return e.e.saveCall(argCount, e.terms, iter)
+		}
+
 		var mustGenerateSupport bool
 
 		if defRule := e.ir.Default; defRule != nil {
@@ -2109,7 +2153,7 @@ func (e evalFunc) eval(iter unifyIterator) error {
 	return e.evalValue(iter, argCount, e.ir.EarlyExit)
 }
 
-func (e evalFunc) evalValue(iter unifyIterator, argCount int, findOne bool) error {
+func (e *evalFunc) evalValue(iter unifyIterator, argCount int, findOne bool) error {
 	var cacheKey ast.Ref
 	if !e.e.partial() {
 		var hit bool
@@ -2194,7 +2238,7 @@ func (e evalFunc) evalValue(iter unifyIterator, argCount int, findOne bool) erro
 	})
 }
 
-func (e evalFunc) evalCache(argCount int, iter unifyIterator) (ast.Ref, bool, error) {
+func (e *evalFunc) evalCache(argCount int, iter unifyIterator) (ast.Ref, bool, error) {
 	plen := len(e.terms)
 	if plen == argCount+2 { // func name + output = 2
 		plen -= 1
@@ -2226,7 +2270,7 @@ func (e evalFunc) evalCache(argCount int, iter unifyIterator) (ast.Ref, bool, er
 	return cacheKey, false, nil
 }
 
-func (e evalFunc) evalOneRule(iter unifyIterator, rule *ast.Rule, args []*ast.Term, cacheKey ast.Ref, prev *ast.Term, findOne bool) (*ast.Term, error) {
+func (e *evalFunc) evalOneRule(iter unifyIterator, rule *ast.Rule, args []*ast.Term, cacheKey ast.Ref, prev *ast.Term, findOne bool) (*ast.Term, error) {
 	child := evalPool.Get()
 	defer evalPool.Put(child)
 
@@ -2288,7 +2332,7 @@ func (e evalFunc) evalOneRule(iter unifyIterator, rule *ast.Rule, args []*ast.Te
 	return result, err
 }
 
-func (e evalFunc) partialEvalSupport(declArgsLen int, iter unifyIterator) error {
+func (e *evalFunc) partialEvalSupport(declArgsLen int, iter unifyIterator) error {
 	path := e.e.namespaceRef(e.terms[0].Value.(ast.Ref))
 
 	if !e.e.saveSupport.Exists(path) {
@@ -2316,7 +2360,7 @@ func (e evalFunc) partialEvalSupport(declArgsLen int, iter unifyIterator) error
 	return e.e.saveCall(declArgsLen, append([]*ast.Term{term}, e.terms[1:]...), iter)
 }
 
-func (e evalFunc) partialEvalSupportRule(rule *ast.Rule, path ast.Ref) error {
+func (e *evalFunc) partialEvalSupportRule(rule *ast.Rule, path ast.Ref) error {
 	child := evalPool.Get()
 	defer evalPool.Put(child)
 

v1/topdown/regex.go

@@ -16,11 +16,15 @@ import (
 	"github.com/open-policy-agent/opa/v1/topdown/builtins"
 )
 
-const regexCacheMaxSize = 100
-const regexInterQueryValueCacheHits = "rego_builtin_regex_interquery_value_cache_hits"
+const (
+	regexCacheMaxSize             = 100
+	regexInterQueryValueCacheHits = "rego_builtin_regex_interquery_value_cache_hits"
+)
 
-var regexpCacheLock = sync.Mutex{}
-var regexpCache map[string]*regexp.Regexp
+var (
+	regexpCacheLock = sync.RWMutex{}
+	regexpCache     = make(map[string]*regexp.Regexp)
+)
 
 func builtinRegexIsValid(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
 	if s, err := builtins.StringOperand(operands[0].Value, 1); err == nil {
@@ -103,7 +107,8 @@ func builtinRegexSplit(bctx BuiltinContext, operands []*ast.Term, iter func(*ast
 func getRegexp(bctx BuiltinContext, pat string) (*regexp.Regexp, error) {
 	if bctx.InterQueryBuiltinValueCache != nil {
 		// TODO: Use named cache
-		val, ok := bctx.InterQueryBuiltinValueCache.Get(ast.String(pat))
+		var key ast.Value = ast.String(pat)
+		val, ok := bctx.InterQueryBuiltinValueCache.Get(key)
 		if ok {
 			res, valid := val.(*regexp.Regexp)
 			if !valid {
@@ -120,42 +125,48 @@ func getRegexp(bctx BuiltinContext, pat string) (*regexp.Regexp, error) {
 		if err != nil {
 			return nil, err
 		}
-		bctx.InterQueryBuiltinValueCache.Insert(ast.String(pat), re)
+		bctx.InterQueryBuiltinValueCache.Insert(key, re)
 		return re, nil
 	}
 
-	regexpCacheLock.Lock()
-	defer regexpCacheLock.Unlock()
+	regexpCacheLock.RLock()
 	re, ok := regexpCache[pat]
+	numCached := len(regexpCache)
+	regexpCacheLock.RUnlock()
 	if !ok {
 		var err error
 		re, err = regexp.Compile(pat)
 		if err != nil {
 			return nil, err
 		}
-		if len(regexpCache) >= regexCacheMaxSize {
+
+		regexpCacheLock.Lock()
+		if numCached >= regexCacheMaxSize {
 			// Delete a (semi-)random key to make room for the new one.
 			for k := range regexpCache {
 				delete(regexpCache, k)
 				break
 			}
 		}
 		regexpCache[pat] = re
+		regexpCacheLock.Unlock()
 	}
 	return re, nil
 }
 
 func getRegexpTemplate(pat string, delimStart, delimEnd byte) (*regexp.Regexp, error) {
-	regexpCacheLock.Lock()
-	defer regexpCacheLock.Unlock()
+	regexpCacheLock.RLock()
 	re, ok := regexpCache[pat]
+	regexpCacheLock.RUnlock()
 	if !ok {
 		var err error
 		re, err = compileRegexTemplate(pat, delimStart, delimEnd)
 		if err != nil {
 			return nil, err
 		}
+		regexpCacheLock.Lock()
 		regexpCache[pat] = re
+		regexpCacheLock.Unlock()
 	}
 	return re, nil
 }
@@ -264,7 +275,6 @@ func builtinRegexReplace(bctx BuiltinContext, operands []*ast.Term, iter func(*a
 }
 
 func init() {
-	regexpCache = map[string]*regexp.Regexp{}
 	RegisterBuiltinFunc(ast.RegexIsValid.Name, builtinRegexIsValid)
 	RegisterBuiltinFunc(ast.RegexMatch.Name, builtinRegexMatch)
 	RegisterBuiltinFunc(ast.RegexMatchDeprecated.Name, builtinRegexMatch)