eval+rego: Support caching output of non-deterministic builtins. (#4926)

This commit includes evaluator support for an opt-in, non-deterministic
builtins caching system, designed to help with future replay of decision
logs.

The cache allows early-exit in the evaluator if the builtin is
non-deterministic, and has already cached a result. Since the cache can
be pre-populated by `rego` module users, this should make offline policy
testing and future work around decision replay more straightforward.

Fixes: #1514

Signed-off-by: Philip Conrad <philipaconrad@gmail.com>

ast/builtins.go

@@ -1288,6 +1288,7 @@ var StringReverse = &Builtin{
  */
 
 // RandIntn returns a random number 0 - n
+// Marked non-deterministic because it relies on RNG internally.
 var RandIntn = &Builtin{
 	Name:        "rand.intn",
 	Description: "Returns a random integer between `0` and `n` (`n` exlusive). If `n` is `0`, then `y` is always `0`. For any given argument pair (`str`, `n`), the output will be consistent throughout a query evaluation.",
@@ -1298,7 +1299,8 @@ var RandIntn = &Builtin{
 		),
 		types.Named("y", types.N).Description("random integer in the range `[0, abs(n))`"),
 	),
-	Categories: number,
+	Categories:       number,
+	Nondeterministic: true,
 }
 
 var NumbersRange = &Builtin{
@@ -1353,6 +1355,7 @@ unit is optional and omitting it wil give the same result (e.g. Mi and MiB).`,
  */
 
 // UUIDRFC4122 returns a version 4 UUID string.
+// Marked non-deterministic because it relies on RNG internally.
 var UUIDRFC4122 = &Builtin{
 	Name:        "uuid.rfc4122",
 	Description: "Returns a new UUIDv4.",
@@ -1362,6 +1365,7 @@ var UUIDRFC4122 = &Builtin{
 		),
 		types.Named("output", types.S).Description("a version 4 UUID; for any given `k`, the output will be consistent throughout a query evaluation"),
 	),
+	Nondeterministic: true,
 }
 
 /**
@@ -1999,6 +2003,7 @@ var JWTVerifyHS512 = &Builtin{
 	Categories: tokensCat,
 }
 
+// Marked non-deterministic because it relies on time internally.
 var JWTDecodeVerify = &Builtin{
 	Name: "io.jwt.decode_verify",
 	Description: `Verifies a JWT signature under parameterized constraints and decodes the claims if it is valid.
@@ -2014,11 +2019,13 @@ Supports the following algorithms: HS256, HS384, HS512, RS256, RS384, RS512, ES2
 			types.NewObject(nil, types.NewDynamicProperty(types.A, types.A)),
 		}, nil)).Description("`[valid, header, payload]`:  if the input token is verified and meets the requirements of `constraints` then `valid` is `true`; `header` and `payload` are objects containing the JOSE header and the JWT claim set; otherwise, `valid` is `false`, `header` and `payload` are `{}`"),
 	),
-	Categories: tokensCat,
+	Categories:       tokensCat,
+	Nondeterministic: true,
 }
 
 var tokenSign = category("tokensign")
 
+// Marked non-deterministic because it relies on RNG internally.
 var JWTEncodeSignRaw = &Builtin{
 	Name:        "io.jwt.encode_sign_raw",
 	Description: "Encodes and optionally signs a JSON Web Token.",
@@ -2030,9 +2037,11 @@ var JWTEncodeSignRaw = &Builtin{
 		),
 		types.Named("output", types.S).Description("signed JWT"),
 	),
-	Categories: tokenSign,
+	Categories:       tokenSign,
+	Nondeterministic: true,
 }
 
+// Marked non-deterministic because it relies on RNG internally.
 var JWTEncodeSign = &Builtin{
 	Name:        "io.jwt.encode_sign",
 	Description: "Encodes and optionally signs a JSON Web Token. Inputs are taken as objects, not encoded strings (see `io.jwt.encode_sign_raw`).",
@@ -2044,13 +2053,15 @@ var JWTEncodeSign = &Builtin{
 		),
 		types.Named("output", types.S).Description("signed JWT"),
 	),
-	Categories: tokenSign,
+	Categories:       tokenSign,
+	Nondeterministic: true,
 }
 
 /**
  * Time
  */
 
