WIP: Pass mutex to built-in functions via BuiltinContext

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>

topdown/builtins.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
+	"sync"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/metrics"
@@ -54,6 +55,7 @@ type (
 		DistributedTracingOpts tracing.Options       // options to be used by distributed tracing.
 		rand                   *rand.Rand            // randomization source for non-security-sensitive operations
 		Capabilities           *ast.Capabilities
+		Mtx                    *sync.Mutex
 	}
 
 	// BuiltinFunc defines an interface for implementing built-in functions.

topdown/eval.go

@@ -7,6 +7,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/metrics"
@@ -98,6 +99,7 @@ type eval struct {
 	tracingOpts            tracing.Options
 	findOne                bool
 	strictObjects          bool
+	mtx                    *sync.Mutex
 }
 
 func (e *eval) Run(iter evalIterator) error {
@@ -811,6 +813,7 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error {
 		PrintHook:              e.printHook,
 		DistributedTracingOpts: e.tracingOpts,
 		Capabilities:           capabilities,
+		Mtx:                    e.mtx,
 	}
 
 	eval := evalBuiltin{

topdown/http.go

@@ -124,7 +124,15 @@ func builtinHTTPSend(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.T
 		return handleBuiltinErr(ast.HTTPSend.Name, bctx.Location, err)
 	}
 
+	if bctx.Mtx != nil {
+		bctx.Mtx.Lock()
+	}
 	result, err := getHTTPResponse(bctx, req)
+
+	if bctx.Mtx != nil {
+		bctx.Mtx.Unlock()
+	}
+
 	if err != nil {
 		if raiseError {
 			return handleHTTPSendErr(bctx, err)
@@ -807,6 +815,7 @@ func (c *interQueryCache) checkHTTPSendInterQueryCache() (ast.Value, error) {
 	if !found {
 		return nil, nil
 	}
+
 	c.bctx.Metrics.Counter(httpSendInterQueryCacheHits).Incr()
 	var cachedRespData *interQueryCacheData
 

topdown/http_test.go

@@ -23,6 +23,7 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -2830,6 +2831,57 @@ func TestHTTPSendCacheDefaultStatusCodesInterQueryCache(t *testing.T) {
 	})
 }
 
+func TestHTTPSendInterQueryCacheConcurrentRequests(t *testing.T) {
+
+	// run test server
+	var requests []*http.Request
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requests = append(requests, r)
+		headers := w.Header()
+		headers["Cache-Control"] = []string{"max-age=290304000, public"}
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	defer ts.Close()
+
+	// add an inter-query cache
+	config, _ := iCache.ParseCachingConfig(nil)
+	interQueryCache := iCache.NewInterQueryCache(config)
+
+	m := metrics.New()
+	var mtx sync.Mutex
+
+	done := make(chan struct{})
+	go func() {
+		q := NewQuery(ast.MustParseBody(fmt.Sprintf(`http.send({"method": "get", "url": %q, "cache": true})`, ts.URL))).
+			WithMetrics(m).WithInterQueryBuiltinCache(interQueryCache).WithMutex(&mtx)
+
+		_, err := q.Run(context.Background())
+		if err != nil {
+			t.Error(err)
+		}
+		done <- struct{}{}
+	}()
+
+	q := NewQuery(ast.MustParseBody(fmt.Sprintf(`http.send({"method": "get", "url": %q, "cache": true})`, ts.URL))).
+		WithMetrics(m).WithInterQueryBuiltinCache(interQueryCache).WithMutex(&mtx)
+
+	_, err := q.Run(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	<-done
+
+	expectedReqCount := 1
+	if len(requests) != expectedReqCount {
+		t.Fatalf("expected to get %d requests, got %d", expectedReqCount, len(requests))
+	}
+
+	if exp, act := uint64(1), m.Counter(httpSendInterQueryCacheHits).Value(); exp != act {
+		t.Fatalf("expected %d cache hits, got %d", exp, act)
+	}
+}
+
 type onlyOnceInterQueryCache struct {
 	value   *interQueryCacheData
 	counter int

topdown/query.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"io"
 	"sort"
+	"sync"
 	"time"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -58,6 +59,7 @@ type Query struct {
 	strictObjects          bool
 	printHook              print.Hook
 	tracingOpts            tracing.Options
+	mtx                    *sync.Mutex
 }
 
 // Builtin represents a built-in function that queries can call.
@@ -291,6 +293,12 @@ func (q *Query) WithStrictObjects(yes bool) *Query {
 	return q
 }
 
+// WithMutex sets the mutual exclusion lock that built-in functions can utilize.
+func (q *Query) WithMutex(mtx *sync.Mutex) *Query {
+	q.mtx = mtx
+	return q
+}
+
 // PartialRun executes partial evaluation on the query with respect to unknown
 // values. Partial evaluation attempts to evaluate as much of the query as
 // possible without requiring values for the unknowns set on the query. The
@@ -357,6 +365,7 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support []
 		builtinErrors: &builtinErrors{},
 		printHook:     q.printHook,
 		strictObjects: q.strictObjects,
+		mtx:           q.mtx,
 	}
 
 	if len(q.disableInlining) > 0 {
@@ -526,6 +535,7 @@ func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error {
 		printHook:              q.printHook,
 		tracingOpts:            q.tracingOpts,
 		strictObjects:          q.strictObjects,
+		mtx:                    q.mtx,
 	}
 	e.caller = e
 	q.metrics.Timer(metrics.RegoQueryEval).Start()