server: Cache PreparedEvalQuery's for data requests

This adds in caching of prepared queries for versioned and unversioned
data queries (POST/GET to `/`, `/data`, and `/v1/data`).

The cache has a max size of 100 and just starts acting as a circular
buffer for queries (FIFO, no smarts for LRU caching or anything). It
seems like it would be unlikely for these API's to hit the cache max
size. Most OPA use-cases have a single query that is re-used over
and over with different inputs.

There is a new metric `counter_server_query_cache_hit` which will
show whether or not a request used the query cache or not. It is there
primarily to help explain away why sometimes a handful of the other
metrics aren't there (the query parse/compile/etc).

In the future this could be added to the query API's too. This change
does not touch anything other than the "data" API's.

Closes: #1567
Signed-off-by: Patrick East <east.patrick@gmail.com>

CHANGELOG.md

@@ -5,6 +5,18 @@ project adheres to [Semantic Versioning](http://semver.org/).
 
 ## Unreleased
 
+### Fixes
+- Rego parse and compile metrics are now correctly tracked and returned
+  when using bundles.
+  
+### Compatibility Notes
+- The `Path` string provided to decision log implementations via the
+  `server.Info#Path` parameter is no longer dot-notation rego reference
+  string. It is now a `/` separated path which has had the `data` prefix
+  removed. This change corrects is issue described in
+  [1532](https://github.com/open-policy-agent/opa/issues/1532) and may
+  affect custom decision log plugins that rely on the `Path`.
+
 ## 0.15.1
 
 In this release we reached a milestone for Wasm: any Rego policy can

metrics/metrics.go

@@ -19,17 +19,18 @@ import (
 
 // Well-known metric names.
 const (
-	ServerHandler     = "server_handler"
-	RegoQueryCompile  = "rego_query_compile"
-	RegoQueryEval     = "rego_query_eval"
-	RegoQueryParse    = "rego_query_parse"
-	RegoModuleParse   = "rego_module_parse"
-	RegoDataParse     = "rego_data_parse"
-	RegoModuleCompile = "rego_module_compile"
-	RegoPartialEval   = "rego_partial_eval"
-	RegoInputParse    = "rego_input_parse"
-	RegoLoadFiles     = "rego_load_files"
-	RegoLoadBundles   = "rego_load_bundles"
+	ServerHandler       = "server_handler"
+	ServerQueryCacheHit = "server_query_cache_hit"
+	RegoQueryCompile    = "rego_query_compile"
+	RegoQueryEval       = "rego_query_eval"
+	RegoQueryParse      = "rego_query_parse"
+	RegoModuleParse     = "rego_module_parse"
+	RegoDataParse       = "rego_data_parse"
+	RegoModuleCompile   = "rego_module_compile"
+	RegoPartialEval     = "rego_partial_eval"
+	RegoInputParse      = "rego_input_parse"
+	RegoLoadFiles       = "rego_load_files"
+	RegoLoadBundles     = "rego_load_bundles"
 )
 
 // Info contains attributes describing the underlying metrics provider.

server/cache.go

@@ -0,0 +1,53 @@
+package server
+
+import "sync"
+
+type cache struct {
+	data    map[string]interface{}
+	keylist []string
+	idx     int
+	maxSize int
+	mtx     sync.RWMutex
+}
+
+func newCache(maxSize int) *cache {
+	return &cache{
+		data:    map[string]interface{}{},
+		keylist: []string{},
+		maxSize: maxSize,
+	}
+}
+
+func (c *cache) Get(k string) (interface{}, bool) {
+	c.mtx.RLock()
+	v, ok := c.data[k]
+	c.mtx.RUnlock()
+	return v, ok
+}
+
+func (c *cache) Insert(k string, v interface{}) {
+
+	// Short path if its already in the cache
+	_, ok := c.Get(k)
+	if ok {
+		return
+	}
+
+	// Slow path, grab the write lock and insert
+	c.mtx.Lock()
+	_, ok = c.data[k]
+	if !ok {
+		c.data[k] = v
+		if len(c.keylist) < c.maxSize {
+			// Haven't reached max size yet, keep adding keys.
+			c.keylist = append(c.keylist, k)
+		} else {
+			// Start recycling spots in the key list and
+			// dropping cache entries for them.
+			delete(c.data, c.keylist[c.idx])
+			c.keylist[c.idx] = k
+			c.idx = (c.idx + 1) % c.maxSize
+		}
+	}
+	c.mtx.Unlock()
+}

server/cache_test.go

@@ -0,0 +1,78 @@
+package server
+
+import (
+	"strconv"
+	"testing"
+)
+
+func TestCacheBase(t *testing.T) {
+	c := newCache(5)
+	foo := struct{}{}
+	c.Insert("foo", foo)
+	ensureCacheKey(t, c, "foo", foo)
+}
+
+func TestCacheLimit(t *testing.T) {
+	max := 10
+	c := newCache(max)
+
+	// Fill the cache with values
+	var i int
+	for i = 0; i < max; i++ {
+		c.Insert(strconv.Itoa(i), i)
+	}
+
+	// Ensure its at the max size
+	ensureCacheSize(t, c, max)
+
+	// Ensure they are all stored..
+	for j := 0; j < max; j++ {
+		ensureCacheKey(t, c, strconv.Itoa(j), j)
+	}
+
+	// Continue filling the cache, expect old keys to be dropped
+	c.Insert(strconv.Itoa(i), i)
+	ensureCacheKey(t, c, strconv.Itoa(i), i)
+
+	// Should still be at max size
+	ensureCacheSize(t, c, max)
+
+	// Expect that "0" got dropped
+	_, ok := c.Get("0")
+	if ok {
+		t.Fatal("Expected key '0' to not be found")
+	}
+
+	// Load the cache with many more than max
+	for ; i < max*20; i++ {
+		c.Insert(strconv.Itoa(i), i)
+	}
+	ensureCacheSize(t, c, max)
+
+	// Ensure the last set of "max" number are available, and everything else is not
+	for j := 0; j < i; j++ {
+		k := strconv.Itoa(j)
+		if j >= (i - max) {
+			ensureCacheKey(t, c, k, j)
+		} else {
+			_, ok := c.Get(k)
+			if ok {
+				t.Fatalf("Expected key %s to not be found", k)
+			}
+		}
+	}
+}
+
+func ensureCacheKey(t *testing.T, c *cache, k string, v interface{}) {
+	t.Helper()
+	actual, ok := c.Get(k)
+	if !ok || v != actual {
+		t.Fatalf("expected to retrieve value %v for key %s, got %v ok==%t", v, k, actual, ok)
+	}
+}
+
+func ensureCacheSize(t *testing.T, c *cache, size int) {
+	if len(c.data) != size && len(c.keylist) != size {
+		t.Fatalf("Unexpected cache size len(data)=%d len(keylist)=%d, expected %d", size, len(c.data), len(c.keylist))
+	}
+}

server/server.go

@@ -79,6 +79,8 @@ const (
 	PromHandlerHealth     = "health"
 )
 
+const pqMaxCacheSize = 100
+
 // map of unsafe builtins
 var unsafeBuiltinsMap = map[string]struct{}{ast.HTTPSend.Name: struct{}{}}
 
@@ -86,30 +88,32 @@ var unsafeBuiltinsMap = map[string]struct{}{ast.HTTPSend.Name: struct{}{}}
 type Server struct {
 	Handler http.Handler
 
-	router            *mux.Router
-	addrs             []string
-	insecureAddr      string
-	authentication    AuthenticationScheme
-	authorization     AuthorizationScheme
-	cert              *tls.Certificate
-	certPool          *x509.CertPool
-	mtx               sync.RWMutex
-	partials          map[string]rego.PartialResult
-	store             storage.Store
-	manager           *plugins.Manager
-	watcher           *watch.Watcher
-	decisionIDFactory func() string
-	revisions         map[string]string
-	legacyRevision    string
-	buffer            Buffer
-	logger            func(context.Context, *Info) error
-	errLimit          int
-	pprofEnabled      bool
-	runtime           *ast.Term
-	httpListeners     []httpListener
-	bundleStatuses    map[string]*bundlePlugin.Status
-	bundleStatusMtx   sync.RWMutex
-	metrics           Metrics
+	router              *mux.Router
+	addrs               []string
+	insecureAddr        string
+	authentication      AuthenticationScheme
+	authorization       AuthorizationScheme
+	cert                *tls.Certificate
+	certPool            *x509.CertPool
+	mtx                 sync.RWMutex
+	partials            map[string]rego.PartialResult
+	preparedEvalQueries *cache
+	store               storage.Store
+	manager             *plugins.Manager
+	watcher             *watch.Watcher
+	decisionIDFactory   func() string
+	revisions           map[string]string
+	legacyRevision      string
+	buffer              Buffer
+	logger              func(context.Context, *Info) error
+	errLimit            int
+	pprofEnabled        bool
+	runtime             *ast.Term
+	httpListeners       []httpListener
+	bundleStatuses      map[string]*bundlePlugin.Status
+	bundleStatusMtx     sync.RWMutex
+	metrics             Metrics
+	defaultDecisionPath string
 }
 
 // Metrics defines the interface that the server requires for recording HTTP
@@ -174,6 +178,7 @@ func (s *Server) Init(ctx context.Context) (*Server, error) {
 	}
 
 	s.partials = map[string]rego.PartialResult{}
+	s.preparedEvalQueries = newCache(pqMaxCacheSize)
 
 	bp := bundlePlugin.Lookup(s.manager)
 	if bp != nil {
@@ -193,6 +198,8 @@ func (s *Server) Init(ctx context.Context) (*Server, error) {
 		s.legacyRevision = rev
 	}
 
+	s.defaultDecisionPath = s.generateDefaultDecisionPath()
+
 	return s, s.store.Commit(ctx, txn)
 }
 
@@ -716,6 +723,8 @@ func (s *Server) reload(ctx context.Context, txn storage.Transaction, event stor
 	// reset some cached info
 	s.partials = map[string]rego.PartialResult{}
 	s.revisions = map[string]string{}
+	s.preparedEvalQueries = newCache(pqMaxCacheSize)
+	s.defaultDecisionPath = s.generateDefaultDecisionPath()
 
 	// read all bundle revisions from storage (if any exist)
 	names, err := bundle.ReadBundleNamesFromStore(ctx, s.store, txn)
@@ -751,15 +760,14 @@ func (s *Server) migrateWatcher(txn storage.Transaction) {
 }
 
 func (s *Server) unversionedPost(w http.ResponseWriter, r *http.Request) {
-	s.v0QueryPath(w, r, s.manager.Config.DefaultDecisionRef())
+	s.v0QueryPath(w, r, s.defaultDecisionPath)
 }
 
 func (s *Server) v0DataPost(w http.ResponseWriter, r *http.Request) {
-	path := stringPathToDataRef(mux.Vars(r)["path"])
-	s.v0QueryPath(w, r, path)
+	s.v0QueryPath(w, r, mux.Vars(r)["path"])
 }
 
-func (s *Server) v0QueryPath(w http.ResponseWriter, r *http.Request, path ast.Ref) {
+func (s *Server) v0QueryPath(w http.ResponseWriter, r *http.Request, urlPath string) {
 	m := metrics.New()
 	m.Timer(metrics.ServerHandler).Start()
 
@@ -792,32 +800,50 @@ func (s *Server) v0QueryPath(w http.ResponseWriter, r *http.Request, path ast.Re
 
 	defer s.store.Abort(ctx, txn)
 
-	rego := rego.New(
-		rego.Compiler(s.getCompiler()),
-		rego.Store(s.store),
-		rego.Transaction(txn),
-		rego.ParsedInput(input),
-		rego.Query(path.String()),
-		rego.Metrics(m),
-		rego.Runtime(s.runtime),
-		rego.UnsafeBuiltins(unsafeBuiltinsMap),
-	)
+	pqID := "v0QueryPath::" + urlPath
+	preparedQuery, ok := s.getCachedPreparedEvalQuery(pqID, m)
+	if !ok {
+		path := stringPathToDataRef(urlPath)
 
-	rs, err := rego.Eval(ctx)
+		rego := rego.New(
+			rego.Compiler(s.getCompiler()),
+			rego.Store(s.store),
+			rego.Transaction(txn),
+			rego.Query(path.String()),
+			rego.Metrics(m),
+			rego.Runtime(s.runtime),
+			rego.UnsafeBuiltins(unsafeBuiltinsMap),
+		)
+		pq, err := rego.PrepareForEval(ctx)
+		if err != nil {
+			_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, err, m)
+			writer.ErrorAuto(w, err)
+			return
+		}
+		preparedQuery = &pq
+		s.preparedEvalQueries.Insert(pqID, preparedQuery)
+	}
+
+	rs, err := preparedQuery.Eval(
+		ctx,
+		rego.EvalTransaction(txn),
+		rego.EvalParsedInput(input),
+		rego.EvalMetrics(m),
+	)
 
 	m.Timer(metrics.ServerHandler).Stop()
 
 	// Handle results.
 	if err != nil {
-		_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, err, m)
+		_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, err, m)
 		writer.ErrorAuto(w, err)
 		return
 	}
 
 	if len(rs) == 0 {
-		err := types.NewErrorV1(types.CodeUndefinedDocument, fmt.Sprintf("%v: %v", types.MsgUndefinedError, path))
+		err := types.NewErrorV1(types.CodeUndefinedDocument, fmt.Sprintf("%v: %v", types.MsgUndefinedError, stringPathToDataRef(urlPath)))
 
-		if logErr := logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, err, m); logErr != nil {
+		if logErr := logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, err, m); logErr != nil {
 			writer.ErrorAuto(w, logErr)
 			return
 		}
@@ -826,7 +852,7 @@ func (s *Server) v0QueryPath(w http.ResponseWriter, r *http.Request, path ast.Re
 		return
 	}
 
-	err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, &rs[0].Expressions[0].Value, nil, m)
+	err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, &rs[0].Expressions[0].Value, nil, m)
 	if err != nil {
 		writer.ErrorAuto(w, err)
 		return
@@ -835,6 +861,16 @@ func (s *Server) v0QueryPath(w http.ResponseWriter, r *http.Request, path ast.Re
 	writer.JSON(w, 200, rs[0].Expressions[0].Value, false)
 }
 
+func (s *Server) getCachedPreparedEvalQuery(key string, m metrics.Metrics) (*rego.PreparedEvalQuery, bool) {
+	pq, ok := s.preparedEvalQueries.Get(key)
+	m.Counter(metrics.ServerQueryCacheHit) // Creates the counter on the metrics if it doesn't exist, starts at 0
+	if ok {
+		m.Counter(metrics.ServerQueryCacheHit).Incr() // Increment counter on hit
+		return pq.(*rego.PreparedEvalQuery), true
+	}
+	return nil, false
+}
+
 func (s *Server) updateBundleStatus(status map[string]*bundlePlugin.Status) {
 	s.bundleStatusMtx.Lock()
 	defer s.bundleStatusMtx.Unlock()
@@ -982,12 +1018,12 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 
 	ctx := r.Context()
 	vars := mux.Vars(r)
-	path := stringPathToDataRef(vars["path"])
+	urlPath := vars["path"]
 	logger := s.getDecisionLogger()
 
 	watch := getWatch(r.URL.Query()[types.ParamWatchV1])
 	if watch {
-		s.watchQuery(path.String(), w, r, true)
+		s.watchQuery(stringPathToDataRef(urlPath).String(), w, r, true)
 		return
 	}
 
@@ -1030,7 +1066,6 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 		writer.ErrorAuto(w, err)
 		return
 	}
-
 	defer s.store.Abort(ctx, txn)
 
 	var buf *topdown.BufferTracer
@@ -1039,24 +1074,45 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 		buf = topdown.NewBufferTracer()
 	}
 
-	rego := rego.New(
-		rego.Compiler(s.getCompiler()),
-		rego.Store(s.store),
-		rego.Transaction(txn),
-		rego.ParsedInput(input),
-		rego.Query(path.String()),
-		rego.Metrics(m),
-		rego.Tracer(buf),
-		rego.Instrument(includeInstrumentation),
-		rego.Runtime(s.runtime),
-		rego.UnsafeBuiltins(unsafeBuiltinsMap),
+	pqID := "v1DataGet::" + urlPath
+	preparedQuery, ok := s.getCachedPreparedEvalQuery(pqID, m)
+	if !ok {
+		rego := rego.New(
+			rego.Compiler(s.getCompiler()),
+			rego.Store(s.store),
+			rego.Transaction(txn),
+			rego.ParsedInput(input),
+			rego.Query(stringPathToDataRef(urlPath).String()),
+			rego.Metrics(m),
+			rego.Tracer(buf),
+			rego.Instrument(includeInstrumentation),
+			rego.Runtime(s.runtime),
+			rego.UnsafeBuiltins(unsafeBuiltinsMap),
+		)
+
+		pq, err := rego.PrepareForEval(ctx)
+		if err != nil {
+			_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, err, m)
+			writer.ErrorAuto(w, err)
+			return
+		}
+		preparedQuery = &pq
+		s.preparedEvalQueries.Insert(pqID, preparedQuery)
+	}
+
+	rs, err := preparedQuery.Eval(
+		ctx,
+		rego.EvalTransaction(txn),
+		rego.EvalParsedInput(input),
+		rego.EvalMetrics(m),
+		rego.EvalTracer(buf),
 	)
 
-	rs, err := rego.Eval(ctx)
+	m.Timer(metrics.ServerHandler).Stop()
 
 	// Handle results.
 	if err != nil {
-		_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, err, m)
+		_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, err, m)
 		writer.ErrorAuto(w, err)
 		return
 	}