+// Marked non-deterministic because it relies on time directly.
 var NowNanos = &Builtin{
 	Name:        "time.now_ns",
 	Description: "Returns the current time since epoch in nanoseconds.",
@@ -2058,6 +2069,7 @@ var NowNanos = &Builtin{
 		nil,
 		types.Named("now", types.N).Description("nanoseconds since epoch"),
 	),
+	Nondeterministic: true,
 }
 
 var ParseNanos = &Builtin{
@@ -2480,6 +2492,7 @@ var TypeNameBuiltin = &Builtin{
  * HTTP Request
  */
 
+// Marked non-deterministic because HTTP request results can be non-deterministic.
 var HTTPSend = &Builtin{
 	Name:        "http.send",
 	Description: "Returns a HTTP response to the given HTTP request.",
@@ -2489,6 +2502,7 @@ var HTTPSend = &Builtin{
 		),
 		types.Named("response", types.NewObject(nil, types.NewDynamicProperty(types.A, types.A))),
 	),
+	Nondeterministic: true,
 }
 
 /**
@@ -2610,6 +2624,7 @@ var RegoMetadataRule = &Builtin{
  * OPA
  */
 
+// Marked non-deterministic because of unpredictable config/environment-dependent results.
 var OPARuntime = &Builtin{
 	Name:        "opa.runtime",
 	Description: "Returns an object that describes the runtime environment where OPA is deployed.",
@@ -2618,6 +2633,7 @@ var OPARuntime = &Builtin{
 		types.Named("output", types.NewObject(nil, types.NewDynamicProperty(types.S, types.A))).
 			Description("includes a `config` key if OPA was started with a configuration file; an `env` key containing the environment variables that the OPA process was started with; includes `version` and `commit` keys containing the version and build commit of OPA."),
 	),
+	Nondeterministic: true,
 }
 
 /**
@@ -2756,6 +2772,7 @@ var netCidrContainsMatchesOperandType = types.NewAny(
 	)),
 )
 
+// Marked non-deterministic because DNS resolution results can be non-deterministic.
 var NetLookupIPAddr = &Builtin{
 	Name:        "net.lookup_ip_addr",
 	Description: "Returns the set of IP addresses (both v4 and v6) that the passed-in `name` resolves to using the standard name resolution mechanisms available.",
@@ -2765,6 +2782,7 @@ var NetLookupIPAddr = &Builtin{
 		),
 		types.Named("addrs", types.NewSet(types.S)).Description("IP addresses (v4 and v6) that `name` resolves to"),
 	),
+	Nondeterministic: true,
 }
 
 /**
@@ -2966,10 +2984,11 @@ type Builtin struct {
 	// "minus" for example, is part of two categories: numbers and sets. (NOTE(sr): aspirational)
 	Categories []string `json:"categories,omitempty"`
 
-	Decl       *types.Function `json:"decl"`               // Built-in function type declaration.
-	Infix      string          `json:"infix,omitempty"`    // Unique name of infix operator. Default should be unset.
-	Relation   bool            `json:"relation,omitempty"` // Indicates if the built-in acts as a relation.
-	deprecated bool            // Indicates if the built-in has been deprecated.
+	Decl             *types.Function `json:"decl"`               // Built-in function type declaration.
+	Infix            string          `json:"infix,omitempty"`    // Unique name of infix operator. Default should be unset.
+	Relation         bool            `json:"relation,omitempty"` // Indicates if the built-in acts as a relation.
+	deprecated       bool            // Indicates if the built-in has been deprecated.
+	Nondeterministic bool            `json:"nondeterministic,omitempty"` // Indicates if the built-in returns non-deterministic results.
 }
 
 // category is a helper for specifying a Builtin's Categories
@@ -2982,6 +3001,11 @@ func (b *Builtin) IsDeprecated() bool {
 	return b.deprecated
 }
 
+// IsDeterministic returns true if the Builtin function returns non-deterministic results.
+func (b *Builtin) IsNondeterministic() bool {
+	return b.Nondeterministic
+}
+
 // Expr creates a new expression for the built-in with the given operands.
 func (b *Builtin) Expr(operands ...*Term) *Expr {
 	ts := make([]*Term, len(operands)+1)

capabilities.json

@@ -1284,7 +1284,8 @@
           "type": "object"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "indexof",
@@ -1493,7 +1494,8 @@
           "type": "array"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "io.jwt.encode_sign",
@@ -1537,7 +1539,8 @@
           "type": "string"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "io.jwt.encode_sign_raw",
@@ -1557,7 +1560,8 @@
           "type": "string"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "io.jwt.verify_es256",
@@ -2529,7 +2533,8 @@
           "type": "set"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "numbers.range",
@@ -2819,7 +2824,8 @@
           "type": "object"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "or",
@@ -2918,7 +2924,8 @@
           "type": "number"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "re_match",
@@ -3690,7 +3697,8 @@
           "type": "number"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "time.parse_duration_ns",
@@ -4091,7 +4099,8 @@
           "type": "string"
         },
         "type": "function"
-      }
+      },
+      "nondeterministic": true
     },
     {
       "name": "walk",

features/wasm/wasm.go

@@ -64,6 +64,7 @@ func (o *OPA) Eval(ctx context.Context, opts opa.EvalOpts) (*opa.Result, error)
 		Time:                   opts.Time,
 		Seed:                   opts.Seed,
 		InterQueryBuiltinCache: opts.InterQueryBuiltinCache,
+		NDBuiltinCache:         opts.NDBuiltinCache,
 		PrintHook:              opts.PrintHook,
 		Capabilities:           opts.Capabilities,
 	}

internal/rego/opa/options.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/metrics"
+	"github.com/open-policy-agent/opa/topdown/builtins"
 	"github.com/open-policy-agent/opa/topdown/cache"
 	"github.com/open-policy-agent/opa/topdown/print"
 )
@@ -23,6 +24,7 @@ type EvalOpts struct {
 	Time                   time.Time
 	Seed                   io.Reader
 	InterQueryBuiltinCache cache.InterQueryCache
+	NDBuiltinCache         builtins.NDBCache
 	PrintHook              print.Hook
 	Capabilities           *ast.Capabilities
 }

internal/wasm/sdk/internal/wasm/bindings.go

@@ -85,6 +85,7 @@ func (d *builtinDispatcher) Reset(ctx context.Context,
 	seed io.Reader,
 	ns time.Time,
 	iqbCache cache.InterQueryCache,
+	ndbCache builtins.NDBCache,
 	ph print.Hook,
 	capabilities *ast.Capabilities) {
 	if ns.IsZero() {
@@ -107,6 +108,7 @@ func (d *builtinDispatcher) Reset(ctx context.Context,
 		QueryID:                0,
 		ParentID:               0,
 		InterQueryBuiltinCache: iqbCache,
+		NDBuiltinCache:         ndbCache,
 		PrintHook:              ph,
 		Capabilities:           capabilities,
 	}

internal/wasm/sdk/internal/wasm/pool_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/open-policy-agent/opa/internal/wasm/sdk/internal/wasm"
 	wasm_util "github.com/open-policy-agent/opa/internal/wasm/util"
 	"github.com/open-policy-agent/opa/metrics"
+	"github.com/open-policy-agent/opa/topdown/builtins"
 	"github.com/open-policy-agent/opa/topdown/cache"
 	"github.com/open-policy-agent/opa/util"
 )
@@ -176,7 +177,7 @@ func ensurePoolResults(t *testing.T, ctx context.Context, testPool *wasm.Pool, p
 		toRelease = append(toRelease, vm)
 
 		cfg, _ := cache.ParseCachingConfig(nil)
-		result, err := vm.Eval(ctx, 0, input, metrics.New(), rand.New(rand.NewSource(0)), time.Now(), cache.NewInterQueryCache(cfg), nil, nil)
+		result, err := vm.Eval(ctx, 0, input, metrics.New(), rand.New(rand.NewSource(0)), time.Now(), cache.NewInterQueryCache(cfg), builtins.NDBCache{}, nil, nil)
 		if err != nil {
 			t.Fatalf("Unexpected error: %s", err)
 		}

internal/wasm/sdk/internal/wasm/vm.go

@@ -21,6 +21,7 @@ import (
 	"github.com/open-policy-agent/opa/internal/wasm/util"
 	"github.com/open-policy-agent/opa/metrics"
 	"github.com/open-policy-agent/opa/topdown"
+	"github.com/open-policy-agent/opa/topdown/builtins"
 	"github.com/open-policy-agent/opa/topdown/cache"
 	"github.com/open-policy-agent/opa/topdown/print"
 )
@@ -271,10 +272,11 @@ func (i *VM) Eval(ctx context.Context,
 	seed io.Reader,
 	ns time.Time,
 	iqbCache cache.InterQueryCache,
+	ndbCache builtins.NDBCache,
 	ph print.Hook,
 	capabilities *ast.Capabilities) ([]byte, error) {
 	if i.abiMinorVersion < int32(2) {
-		return i.evalCompat(ctx, entrypoint, input, metrics, seed, ns, iqbCache, ph, capabilities)
+		return i.evalCompat(ctx, entrypoint, input, metrics, seed, ns, iqbCache, ndbCache, ph, capabilities)
 	}
 
 	metrics.Timer("wasm_vm_eval").Start()
@@ -325,7 +327,7 @@ func (i *VM) Eval(ctx context.Context,
 	// make use of it (e.g. `http.send`); and it will spawn a go routine
 	// cancelling the builtins that use topdown.Cancel, when the context is
 	// cancelled.
-	i.dispatcher.Reset(ctx, seed, ns, iqbCache, ph, capabilities)
+	i.dispatcher.Reset(ctx, seed, ns, iqbCache, ndbCache, ph, capabilities)
 
 	metrics.Timer("wasm_vm_eval_call").Start()
 	resultAddr, err := i.evalOneOff(ctx, int32(entrypoint), i.dataAddr, inputAddr, inputLen, heapPtr)
@@ -354,6 +356,7 @@ func (i *VM) evalCompat(ctx context.Context,
 	seed io.Reader,
 	ns time.Time,
 	iqbCache cache.InterQueryCache,
+	ndbCache builtins.NDBCache,
 	ph print.Hook,
 	capabilities *ast.Capabilities) ([]byte, error) {
 	metrics.Timer("wasm_vm_eval").Start()
@@ -365,7 +368,7 @@ func (i *VM) evalCompat(ctx context.Context,
 	// make use of it (e.g. `http.send`); and it will spawn a go routine
 	// cancelling the builtins that use topdown.Cancel, when the context is
 	// cancelled.
-	i.dispatcher.Reset(ctx, seed, ns, iqbCache, ph, capabilities)
+	i.dispatcher.Reset(ctx, seed, ns, iqbCache, ndbCache, ph, capabilities)
 
 	err := i.setHeapState(ctx, i.evalHeapPtr)
 	if err != nil {

internal/wasm/sdk/opa/opa.go

@@ -17,6 +17,7 @@ import (
 	"github.com/open-policy-agent/opa/internal/wasm/sdk/opa/errors"
 	sdk_errors "github.com/open-policy-agent/opa/internal/wasm/sdk/opa/errors"
 	"github.com/open-policy-agent/opa/metrics"
+	"github.com/open-policy-agent/opa/topdown/builtins"
 	"github.com/open-policy-agent/opa/topdown/cache"
 	"github.com/open-policy-agent/opa/topdown/print"
 )
@@ -165,6 +166,7 @@ type EvalOpts struct {
 	Time                   time.Time
 	Seed                   io.Reader
 	InterQueryBuiltinCache cache.InterQueryCache
+	NDBuiltinCache         builtins.NDBCache
 	PrintHook              print.Hook
 	Capabilities           *ast.Capabilities
 }
@@ -190,8 +192,7 @@ func (o *OPA) Eval(ctx context.Context, opts EvalOpts) (*Result, error) {
 
 	defer o.pool.Release(instance, m)
 
-	result, err := instance.Eval(ctx, opts.Entrypoint, opts.Input, m, opts.Seed, opts.Time, opts.InterQueryBuiltinCache,
-		opts.PrintHook, opts.Capabilities)
+	result, err := instance.Eval(ctx, opts.Entrypoint, opts.Input, m, opts.Seed, opts.Time, opts.InterQueryBuiltinCache, opts.NDBuiltinCache, opts.PrintHook, opts.Capabilities)
 	if err != nil {
 		return nil, err
 	}

rego/rego.go

@@ -29,6 +29,7 @@ import (
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/storage/inmem"
 	"github.com/open-policy-agent/opa/topdown"
+	"github.com/open-policy-agent/opa/topdown/builtins"
 	"github.com/open-policy-agent/opa/topdown/cache"
 	"github.com/open-policy-agent/opa/topdown/print"
 	"github.com/open-policy-agent/opa/tracing"
@@ -115,6 +116,7 @@ type EvalContext struct {
 	indexing               bool
 	earlyExit              bool
 	interQueryBuiltinCache cache.InterQueryCache
+	ndBuiltinCache         builtins.NDBCache
 	resolvers              []refResolver
 	sortSets               bool
 	printHook              print.Hook
@@ -253,6 +255,14 @@ func EvalInterQueryBuiltinCache(c cache.InterQueryCache) EvalOption {
 	}
 }
 
+// EvalNDBuiltinCache sets the non-deterministic builtin cache that built-in functions can
+// use during evaluation.
+func EvalNDBuiltinCache(c builtins.NDBCache) EvalOption {
+	return func(e *EvalContext) {
+		e.ndBuiltinCache = c
+	}
+}
+
 // EvalResolver sets a Resolver for a specified ref path for this evaluation.
 func EvalResolver(ref ast.Ref, r resolver.Resolver) EvalOption {
 	return func(e *EvalContext) {
@@ -310,6 +320,7 @@ func (pq preparedQuery) newEvalContext(ctx context.Context, options []EvalOption
 		compiledQuery:    compiledQuery{},
 		indexing:         true,
 		earlyExit:        true,
+		ndBuiltinCache:   builtins.NDBCache{},
 		resolvers:        pq.r.resolvers,
 		printHook:        pq.r.printHook,
 		capabilities:     pq.r.capabilities,
@@ -508,6 +519,7 @@ type Rego struct {
 	bundles                map[string]*bundle.Bundle
 	skipBundleVerification bool
 	interQueryBuiltinCache cache.InterQueryCache
+	ndBuiltinCache         builtins.NDBCache
 	strictBuiltinErrors    bool
 	resolvers              []refResolver
 	schemaSet              *ast.SchemaSet
@@ -1015,6 +1027,13 @@ func InterQueryBuiltinCache(c cache.InterQueryCache) func(r *Rego) {
 	}
 }
 
+// NDBuiltinCache sets the non-deterministic builtins cache.
+func NDBuiltinCache(c builtins.NDBCache) func(r *Rego) {
+	return func(r *Rego) {
+		r.ndBuiltinCache = c
+	}
+}
+
 // StrictBuiltinErrors tells the evaluator to treat all built-in function errors as fatal errors.
 func StrictBuiltinErrors(yes bool) func(r *Rego) {
 	return func(r *Rego) {
@@ -1093,6 +1112,7 @@ func New(options ...func(r *Rego)) *Rego {
 		builtinDecls:    map[string]*ast.Builtin{},
 		builtinFuncs:    map[string]*topdown.Builtin{},
 		bundles:         map[string]*bundle.Bundle{},
+		ndBuiltinCache:  builtins.NDBCache{},
 	}
 
 	for _, option := range options {
@@ -1162,6 +1182,7 @@ func (r *Rego) Eval(ctx context.Context) (ResultSet, error) {
 		EvalInstrument(r.instrument),
 		EvalTime(r.time),
 		EvalInterQueryBuiltinCache(r.interQueryBuiltinCache),
+		EvalNDBuiltinCache(r.ndBuiltinCache),
 		EvalSeed(r.seed),
 	}
 
@@ -1235,6 +1256,7 @@ func (r *Rego) Partial(ctx context.Context) (*PartialQueries, error) {
 		EvalMetrics(r.metrics),
 		EvalInstrument(r.instrument),
 		EvalInterQueryBuiltinCache(r.interQueryBuiltinCache),
+		EvalNDBuiltinCache(r.ndBuiltinCache),
 	}
 
 	for _, t := range r.queryTracers {
@@ -1906,6 +1928,7 @@ func (r *Rego) eval(ctx context.Context, ectx *EvalContext) (ResultSet, error) {
 		WithIndexing(ectx.indexing).
 		WithEarlyExit(ectx.earlyExit).
 		WithInterQueryBuiltinCache(ectx.interQueryBuiltinCache).
+		WithNDBuiltinCache(ectx.ndBuiltinCache).
 		WithStrictBuiltinErrors(r.strictBuiltinErrors).
 		WithSeed(ectx.seed).
 		WithPrintHook(ectx.printHook).
@@ -1970,6 +1993,7 @@ func (r *Rego) evalWasm(ctx context.Context, ectx *EvalContext) (ResultSet, erro
 		Time:                   ectx.time,
 		Seed:                   ectx.seed,
 		InterQueryBuiltinCache: ectx.interQueryBuiltinCache,
+		NDBuiltinCache:         ectx.ndBuiltinCache,
 		PrintHook:              ectx.printHook,
 		Capabilities:           ectx.capabilities,
 	})
@@ -2181,6 +2205,7 @@ func (r *Rego) partial(ctx context.Context, ectx *EvalContext) (*PartialQueries,
 		WithSkipPartialNamespace(r.skipPartialNamespace).
 		WithShallowInlining(r.shallowInlining).
 		WithInterQueryBuiltinCache(ectx.interQueryBuiltinCache).
+		WithNDBuiltinCache(ectx.ndBuiltinCache).
 		WithStrictBuiltinErrors(r.strictBuiltinErrors).
 		WithSeed(ectx.seed).
 		WithPrintHook(ectx.printHook)

rego/rego_test.go

@@ -30,6 +30,7 @@ import (
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/storage/inmem"
 	"github.com/open-policy-agent/opa/topdown"
+	"github.com/open-policy-agent/opa/topdown/builtins"
 	"github.com/open-policy-agent/opa/topdown/cache"
 	"github.com/open-policy-agent/opa/types"
 	"github.com/open-policy-agent/opa/util"
@@ -2010,6 +2011,132 @@ func TestEvalWithInterQueryCache(t *testing.T) {
 	}
 }
 
+// We use http.send to ensure the NDBuiltinCache is involved.
+func TestEvalWithNDCache(t *testing.T) {
+	var requests []*http.Request
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requests = append(requests, r)
+		_, _ = w.Write([]byte(`{"x": 1}`))
+	}))
+	defer ts.Close()
+	query := fmt.Sprintf(`http.send({"method": "get", "url": "%s", "force_json_decode": true})`, ts.URL)
+
+	// Set up the ND cache, and put in some arbitrary constants for the first K/V pair.
+	arbitraryKey := ast.Number(strconv.Itoa(2015))
+	arbitraryValue := ast.String("First commit year")
+	ndBC := builtins.NDBCache{}
+	ndBC.Put("arbitrary_experiment", arbitraryKey, arbitraryValue)
+
+	// Query execution of http.send should add an entry to the NDBuiltinCache.
+	ctx := context.Background()
+	_, err := New(Query(query), NDBuiltinCache(ndBC)).Eval(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Check and make sure we got exactly 2x items back in the ND builtin cache.
+	// NDBuiltinsCache always has the structure: map[ast.String]map[ast.Array]ast.Value
+	if len(ndBC) != 2 {
+		t.Fatalf("Expected exactly 2 items in non-deterministic builtin cache. Found %d items.\n", len(ndBC))
+	}
+	// Check the cached k/v types for the HTTP section of the cache.
+	if cachedResults, ok := ndBC["http.send"]; ok {
+		err := cachedResults.Iter(func(k, v *ast.Term) error {
+			if _, ok := k.Value.(*ast.Array); !ok {
+				t.Fatalf("http.send failed to store Object key in the ND builtins cache")
+			}
+			if _, ok := v.Value.(ast.Object); !ok {
+				t.Fatalf("http.send failed to store Object value in the ND builtins cache")
+			}
+			return nil
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// Ensure our original arbitrary data in the cache was preserved.
+	if v, ok := ndBC.Get("arbitrary_experiment", arbitraryKey); ok {
+		if v != arbitraryValue {
+			t.Fatalf("Non-deterministic builtins cache value was mangled. Expected: %v, got: %v\n", arbitraryValue, v)
+		}
+	} else {
+		t.Fatal("Non-deterministic builtins cache lookup failed.")
+	}
+}
+
+func TestEvalWithPrebuiltNDCache(t *testing.T) {
+	query := "time.now_ns()"
+	ndBC := builtins.NDBCache{}
+
+	// Populate the cache for time.now_ns with an arbitrary timestamp.
+	timeValue, err := time.Parse("2006-01-02T15:04:05Z", "2015-12-28T14:08:25Z")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Timestamp ns value will be: 1451311705000000000
+	ndBC.Put("time.now_ns", ast.NewArray(), ast.Number(json.Number(strconv.FormatInt(timeValue.UnixNano(), 10))))
+	// time.now_ns should use the cached entry instead of the current time.
+	ctx := context.Background()
+	rs, err := New(Query(query), NDBuiltinCache(ndBC)).Eval(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Check that we got the correct time value in the result set.
+	assertResultSet(t, rs, "[[1451311705000000000]]")
+}
+
+func TestNDBCacheWithRuleBody(t *testing.T) {
+	ctx := context.Background()
+	ndBC := builtins.NDBCache{}
+	query := "data.foo.p = x"
+	_, err := New(
+		Query(query),
+		NDBuiltinCache(ndBC),
+		Module("test.rego", `package foo
+p {
+	http.send({"url": "http://httpbin.org", "method":"get"})
+}`),
+	).Eval(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, ok := ndBC["http.send"]
+	if !ok {
+		t.Errorf("expected http.send cache entry")
+	}
+}
+
+// This test ensures that the NDBCache correctly serializes/deserializes.
+func TestNDBCacheMarshalUnmarshalJSON(t *testing.T) {
+	original := builtins.NDBCache{}
+
+	// Populate the cache for time.now_ns with an arbitrary timestamp.
+	original.Put("time.now_ns", ast.NewArray(), ast.Number(json.Number(strconv.FormatInt(1451311705000000000, 10))))
+	jOriginal, err := json.Marshal(original)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var other builtins.NDBCache
+	err = json.Unmarshal(jOriginal, &other)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	jOther, err := json.Marshal(other)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Check that the two NDBCache value's JSONified forms match exactly.
+	if !bytes.Equal(jOriginal, jOther) {
+		t.Fatalf("JSONified values of NDBCaches do not match; expected %s, got %s", string(jOriginal), string(jOther))
+	}
+}
+
 func TestStrictBuiltinErrors(t *testing.T) {
 	_, err := New(Query("1/0"), StrictBuiltinErrors(true)).Eval(context.Background())
 	if err == nil {

topdown/builtins.go

@@ -43,6 +43,7 @@ type (
 		Runtime                *ast.Term             // runtime information on the OPA instance
 		Cache                  builtins.Cache        // built-in function state cache
 		InterQueryBuiltinCache cache.InterQueryCache // cross-query built-in function state cache
+		NDBuiltinCache         builtins.NDBCache     // cache for non-deterministic built-in state
 		Location               *ast.Location         // location of built-in call
 		Tracers                []Tracer              // Deprecated: Use QueryTracers instead
 		QueryTracers           []QueryTracer         // tracer objects for trace() built-in function

topdown/builtins/builtins.go

@@ -6,6 +6,7 @@
 package builtins
 
 import (
+	"encoding/json"
 	"fmt"
 	"math/big"
 	"strings"
@@ -28,6 +29,67 @@ func (c Cache) Get(k interface{}) (interface{}, bool) {
 	return v, ok
 }
 
+// We use an ast.Object for the cached keys/values because a naive
+// map[ast.Value]ast.Value will not correctly detect value equality of
+// the member keys.
+type NDBCache map[string]ast.Object
+
+// Put updates the cache for the named built-in.
+// Automatically creates the 2-level hierarchy as needed.
+func (c NDBCache) Put(name string, k, v ast.Value) {
+	if _, ok := c[name]; !ok {
+		c[name] = ast.NewObject()
+	}
+	c[name].Insert(ast.NewTerm(k), ast.NewTerm(v))
+}
+
+// Get returns the cached value for k for the named builtin.
+func (c NDBCache) Get(name string, k ast.Value) (ast.Value, bool) {
+	if m, ok := c[name]; ok {
+		v := m.Get(ast.NewTerm(k))
+		if v != nil {
+			return v.Value, true
+		}
+		return nil, false
+	}
+	return nil, false
+}
+
+// Convenience functions for serializing the data structure.
+func (c NDBCache) MarshalJSON() ([]byte, error) {
+	out := make(map[string]json.RawMessage)
+	for bname, obj := range c {
+		j, err := json.Marshal(ast.NewTerm(obj))
+		if err != nil {
+			return nil, err
+		}
+		out[bname] = j
+	}
+	return json.Marshal(out)
+}
+
+func (c *NDBCache) UnmarshalJSON(data []byte) error {
+	out := map[string]ast.Object{}
+	var incoming map[string]ast.Term
+
+	// We deserialize into a map of Terms, and then extract out the Objects.
+	err := json.Unmarshal(data, &incoming)
+	if err != nil {
+		return err
+	}
+	for k, v := range incoming {
+		if obj, ok := v.Value.(ast.Object); ok {
+			out[k] = obj
+		} else {
+			return fmt.Errorf("expected Object, got other Value type in conversion")
+		}
+	}
+
+	*c = out
+
+	return nil
+}
+
 // ErrOperand represents an invalid operand has been passed to a built-in
 // function. Built-ins should return ErrOperand to indicate a type error has
 // occurred.

topdown/eval.go

@@ -78,6 +78,7 @@ type eval struct {
 	instr                  *Instrumentation
 	builtins               map[string]*Builtin
 	builtinCache           builtins.Cache
+	ndBuiltinCache         builtins.NDBCache
 	functionMocks          *functionMocksStack
 	virtualCache           *virtualCache
 	comprehensionCache     *comprehensionCache
@@ -793,6 +794,7 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
 		Runtime:                e.runtime,
 		Cache:                  e.builtinCache,
 		InterQueryBuiltinCache: e.interQueryBuiltinCache,
+		NDBuiltinCache:         e.ndBuiltinCache,
 		Location:               e.query[e.index].Location,
 		QueryTracers:           e.tracers,
 		TraceEnabled:           e.traceEnabled,
@@ -810,6 +812,7 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
 		f:     f,
 		terms: terms[1:],
 	}
+
 	return eval.eval(iter)
 }
 
@@ -1671,6 +1674,11 @@ type evalBuiltin struct {
 	terms []*ast.Term
 }
 
+// Is this builtin non-deterministic, and did the caller provide an NDBCache?
+func (e *evalBuiltin) canUseNDBCache(bi *ast.Builtin) bool {
+	return bi.Nondeterministic && e.bctx.NDBuiltinCache != nil
+}
+
 func (e evalBuiltin) eval(iter unifyIterator) error {
 
 	operands := make([]*ast.Term, len(e.terms))
@@ -1682,21 +1690,70 @@ func (e evalBuiltin) eval(iter unifyIterator) error {
 	numDeclArgs := len(e.bi.Decl.FuncArgs().Args)
 
 	e.e.instr.startTimer(evalOpBuiltinCall)
+	var err error
 
-	err := e.f(e.bctx, operands, func(output *ast.Term) error {
+	// NOTE(philipc): We sometimes have to drop the very last term off
+	// the args list for cases where a builtin's result is used/assigned,
+	// because the last term will be a generated term, not an actual
+	// argument to the builtin.
+	endIndex := len(operands)
+	if len(operands) > numDeclArgs {
+		endIndex--
+	}
+
+	// We skip evaluation of the builtin entirely if the NDBCache is
+	// present, and we have a non-deterministic builtin already cached.
+	if e.canUseNDBCache(e.bi) {
+		e.e.instr.stopTimer(evalOpBuiltinCall)
+
+		// Unify against the NDBCache result if present.
+		if v, ok := e.bctx.NDBuiltinCache.Get(e.bi.Name, ast.NewArray(e.terms[:endIndex]...)); ok {
+			switch {
+			case e.bi.Decl.Result() == nil:
+				err = iter()
+			case len(operands) == numDeclArgs:
+				if v.Compare(ast.Boolean(false)) != 0 {
+					err = iter()
+				} // else: nothing to do, don't iter()
+			default:
+				err = e.e.unify(e.terms[endIndex], ast.NewTerm(v), iter)
+			}
+
+			if err != nil {
+				return Halt{Err: err}
+			}
+
+			return nil
+		}
+
+		e.e.instr.startTimer(evalOpBuiltinCall)
+
+		// Otherwise, we'll need to go through the normal unify flow.
+	}
+
+	// Normal unification flow for builtins:
+	err = e.f(e.bctx, operands, func(output *ast.Term) error {
 
 		e.e.instr.stopTimer(evalOpBuiltinCall)
 
 		var err error
 
-		if e.bi.Decl.Result() == nil {
+		switch {
+		case e.bi.Decl.Result() == nil:
 			err = iter()
-		} else if len(operands) == numDeclArgs {
+		case len(operands) == numDeclArgs:
 			if output.Value.Compare(ast.Boolean(false)) != 0 {
 				err = iter()
-			}
-		} else {
-			err = e.e.unify(e.terms[len(e.terms)-1], output, iter)
+			} // else: nothing to do, don't iter()
+		default:
+			err = e.e.unify(e.terms[endIndex], output, iter)
+		}
+
+		// If the NDBCache is present, we can assume this builtin
+		// call was not cached earlier.
+		if e.canUseNDBCache(e.bi) {
+			// Populate the NDBCache from the output term.
+			e.bctx.NDBuiltinCache.Put(e.bi.Name, ast.NewArray(e.terms[:endIndex]...), output.Value)
 		}
 
 		if err != nil {

topdown/http.go

@@ -136,7 +136,7 @@ func getHTTPResponse(bctx BuiltinContext, req ast.Object) (*ast.Term, error) {
 		return nil, err
 	}
 
-	// check if cache already has a response for this query
+	// Check if cache already has a response for this query
 	resp, err := reqExecutor.CheckCache()
 	if err != nil {
 		return nil, err
@@ -148,7 +148,7 @@ func getHTTPResponse(bctx BuiltinContext, req ast.Object) (*ast.Term, error) {
 			return nil, err
 		}
 		defer util.Close(httpResp)
-		// add result to cache
+		// Add result to intra/inter-query cache.
 		resp, err = reqExecutor.InsertIntoCache(httpResp)
 		if err != nil {
 			return nil, err

topdown/query.go

@@ -52,6 +52,7 @@ type Query struct {
 	indexing               bool
 	earlyExit              bool
 	interQueryBuiltinCache cache.InterQueryCache
+	ndBuiltinCache         builtins.NDBCache
 	strictBuiltinErrors    bool
 	printHook              print.Hook
 	tracingOpts            tracing.Options
@@ -242,6 +243,12 @@ func (q *Query) WithInterQueryBuiltinCache(c cache.InterQueryCache) *Query {
 	return q
 }
 
+// WithNDBuiltinCache sets the non-deterministic builtin cache.
+func (q *Query) WithNDBuiltinCache(c builtins.NDBCache) *Query {
+	q.ndBuiltinCache = c
+	return q
+}
+
 // WithStrictBuiltinErrors tells the evaluator to treat all built-in function errors as fatal errors.
 func (q *Query) WithStrictBuiltinErrors(yes bool) *Query {
 	q.strictBuiltinErrors = yes
@@ -313,6 +320,7 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support []
 		builtinCache:           builtins.Cache{},
 		functionMocks:          newFunctionMocksStack(),
 		interQueryBuiltinCache: q.interQueryBuiltinCache,
+		ndBuiltinCache:         q.ndBuiltinCache,
 		virtualCache:           newVirtualCache(),
 		comprehensionCache:     newComprehensionCache(),
 		saveSet:                newSaveSet(q.unknowns, b, q.instr),
@@ -462,6 +470,7 @@ func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error {
 		builtinCache:           builtins.Cache{},
 		functionMocks:          newFunctionMocksStack(),
 		interQueryBuiltinCache: q.interQueryBuiltinCache,
+		ndBuiltinCache:         q.ndBuiltinCache,
 		virtualCache:           newVirtualCache(),
 		comprehensionCache:     newComprehensionCache(),
 		genvarprefix:           q.genvarprefix,