From b7511820f4ae68d8b53ba77da32dee6d0ed70dc2 Mon Sep 17 00:00:00 2001 From: Anders Eknert Date: Thu, 30 Jan 2025 13:03:22 +0100 Subject: [PATCH] Go API: Allow providing custom base cache (#7329) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The same way it's possible to provide a VirtualCache, it should be possible to bring your own BaseCache. In clients like Regal (when invoked via `regal lint`), the base data is only ever loaded once and doesn't change later. Yet currently, each evaluation (of which there are hundreds) will have a new cache instantiated, leading to unnecessary cache misses. Since our inmem-store is AST-based, we could also try to have the base cache tap directly into the store for a 100% hit ratio, avoiding the more costly storage queries. But that's still untested this point 🤓 (The tiny changes in resolver.go are unrelated to this feature but fix a perf issue I noticed last night as I was testing that functionality. Too small fix to warrant a PR of its own.) Signed-off-by: Anders Eknert --- v1/rego/rego.go | 27 ++++++++++++++++----------- v1/topdown/cache.go | 6 ++++++ v1/topdown/eval.go | 2 +- v1/topdown/query.go | 26 ++++++++++++++++++++++++-- v1/topdown/resolver.go | 11 +++++++++-- 5 files changed, 56 insertions(+), 16 deletions(-) diff --git a/v1/rego/rego.go b/v1/rego/rego.go index 9a7dfbf2a9..daf322e0f4 100644 --- a/v1/rego/rego.go +++ b/v1/rego/rego.go @@ -11,6 +11,7 @@ import ( "errors" "fmt" "io" + "maps" "strings" "time" @@ -128,6 +129,7 @@ type EvalContext struct { capabilities *ast.Capabilities strictBuiltinErrors bool virtualCache topdown.VirtualCache + baseCache topdown.BaseCache } func (e *EvalContext) RawInput() *interface{} { @@ -365,14 +367,22 @@ func EvalPrintHook(ph print.Hook) EvalOption { } } -// EvalVirtualCache sets the topdown.VirtualCache to use for evaluation. This is -// optional, and if not set, the default cache is used. +// EvalVirtualCache sets the topdown.VirtualCache to use for evaluation. +// This is optional, and if not set, the default cache is used. func EvalVirtualCache(vc topdown.VirtualCache) EvalOption { return func(e *EvalContext) { e.virtualCache = vc } } +// EvalBaseCache sets the topdown.BaseCache to use for evaluation. +// This is optional, and if not set, the default cache is used. +func EvalBaseCache(bc topdown.BaseCache) EvalOption { + return func(e *EvalContext) { + e.baseCache = bc + } +} + // EvalNondeterministicBuiltins causes non-deterministic builtins to be evalued // during partial evaluation. This is needed to pull in external data, or validate // a JWT, during PE, so that the result informs what queries are returned. @@ -2183,7 +2193,8 @@ func (r *Rego) eval(ctx context.Context, ectx *EvalContext) (ResultSet, error) { WithSeed(ectx.seed). WithPrintHook(ectx.printHook). WithDistributedTracingOpts(r.distributedTacingOpts). - WithVirtualCache(ectx.virtualCache) + WithVirtualCache(ectx.virtualCache). + WithBaseCache(ectx.baseCache) if !ectx.time.IsZero() { q = q.WithTime(ectx.time) @@ -2895,14 +2906,8 @@ func (r *Rego) planQuery(queries []ast.Body, evalQueryType queryType) (*ir.Polic } decls := make(map[string]*ast.Builtin, len(r.builtinDecls)+len(ast.BuiltinMap)) - - for k, v := range ast.BuiltinMap { - decls[k] = v - } - - for k, v := range r.builtinDecls { - decls[k] = v - } + maps.Copy(decls, ast.BuiltinMap) + maps.Copy(decls, r.builtinDecls) const queryName = "eval" // NOTE(tsandall): the query name is arbitrary diff --git a/v1/topdown/cache.go b/v1/topdown/cache.go index 8767a69578..9d2416b7bd 100644 --- a/v1/topdown/cache.go +++ b/v1/topdown/cache.go @@ -32,6 +32,12 @@ type VirtualCache interface { Keys() []ast.Ref } +// BaseCache defines the interface for a cache that stores cached base documents, i.e. data. +type BaseCache interface { + Get(ast.Ref) ast.Value + Put(ast.Ref, ast.Value) +} + type virtualCache struct { stack []*virtualCacheElem } diff --git a/v1/topdown/eval.go b/v1/topdown/eval.go index bda2bd3a1c..9a1facea16 100644 --- a/v1/topdown/eval.go +++ b/v1/topdown/eval.go @@ -72,6 +72,7 @@ type eval struct { store storage.Store txn storage.Transaction virtualCache VirtualCache + baseCache BaseCache interQueryBuiltinCache cache.InterQueryCache interQueryBuiltinValueCache cache.InterQueryValueCache printHook print.Hook @@ -80,7 +81,6 @@ type eval struct { parent *eval caller *eval bindings *bindings - baseCache *baseCache compiler *ast.Compiler input *ast.Term data *ast.Term diff --git a/v1/topdown/query.go b/v1/topdown/query.go index a008517cca..28af1e4bd4 100644 --- a/v1/topdown/query.go +++ b/v1/topdown/query.go @@ -62,6 +62,7 @@ type Query struct { printHook print.Hook tracingOpts tracing.Options virtualCache VirtualCache + baseCache BaseCache } // Builtin represents a built-in function that queries can call. @@ -314,6 +315,13 @@ func (q *Query) WithVirtualCache(vc VirtualCache) *Query { return q } +// WithBaseCache sets the BaseCache to use during evaluation. This is +// optional, and if not set, the default cache is used. +func (q *Query) WithBaseCache(bc BaseCache) *Query { + q.baseCache = bc + return q +} + // WithNondeterministicBuiltins causes non-deterministic builtins to be evalued // during partial evaluation. This is needed to pull in external data, or validate // a JWT, during PE, so that the result informs what queries are returned. @@ -353,6 +361,13 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support [] vc = NewVirtualCache() } + var bc BaseCache + if q.baseCache != nil { + bc = q.baseCache + } else { + bc = newBaseCache() + } + e := &eval{ ctx: ctx, metrics: q.metrics, @@ -366,7 +381,7 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support [] bindings: b, compiler: q.compiler, store: q.store, - baseCache: newBaseCache(), + baseCache: bc, targetStack: newRefStack(), txn: q.txn, input: q.input, @@ -544,6 +559,13 @@ func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error { vc = NewVirtualCache() } + var bc BaseCache + if q.baseCache != nil { + bc = q.baseCache + } else { + bc = newBaseCache() + } + e := &eval{ ctx: ctx, metrics: q.metrics, @@ -557,7 +579,7 @@ func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error { bindings: newBindings(0, q.instr), compiler: q.compiler, store: q.store, - baseCache: newBaseCache(), + baseCache: bc, targetStack: newRefStack(), txn: q.txn, input: q.input, diff --git a/v1/topdown/resolver.go b/v1/topdown/resolver.go index 170e6e6402..3620168874 100644 --- a/v1/topdown/resolver.go +++ b/v1/topdown/resolver.go @@ -48,7 +48,11 @@ func (t *resolverTrie) Resolve(e *eval, ref ast.Ref) (ast.Value, error) { Input: e.input, Metrics: e.metrics, } - e.traceWasm(e.query[e.index], &in.Ref) + if e.traceEnabled { + // avoid leaking pointer if trace is disabled + cpy := in.Ref + e.traceWasm(e.query[e.index], &cpy) + } if e.data != nil { return nil, errInScopeWithStmt } @@ -75,7 +79,10 @@ func (t *resolverTrie) Resolve(e *eval, ref ast.Ref) (ast.Value, error) { func (t *resolverTrie) mktree(e *eval, in resolver.Input) (ast.Value, error) { if t.r != nil { - e.traceWasm(e.query[e.index], &in.Ref) + if e.traceEnabled { + cpy := in.Ref + e.traceWasm(e.query[e.index], &cpy) + } if e.data != nil { return nil, errInScopeWithStmt }