@@ -1065,8 +1121,6 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 		DecisionID: decisionID,
 	}
 
-	m.Timer(metrics.ServerHandler).Stop()
-
 	if includeMetrics || includeInstrumentation {
 		result.Metrics = m.All()
 	}
@@ -1082,7 +1136,7 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 				writer.ErrorAuto(w, err)
 			}
 		}
-		err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, nil, m)
+		err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, nil, m)
 		if err != nil {
 			writer.ErrorAuto(w, err)
 			return
@@ -1097,7 +1151,7 @@ func (s *Server) v1DataGet(w http.ResponseWriter, r *http.Request) {
 		result.Explanation = s.getExplainResponse(explainMode, *buf, pretty)
 	}
 
-	err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, result.Result, nil, m)
+	err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, result.Result, nil, m)
 	if err != nil {
 		writer.ErrorAuto(w, err)
 		return
@@ -1161,12 +1215,12 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 
 	ctx := r.Context()
 	vars := mux.Vars(r)
-	path := stringPathToDataRef(vars["path"])
+	urlPath := vars["path"]
 	logger := s.getDecisionLogger()
 
 	watch := getWatch(r.URL.Query()[types.ParamWatchV1])
 	if watch {
-		s.watchQuery(path.String(), w, r, true)
+		s.watchQuery(stringPathToDataRef(urlPath).String(), w, r, true)
 		return
 	}
 
@@ -1202,7 +1256,6 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 		writer.ErrorAuto(w, err)
 		return
 	}
