Go API: Allow providing custom base cache (#7329)

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 <anders@styra.com>
This commit is contained in:
Anders Eknert
2025-01-30 13:03:22 +01:00
committed by GitHub
parent b942136a4a
commit b7511820f4
5 changed files with 56 additions and 16 deletions
+16 -11
View File
@@ -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
+6
View File
@@ -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
}
+1 -1
View File
@@ -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
+24 -2
View File
@@ -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,
+9 -2
View File
@@ -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
}