chore: Optimize parial rego execution byte allocations (coder#6144)

Emyrk · web-flow · commit 32fbd10a1f29 · 2023-02-10T08:39:45.000-06:00
* chore: Implement benchmark for authorizer.Prepare

Identify time + alloc cost before optimizing
diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go
@@ -139,7 +139,8 @@ func Filter[O Objecter](ctx context.Context, auth Authorizer, subject Subject, a
 
 // RegoAuthorizer will use a prepared rego query for performing authorize()
 type RegoAuthorizer struct {
-	query rego.PreparedEvalQuery
+	query        rego.PreparedEvalQuery
+	partialQuery rego.PreparedPartialQuery
 
 	authorizeHist *prometheus.HistogramVec
 	prepareHist   prometheus.Histogram
@@ -151,9 +152,10 @@ var (
 	// Load the policy from policy.rego in this directory.
 	//
 	//go:embed policy.rego
-	policy    string
-	queryOnce sync.Once
-	query     rego.PreparedEvalQuery
+	policy       string
+	queryOnce    sync.Once
+	query        rego.PreparedEvalQuery
+	partialQuery rego.PreparedPartialQuery
 )
 
 func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
@@ -166,6 +168,21 @@ func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
 		if err != nil {
 			panic(xerrors.Errorf("compile rego: %w", err))
 		}
+
+		partialQuery, err = rego.New(
+			rego.Unknowns([]string{
+				"input.object.id",
+				"input.object.owner",
+				"input.object.org_owner",
+				"input.object.acl_user_list",
+				"input.object.acl_group_list",
+			}),
+			rego.Query("data.authz.allow = true"),
+			rego.Module("policy.rego", policy),
+		).PrepareForPartial(context.Background())
+		if err != nil {
+			panic(xerrors.Errorf("compile partial rego: %w", err))
+		}
 	})
 
 	// Register metrics to prometheus.
@@ -207,7 +224,8 @@ func NewAuthorizer(registry prometheus.Registerer) *RegoAuthorizer {
 	})
 
 	return &RegoAuthorizer{
-		query: query,
+		query:        query,
+		partialQuery: partialQuery,
 
 		authorizeHist: authorizeHistogram,
 		prepareHist:   prepareHistogram,
@@ -289,7 +307,7 @@ func (a RegoAuthorizer) Prepare(ctx context.Context, subject Subject, action Act
 	)
 	defer span.End()
 
-	prepared, err := newPartialAuthorizer(ctx, subject, action, objectType)
+	prepared, err := a.newPartialAuthorizer(ctx, subject, action, objectType)
 	if err != nil {
 		return nil, xerrors.Errorf("new partial authorizer: %w", err)
 	}
diff --git a/coderd/rbac/authz_internal_test.go b/coderd/rbac/authz_internal_test.go
@@ -977,9 +977,14 @@ func testAuthorize(t *testing.T, name string, subject Subject, sets ...[]authTes
 
 					d, _ := json.Marshal(map[string]interface{}{
 						// This is not perfect marshal, but it is good enough for debugging this test.
-						"subject": subject,
-						"object":  c.resource,
-						"action":  a,
+						"subject": authSubject{
+							ID:     subject.ID,
+							Roles:  must(subject.Roles.Expand()),
+							Groups: subject.Groups,
+							Scope:  must(subject.Scope.Expand()),
+						},
+						"object": c.resource,
+						"action": a,
 					})
 
 					// Logging only
diff --git a/coderd/rbac/authz_test.go b/coderd/rbac/authz_test.go
@@ -189,6 +189,17 @@ func BenchmarkRBACFilter(b *testing.B) {
 	)
 
 	authorizer := rbac.NewAuthorizer(prometheus.NewRegistry())
+
+	for _, c := range benchCases {
+		b.Run("PrepareOnly-"+c.Name, func(b *testing.B) {
+			obType := rbac.ResourceWorkspace.Type
+			for i := 0; i < b.N; i++ {
+				_, err := authorizer.Prepare(context.Background(), c.Actor, rbac.ActionRead, obType)
+				require.NoError(b, err)
+			}
+		})
+	}
+
 	for _, c := range benchCases {
 		b.Run(c.Name, func(b *testing.B) {
 			objects := benchmarkSetup(orgs, users, b.N)
diff --git a/coderd/rbac/builtin_internal_test.go b/coderd/rbac/builtin_internal_test.go
@@ -51,7 +51,11 @@ func BenchmarkRBACValueAllocation(b *testing.B) {
 	})
 	b.Run("JSONRegoValue", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
-			_, err := ast.InterfaceToValue(jsonSubject)
+			_, err := ast.InterfaceToValue(map[string]interface{}{
+				"subject": jsonSubject,
+				"action":  ActionRead,
+				"object":  obj,
+			})
 			require.NoError(b, err)
 		}
 	})
diff --git a/coderd/rbac/partial.go b/coderd/rbac/partial.go
@@ -127,7 +127,7 @@ EachQueryLoop:
 		pa.subjectInput, pa.subjectAction, pa.subjectResourceType, nil)
 }
 
-func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, objectType string) (*PartialAuthorizer, error) {
+func (a RegoAuthorizer) newPartialAuthorizer(ctx context.Context, subject Subject, action Action, objectType string) (*PartialAuthorizer, error) {
 	if subject.Roles == nil {
 		return nil, xerrors.Errorf("subject must have roles")
 	}
@@ -140,20 +140,8 @@ func newPartialAuthorizer(ctx context.Context, subject Subject, action Action, o
 		return nil, xerrors.Errorf("prepare input: %w", err)
 	}
 
-	// Run the rego policy with a few unknown fields. This should simplify our
-	// policy to a set of queries.
-	partialQueries, err := rego.New(
-		rego.Query("data.authz.allow = true"),
-		rego.Module("policy.rego", policy),
-		rego.Unknowns([]string{
-			"input.object.id",
-			"input.object.owner",
-			"input.object.org_owner",
-			"input.object.acl_user_list",
-			"input.object.acl_group_list",
-		}),
-		rego.ParsedInput(input),
-	).Partial(ctx)
+	partialQueries, err := a.partialQuery.Partial(ctx, rego.EvalParsedInput(input))
+
 	if err != nil {
 		return nil, xerrors.Errorf("prepare: %w", err)
 	}