-
 	defer s.store.Abort(ctx, txn)
 
 	opts := []func(*rego.Rego){
@@ -1216,21 +1269,44 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 		buf = topdown.NewBufferTracer()
 	}
 
-	rego, err := s.makeRego(ctx, partial, txn, input, path.String(), m, includeInstrumentation, buf, opts)
+	pqID := "v1DataPost::"
+	if partial {
+		pqID += "partial::"
+	}
+	pqID += urlPath
+	preparedQuery, ok := s.getCachedPreparedEvalQuery(pqID, m)
+	if !ok {
+		rego, err := s.makeRego(ctx, partial, txn, input, stringPathToDataRef(urlPath).String(), m, includeInstrumentation, buf, opts)
 
-	if err != nil {
-		_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, err, m)
-		writer.ErrorAuto(w, err)
-		return
+		if err != nil {
+			_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, err, m)
+			writer.ErrorAuto(w, err)
+			return
+		}
+
+		pq, err := rego.PrepareForEval(ctx)
+		if err != nil {
+			_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, err, m)
+			writer.ErrorAuto(w, err)
+			return
+		}
+		preparedQuery = &pq
+		s.preparedEvalQueries.Insert(pqID, preparedQuery)
 	}
 
-	rs, err := rego.Eval(ctx)
+	rs, err := preparedQuery.Eval(
+		ctx,
+		rego.EvalTransaction(txn),
+		rego.EvalParsedInput(input),
+		rego.EvalMetrics(m),
+		rego.EvalTracer(buf),
+	)
 
 	m.Timer(metrics.ServerHandler).Stop()
 
 	// Handle results.
 	if err != nil {
-		_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, err, m)
+		_ = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, err, m)
 		writer.ErrorAuto(w, err)
 		return
 	}
