mirror of
https://github.com/open-policy-agent/opa.git
synced 2026-08-12 19:32:48 -06:00
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:
+16
-11
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user