chore: Optimize rego policy input allocations (#6135)

* chore: Optimize rego policy evaluation allocations

Manually convert to ast.Value instead of using generic
json.Marshal conversion.

* Add a unit test that prevents regressions of rego input

The optimized input is always compared to the normal json
marshal parser.

coderd/rbac/builtin_internal_test.go

@@ -3,11 +3,191 @@ package rbac
 import (
 	"testing"
 
-	"github.com/stretchr/testify/require"
-
 	"github.com/google/uuid"
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/stretchr/testify/require"
 )
 
+// BenchmarkRBACValueAllocation benchmarks the cost of allocating a rego input
+// value. By default, `ast.InterfaceToValue` is used to convert the input,
+// which uses json marshalling under the hood.
+//
+// Currently ast.Object.insert() is the slowest part of the process and allocates
+// the most amount of bytes. This general approach copies all of our struct
+// data and uses a lot of extra memory for handling things like sort order.
+// A possible large improvement would be to implement the ast.Value interface directly.
+func BenchmarkRBACValueAllocation(b *testing.B) {
+	actor := Subject{
+		Roles:  RoleNames{RoleOrgMember(uuid.New()), RoleOrgAdmin(uuid.New()), RoleMember()},
+		ID:     uuid.NewString(),
+		Scope:  ScopeAll,
+		Groups: []string{uuid.NewString(), uuid.NewString(), uuid.NewString()},
+	}
+	obj := ResourceTemplate.
+		WithID(uuid.New()).
+		InOrg(uuid.New()).
+		WithOwner(uuid.NewString()).
+		WithGroupACL(map[string][]Action{
+			uuid.NewString(): {ActionRead, ActionCreate},
+			uuid.NewString(): {ActionRead, ActionCreate},
+			uuid.NewString(): {ActionRead, ActionCreate},
+		}).WithACLUserList(map[string][]Action{
+		uuid.NewString(): {ActionRead, ActionCreate},
+		uuid.NewString(): {ActionRead, ActionCreate},
+	})
+
+	jsonSubject := authSubject{
+		ID:     actor.ID,
+		Roles:  must(actor.Roles.Expand()),
+		Groups: actor.Groups,
+		Scope:  must(actor.Scope.Expand()),
+	}
+
+	b.Run("ManualRegoValue", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			_, err := regoInputValue(actor, ActionRead, obj)
+			require.NoError(b, err)
+		}
+	})
+	b.Run("JSONRegoValue", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			_, err := ast.InterfaceToValue(jsonSubject)
+			require.NoError(b, err)
+		}
+	})
+}
+
+// TestRegoInputValue ensures the custom rego input parser returns the
+// same value as the default json parser. The json parser is always correct,
+// and the custom parser is used to reduce allocations. This optimization
+// should yield the same results. Anything different is a bug.
+func TestRegoInputValue(t *testing.T) {
+	t.Parallel()
+
+	actor := Subject{
+		Roles:  RoleNames{RoleOrgMember(uuid.New()), RoleOrgAdmin(uuid.New()), RoleMember()},
+		ID:     uuid.NewString(),
+		Scope:  ScopeAll,
+		Groups: []string{uuid.NewString(), uuid.NewString(), uuid.NewString()},
+	}
+
+	obj := ResourceTemplate.
+		WithID(uuid.New()).
+		InOrg(uuid.New()).
+		WithOwner(uuid.NewString()).
+		WithGroupACL(map[string][]Action{
+			uuid.NewString(): {ActionRead, ActionCreate},
+			uuid.NewString(): {ActionRead, ActionCreate},
+			uuid.NewString(): {ActionRead, ActionCreate},
+		}).WithACLUserList(map[string][]Action{
+		uuid.NewString(): {ActionRead, ActionCreate},
+		uuid.NewString(): {ActionRead, ActionCreate},
+	})
+
+	action := ActionRead
+
+	t.Run("InputValue", func(t *testing.T) {
+		t.Parallel()
+
+		// This is the input that would be passed to the rego policy.
+		jsonInput := map[string]interface{}{
+			"subject": authSubject{
+				ID:     actor.ID,
+				Roles:  must(actor.Roles.Expand()),
+				Groups: actor.Groups,
+				Scope:  must(actor.Scope.Expand()),
+			},
+			"action": action,
+			"object": obj,
+		}
+
+		manual, err := regoInputValue(actor, action, obj)
+		require.NoError(t, err)
+
+		general, err := ast.InterfaceToValue(jsonInput)
+		require.NoError(t, err)
+
+		// The custom parser does not set these fields because they are not needed.
+		// To ensure the outputs are identical, intentionally overwrite all names
+		// to the same values.
+		ignoreNames(t, manual)
+		ignoreNames(t, general)
+
+		cmp := manual.Compare(general)
+		require.Equal(t, 0, cmp, "manual and general input values should be equal")
+	})
+
+	t.Run("PartialInputValue", func(t *testing.T) {
+		t.Parallel()
+
+		// This is the input that would be passed to the rego policy.
+		jsonInput := map[string]interface{}{
+			"subject": authSubject{
+				ID:     actor.ID,
+				Roles:  must(actor.Roles.Expand()),
+				Groups: actor.Groups,
+				Scope:  must(actor.Scope.Expand()),
+			},
+			"action": action,
+			"object": map[string]interface{}{
+				"type": obj.Type,
+			},
+		}
+
+		manual, err := regoPartialInputValue(actor, action, obj.Type)
+		require.NoError(t, err)
+
+		general, err := ast.InterfaceToValue(jsonInput)
+		require.NoError(t, err)
+
+		// The custom parser does not set these fields because they are not needed.
+		// To ensure the outputs are identical, intentionally overwrite all names
+		// to the same values.
+		ignoreNames(t, manual)
+		ignoreNames(t, general)
+
+		cmp := manual.Compare(general)
+		require.Equal(t, 0, cmp, "manual and general input values should be equal")
+	})
+}
+
+// ignoreNames sets all names to "ignore" to ensure the values are identical.
+func ignoreNames(t *testing.T, value ast.Value) {
+	t.Helper()
+
+	// Override the names of the roles
+	ref := ast.Ref{
+		ast.StringTerm("subject"),
+		ast.StringTerm("roles"),
+	}
+	roles, err := value.Find(ref)
+	require.NoError(t, err)
+
+	rolesArray, ok := roles.(*ast.Array)
+	require.True(t, ok, "roles is expected to be an array")
+
+	rolesArray.Foreach(func(term *ast.Term) {
+		obj, _ := term.Value.(ast.Object)
+		// Ignore all names
+		obj.Insert(ast.StringTerm("name"), ast.StringTerm("ignore"))
+		obj.Insert(ast.StringTerm("display_name"), ast.StringTerm("ignore"))
+	})
+
+	// Override the names of the scope role
+	ref = ast.Ref{
+		ast.StringTerm("subject"),
+		ast.StringTerm("scope"),
+	}
+	scope, err := value.Find(ref)
+	require.NoError(t, err)
+
+	scopeObj, ok := scope.(ast.Object)
+	require.True(t, ok, "scope is expected to be an object")
+
+	scopeObj.Insert(ast.StringTerm("name"), ast.StringTerm("ignore"))
+	scopeObj.Insert(ast.StringTerm("display_name"), ast.StringTerm("ignore"))
+}
+
 func TestRoleByName(t *testing.T) {
 	t.Parallel()