@@ -1254,7 +1330,7 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 				writer.ErrorAuto(w, err)
 			}
 		}
-		err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, nil, nil, m)
+		err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, nil, nil, m)
 		if err != nil {
 			writer.ErrorAuto(w, err)
 			return
@@ -1269,7 +1345,7 @@ func (s *Server) v1DataPost(w http.ResponseWriter, r *http.Request) {
 		result.Explanation = s.getExplainResponse(explainMode, *buf, pretty)
 	}
 
-	err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, path.String(), "", goInput, result.Result, nil, m)
+	err = logger.Log(ctx, txn, decisionID, r.RemoteAddr, urlPath, "", goInput, result.Result, nil, m)
 	if err != nil {
 		writer.ErrorAuto(w, err)
 		return
@@ -2061,6 +2137,12 @@ func (s *Server) hasLegacyBundle() bool {
 	return s.legacyRevision != "" || (bp != nil && !bp.Config().IsMultiBundle())
 }
 
+func (s *Server) generateDefaultDecisionPath() string {
+	// Assume the path is safe to transition back to a url
+	p, _ := s.manager.Config.DefaultDecisionRef().Ptr()
+	return p
+}
+
 // parsePatchPathEscaped returns a new path for the given escaped str.
 // This is based on storage.ParsePathEscaped so will do URL unescaping of
 // the provided str for backwards compatibility, but also handles the

server/server_test.go

@@ -1875,8 +1875,45 @@ func TestDataMetrics(t *testing.T) {
 
 	f := newFixture(t)
 
-	req := newReqV1(http.MethodPost, "/data?metrics", "")
+	testDataMetrics(t, f, "/data?metrics", []string{
+		"counter_server_query_cache_hit",
+		"timer_rego_input_parse_ns",
+		"timer_rego_load_bundles_ns",
+		"timer_rego_load_files_ns",
+		"timer_rego_module_parse_ns",
+		"timer_rego_query_parse_ns",
+		"timer_rego_query_compile_ns",
+		"timer_rego_query_eval_ns",
+		"timer_server_handler_ns",
+	})
+
+	testDataMetrics(t, f, "/data?metrics", []string{
+		"counter_server_query_cache_hit",
+		"timer_rego_input_parse_ns",
+		"timer_rego_query_parse_ns",
+		"timer_rego_query_eval_ns",
+		"timer_server_handler_ns",
+	})
+
+	testDataMetrics(t, f, "/data?metrics&partial", []string{
+		"counter_server_query_cache_hit",
+		"timer_rego_input_parse_ns",
+		"timer_rego_load_bundles_ns",
+		"timer_rego_load_files_ns",
+		"timer_rego_module_compile_ns",
+		"timer_rego_module_parse_ns",
+		"timer_rego_query_parse_ns",
+		"timer_rego_query_compile_ns",
+		"timer_rego_query_eval_ns",
+		"timer_rego_partial_eval_ns",
+		"timer_server_handler_ns",
+	})
+}
+
+func testDataMetrics(t *testing.T, f *fixture, url string, expected []string) {
+	t.Helper()
 	f.reset()
+	req := newReqV1(http.MethodPost, url, "")
 	f.server.Handler.ServeHTTP(f.recorder, req)
 
 	var result types.DataResponseV1
@@ -1885,45 +1922,20 @@ func TestDataMetrics(t *testing.T) {
 		t.Fatalf("Unexpected JSON decode error: %v", err)
 	}
 
-	// Test some basic well-known metrics.
-	expected := []string{
-		"timer_rego_query_parse_ns",
-		"timer_rego_query_compile_ns",
-		"timer_rego_query_eval_ns",
-		"timer_server_handler_ns",
-	}
-
 	for _, key := range expected {
-		if result.Metrics[key] == nil {
-			t.Fatalf("Expected non-zero metric for %v but got: %v", key, result)
+		v, ok := result.Metrics[key]
+		if !ok {
+			t.Fatalf("Missing expected metric: %s", key)
 		}
-	}
-
-	req = newReqV1(http.MethodPost, "/data?metrics&partial", "")
-
-	f.reset()
-	f.server.Handler.ServeHTTP(f.recorder, req)
-
-	result = types.DataResponseV1{}
-
-	if err := util.NewJSONDecoder(f.recorder.Body).Decode(&result); err != nil {
-		t.Fatalf("Unexpected JSON decode error: %v", err)
-	}
-
-	expected = []string{
-		"timer_rego_query_parse_ns",
-		"timer_rego_query_compile_ns",
-		"timer_rego_query_eval_ns",
-		"timer_rego_partial_eval_ns",
-		"timer_server_handler_ns",
-	}
-
-	for _, key := range expected {
-		if result.Metrics[key] == nil {
-			t.Fatalf("Expected non-zero metric for %v but got: %v", key, result)
+		if v == nil {
+			t.Fatalf("Expected non-nil value for metric: %s", key)
 		}
+
 	}
 
+	if len(expected) != len(result.Metrics) {
+		t.Fatalf("Expected %d metrics, got %d\n\n\tValues: %+v", len(expected), len(result.Metrics), result.Metrics)
+	}
 }
 
 func TestV1Pretty(t *testing.T) {
@@ -2625,7 +2637,7 @@ func TestDecisionLogging(t *testing.T) {
 		nextID++
 		return fmt.Sprint(nextID)
 	}).WithDecisionLoggerWithErr(func(_ context.Context, info *Info) error {
-		if info.Path == "data.fail_closed.decision_logger_err" {
+		if info.Path == "fail_closed/decision_logger_err" {
 			return fmt.Errorf("some error")
 		}
 		decisions = append(decisions, info)
@@ -2753,17 +2765,17 @@ func TestDecisionLogging(t *testing.T) {
 		query   string
 		wantErr bool
 	}{
-		{path: "data"},
-		{path: "data"},
-		{path: "data.nonexistent", input: `{"foo": 1}`},
-		{path: "data"},
-		{path: "data.system.main"},
+		{path: ""},
+		{path: ""},
+		{path: "nonexistent", input: `{"foo": 1}`},
+		{path: ""},
+		{path: "system/main"},
 		{query: "data = x"},
 		{query: "data = x"},
-		{path: "data", wantErr: true},
-		{path: "data", wantErr: true},
-		{path: "data.system.main", wantErr: true},
-		{path: `data.test`, wantErr: true},
+		{path: "", wantErr: true},
+		{path: "", wantErr: true},
+		{path: "system/main", wantErr: true},
+		{path: `test`, wantErr: true},
 	}
 
 	if len(decisions) != len(exp) {
@@ -3405,7 +3417,7 @@ func (f *fixture) executeRequest(req *http.Request, code int, resp string) error
 	f.reset()
 	f.server.Handler.ServeHTTP(f.recorder, req)
 	if f.recorder.Code != code {
-		return fmt.Errorf("Expected code %v from %v %v but got: %v", code, req.Method, req.URL, f.recorder)
+		return fmt.Errorf("Expected code %v from %v %v but got: %+v", code, req.Method, req.URL, f.recorder)
 	}
 	if resp != "" {
 		var result interface{}