Improve storage lookup performance

Some code I had left from the weekend, where I looked into
the cost of querying the store during eval, and if that could
be improved. The code here is mostly benchmarks that I found
useful for that purpose, but also includes a few optimizations
of code that is on the hot path for these lookups and as such
has a high impact — as the benchmarks demonstrate.

The small change in eval.go is the first use of the new feature
of object.Map added some days ago, where a map function returning
a nil key means skipping that key/value pair.

Signed-off-by: Anders Eknert <anders@styra.com>

v1/rego/rego_bench_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/open-policy-agent/opa/internal/runtime"
 	"github.com/open-policy-agent/opa/v1/ast"
 	"github.com/open-policy-agent/opa/v1/loader"
+	"github.com/open-policy-agent/opa/v1/storage"
 	inmem "github.com/open-policy-agent/opa/v1/storage/inmem/test"
 	"github.com/open-policy-agent/opa/v1/util/test"
 )
@@ -318,6 +319,102 @@ func BenchmarkObjectIteration(b *testing.B) {
 	}
 }
 
+// Comparing the cost of referencing not found data in Go vs. AST storage
+//
+// BenchmarkStoreRefNotFound/inmem-go-10         	    5208	    212288 ns/op	  160609 B/op	    2936 allocs/op
+// BenchmarkStoreRefNotFound/inmem-ast-10        	   13929	     90053 ns/op	   39614 B/op	    1012 allocs/op
+func BenchmarkStoreRefNotFound(b *testing.B) {
+	ctx := context.Background()
+
+	things := make(map[string]map[string]string, 100)
+	for i := range 100 {
+		things[strconv.Itoa(i)] = map[string]string{"foo": "bar"}
+	}
+
+	stores := map[string]storage.Store{
+		"inmem-go":  inmem.NewFromObject(map[string]any{"things": things}),
+		"inmem-ast": inmem.NewFromObjectWithASTRead(map[string]any{"things": things}),
+	}
+	policy := `package p
+
+r contains true if {
+	data.things[_].bar
+}
+`
+	for name, store := range stores {
+		b.Run(name, func(b *testing.B) {
+			r := New(
+				Query("data.p.r = x"),
+				Store(store),
+				ParsedModule(ast.MustParseModule(policy)),
+				GenerateJSON(func(*ast.Term, *EvalContext) (any, error) {
+					return nil, nil
+				}),
+			)
+
+			pq, err := r.PrepareForEval(ctx)
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			b.ResetTimer()
+			b.ReportAllocs()
+
+			for range b.N {
+				res, err := pq.Eval(ctx)
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				_ = res
+			}
+		})
+	}
+}
+
+// 242.5 ns/op     168 B/op      7 allocs/op  // original implementation
+// 176.7 ns/op      96 B/op      4 allocs/op  // sync.Pool in ptr.ValuePtr (saving 1 alloc/op per path part)
+func BenchmarkStoreRead(b *testing.B) {
+	ctx := context.Background()
+	store := inmem.NewFromObjectWithASTRead(map[string]any{
+		"foo": map[string]any{
+			"bar": map[string]any{
+				"baz": "qux",
+			},
+		},
+	})
+
+	txn, err := store.NewTransaction(ctx)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	ref := ast.MustParseRef("data.foo.bar.baz")
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for range b.N {
+		// 1 alloc/op
+		path, err := storage.NewPathForRef(ref)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		// 3 allocs/op (down from 6)
+		// turns each string in path into a StringTerm only to use it
+		// for a Get call in storage (ptr.ValuePtr)
+		v, err := store.Read(ctx, txn, path)
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		if v == nil {
+			b.Fatal("expected value")
+		}
+	}
+}
+
 func mustReadFileAsString(b *testing.B, path string) string {
 	b.Helper()
 

v1/storage/inmem/test/testutil.go

@@ -20,3 +20,9 @@ func New() storage.Store {
 func NewFromObject(x map[string]interface{}) storage.Store {
 	return inmem.NewFromObjectWithOpts(x, inmem.OptRoundTripOnWrite(false))
 }
+
+// NewFromObjectWithASTRead returns an inmem store from the passed object, with
+// round-trip on write disabled and AST values returned on read.
+func NewFromObjectWithASTRead(x map[string]interface{}) storage.Store {
+	return inmem.NewFromObjectWithOpts(x, inmem.OptRoundTripOnWrite(false), inmem.OptReturnASTValuesOnRead(true))
+}

v1/storage/internal/errors/errors.go

@@ -20,7 +20,11 @@ func NewNotFoundError(path storage.Path) *storage.Error {
 }
 
 func NewNotFoundErrorWithHint(path storage.Path, hint string) *storage.Error {
-	return NewNotFoundErrorf("%v: %v", path.String(), hint)
+	message := path.String() + ": " + hint
+	return &storage.Error{
+		Code:    storage.NotFoundErr,
+		Message: message,
+	}
 }
 
 func NewNotFoundErrorf(f string, a ...interface{}) *storage.Error {

v1/storage/internal/errors/errors_test.go

@@ -0,0 +1,22 @@
+package errors
+
+import (
+	"testing"
+
+	"github.com/open-policy-agent/opa/v1/storage"
+)
+
+// 160.8 ns/op    96 B/op    5 allocs/op    // using fmt.Sprintf
+// 58.31 ns/op     8 B/op    1 allocs/op    // using string concatenation
+// ...
+func BenchmarkNewNotFoundErrorWithHint(b *testing.B) {
+	path := storage.Path([]string{"a", "b", "c"})
+	hint := "something something"
+
+	for range b.N {
+		err := NewNotFoundErrorWithHint(path, hint)
+		if err == nil {
+			b.Fatal("expected error")
+		}
+	}
+}

v1/storage/path.go

@@ -133,11 +133,22 @@ func (p Path) Ref(head *ast.Term) (ref ast.Ref) {
 }
 
 func (p Path) String() string {
-	buf := make([]string, len(p))
-	for i := range buf {
-		buf[i] = url.PathEscape(p[i])
+	if len(p) == 0 {
+		return "/"
 	}
-	return "/" + strings.Join(buf, "/")
+
+	l := 0
+	for i := range p {
+		l += len(p[i]) + 1
+	}
+
+	sb := strings.Builder{}
+	sb.Grow(l)
+	for i := range p {
+		sb.WriteByte('/')
+		sb.WriteString(url.PathEscape(p[i]))
+	}
+	return sb.String()
 }
 
 // MustParsePath returns a new Path for s. If s cannot be parsed, this function

v1/storage/path_test.go

@@ -192,3 +192,17 @@ func TestPathRef(t *testing.T) {
 		}
 	}
 }
+
+// 108.8 ns/op    80 B/op    3 allocs/op // original implementation concat + Join
+// 68.60 ns/op    24 B/op    2 allocs/op // strings.Builder
+// 50.28 ns/op    16 B/op    1 allocs/op // strings.Builder with pre-allocated buffer
+func BenchmarkPathString(b *testing.B) {
+	path := Path{"foo", "bar", "baz"}
+
+	for range b.N {
+		res := path.String()
+		if res != "/foo/bar/baz" {
+			b.Fatal("unexpected result:", res)
+		}
+	}
+}

v1/topdown/eval.go

@@ -1788,16 +1788,7 @@ func (e *eval) resolveReadFromStorage(ref ast.Ref, a ast.Value) (ast.Value, erro
 				}
 			case ast.Object:
 				if obj.Len() > 0 {
-					cpy := ast.NewObject()
-					if err := obj.Iter(func(k *ast.Term, v *ast.Term) error {
-						if !ast.SystemDocumentKey.Equal(k.Value) {
-							cpy.Insert(k, v)
-						}
-						return nil
-					}); err != nil {
-						return nil, err
-					}
-					blob = cpy
+					blob, _ = obj.Map(systemDocumentKeyRemoveMapper)
 				}
 			}
 		}
@@ -1830,6 +1821,13 @@ func (e *eval) resolveReadFromStorage(ref ast.Ref, a ast.Value) (ast.Value, erro
 	return merged, nil
 }
 
+func systemDocumentKeyRemoveMapper(k, v *ast.Term) (*ast.Term, *ast.Term, error) {
+	if ast.SystemDocumentKey.Equal(k.Value) {
+		return nil, nil, nil
+	}
+	return k, v, nil
+}
+
 func (e *eval) generateVar(suffix string) *ast.Term {
 	buf := make([]byte, 0, len(e.genvarprefix)+len(suffix)+1)