chore: Minor rbac memory optimization (#7391)

* test: Add benchmark for static rbac roles
* static roles should only be allocated once
* A unit test that modifies the ast value should not mess with the globals
* Cache subject AST values to avoid reallocating slices

coderd/database/dbauthz/dbauthz.go

@@ -144,7 +144,8 @@ var (
 			},
 		}),
 		Scope: rbac.ScopeAll,
-	}
+	}.WithCachedASTValue()
+
 	subjectAutostart = rbac.Subject{
 		ID: uuid.Nil.String(),
 		Roles: rbac.Roles([]rbac.Role{
@@ -161,7 +162,8 @@ var (
 			},
 		}),
 		Scope: rbac.ScopeAll,
-	}
+	}.WithCachedASTValue()
+
 	subjectSystemRestricted = rbac.Subject{
 		ID: uuid.Nil.String(),
 		Roles: rbac.Roles([]rbac.Role{
@@ -188,7 +190,7 @@ var (
 			},
 		}),
 		Scope: rbac.ScopeAll,
-	}
+	}.WithCachedASTValue()
 )
 
 // AsProvisionerd returns a context with an actor that has permissions required

coderd/httpmw/apikey.go

@@ -379,7 +379,7 @@ func ExtractAPIKey(rw http.ResponseWriter, r *http.Request, cfg ExtractAPIKeyCon
 			Roles:  rbac.RoleNames(roles.Roles),
 			Groups: roles.Groups,
 			Scope:  rbac.ScopeName(key.Scope),
-		},
+		}.WithCachedASTValue(),
 	}
 
 	return &key, &authz, true

coderd/httpmw/workspaceagent.go

@@ -110,5 +110,5 @@ func getAgentSubject(ctx context.Context, db database.Store, agent database.Work
 		Roles:  rbac.RoleNames(roles.Roles),
 		Groups: roles.Groups,
 		Scope:  rbac.WorkspaceAgentScope(workspace.ID, user.ID),
-	}, nil
+	}.WithCachedASTValue(), nil
 }

coderd/rbac/astvalue.go

@@ -65,6 +65,10 @@ func regoPartialInputValue(subject Subject, action Action, objectType string) (a
 
 // regoValue returns the ast.Object representation of the subject.
 func (s Subject) regoValue() (ast.Value, error) {
+	if s.cachedASTValue != nil {
+		return s.cachedASTValue, nil
+	}
+
 	subjRoles, err := s.Roles.Expand()
 	if err != nil {
 		return nil, xerrors.Errorf("expand roles: %w", err)
@@ -133,7 +137,20 @@ func (z Object) regoValue() ast.Value {
 	)
 }
 
+// withCachedRegoValue returns a copy of the role with the cachedRegoValue.
+// It does not mutate the underlying role.
+// Avoid using this function if possible, it should only be used if the
+// caller can guarantee the role is static and will never change.
+func (role Role) withCachedRegoValue() Role {
+	tmp := role
+	tmp.cachedRegoValue = role.regoValue()
+	return tmp
+}
+
 func (role Role) regoValue() ast.Value {
+	if role.cachedRegoValue != nil {
+		return role.cachedRegoValue
+	}
 	orgMap := ast.NewObject()
 	for k, p := range role.Org {
 		orgMap.Insert(ast.StringTerm(k), ast.NewTerm(regoSlice(p)))

coderd/rbac/authz.go

@@ -49,6 +49,20 @@ type Subject struct {
 	Roles  ExpandableRoles
 	Groups []string
 	Scope  ExpandableScope
+
+	// cachedASTValue is the cached ast value for this subject.
+	cachedASTValue ast.Value
+}
+
+// WithCachedASTValue can be called if the subject is static. This will compute
+// the ast value once and cache it for future calls.
+func (s Subject) WithCachedASTValue() Subject {
+	tmp := s
+	v, err := tmp.regoValue()
+	if err == nil {
+		tmp.cachedASTValue = v
+	}
+	return tmp
 }
 
 func (s Subject) Equal(b Subject) bool {

coderd/rbac/authz_test.go

@@ -86,6 +86,21 @@ func benchmarkUserCases() (cases []benchmarkCase, users uuid.UUID, orgs []uuid.U
 				Groups: noiseGroups,
 			},
 		},
+		{
+			Name: "ManyRolesCachedSubject",
+			Actor: rbac.Subject{
+				// Admin of many orgs
+				Roles: rbac.RoleNames{
+					rbac.RoleOrgMember(orgs[0]), rbac.RoleOrgAdmin(orgs[0]),
+					rbac.RoleOrgMember(orgs[1]), rbac.RoleOrgAdmin(orgs[1]),
+					rbac.RoleOrgMember(orgs[2]), rbac.RoleOrgAdmin(orgs[2]),
+					rbac.RoleMember(),
+				},
+				ID:     user.String(),
+				Scope:  rbac.ScopeAll,
+				Groups: noiseGroups,
+			}.WithCachedASTValue(),
+		},
 		{
 			Name: "AdminWithScope",
 			Actor: rbac.Subject{
@@ -96,13 +111,41 @@ func benchmarkUserCases() (cases []benchmarkCase, users uuid.UUID, orgs []uuid.U
 				Groups: noiseGroups,
 			},
 		},
+		{
+			// This test should only use static roles. AKA no org roles.
+			Name: "StaticRoles",
+			Actor: rbac.Subject{
+				// Give some extra roles that an admin might have
+				Roles: rbac.RoleNames{
+					"auditor", rbac.RoleOwner(), rbac.RoleMember(),
+					rbac.RoleTemplateAdmin(), rbac.RoleUserAdmin(),
+				},
+				ID:     user.String(),
+				Scope:  rbac.ScopeAll,
+				Groups: noiseGroups,
+			},
+		},
+		{
+			// This test should only use static roles. AKA no org roles.
+			Name: "StaticRolesWithCache",
+			Actor: rbac.Subject{
+				// Give some extra roles that an admin might have
+				Roles: rbac.RoleNames{
+					"auditor", rbac.RoleOwner(), rbac.RoleMember(),
+					rbac.RoleTemplateAdmin(), rbac.RoleUserAdmin(),
+				},
+				ID:     user.String(),
+				Scope:  rbac.ScopeAll,
+				Groups: noiseGroups,
+			}.WithCachedASTValue(),
+		},
 	}
 	return benchCases, users, orgs
 }
 
 // BenchmarkRBACAuthorize benchmarks the rbac.Authorize method.
 //
-//	go test -bench BenchmarkRBACAuthorize -benchmem -memprofile memprofile.out -cpuprofile profile.out
+//	go test -run=^$ -bench BenchmarkRBACAuthorize -benchmem -memprofile memprofile.out -cpuprofile profile.out
 func BenchmarkRBACAuthorize(b *testing.B) {
 	benchCases, user, orgs := benchmarkUserCases()
 	users := append([]uuid.UUID{},
@@ -111,6 +154,9 @@ func BenchmarkRBACAuthorize(b *testing.B) {
 		uuid.MustParse("0632b012-49e0-4d70-a5b3-f4398f1dcd52"),
 		uuid.MustParse("70dbaa7a-ea9c-4f68-a781-97b08af8461d"),
 	)
+
+	// There is no caching that occurs because a fresh context is used for each
+	// call. And the context needs 'WithCacheCtx' to work.
 	authorizer := rbac.NewCachingAuthorizer(prometheus.NewRegistry())
 	// This benchmarks all the simple cases using just user permissions. Groups
 	// are added as noise, but do not do anything.

coderd/rbac/roles.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/google/uuid"
+	"github.com/open-policy-agent/opa/ast"
 
 	"golang.org/x/xerrors"
 )
@@ -128,86 +129,100 @@ func ReloadBuiltinRoles(opts *RoleOptions) {
 		)
 	}
 
+	// Static roles that never change should be allocated in a closure.
+	// This is to ensure these data structures are only allocated once and not
+	// on every authorize call. 'withCachedRegoValue' can be used as well to
+	// preallocate the rego value that is used by the rego eval engine.
+	ownerRole := Role{
+		Name:        owner,
+		DisplayName: "Owner",
+		Site:        allPermsExcept(ownerAndAdminExceptions...),
+		Org:         map[string][]Permission{},
+		User:        []Permission{},
+	}.withCachedRegoValue()
+
+	memberRole := Role{
+		Name:        member,
+		DisplayName: "",
+		Site: Permissions(map[string][]Action{
+			// All users can read all other users and know they exist.
+			ResourceUser.Type:           {ActionRead},
+			ResourceRoleAssignment.Type: {ActionRead},
+			// All users can see the provisioner daemons.
+			ResourceProvisionerDaemon.Type: {ActionRead},
+		}),
+		Org:  map[string][]Permission{},
+		User: allPermsExcept(),
+	}.withCachedRegoValue()
+
+	auditorRole := Role{
+		Name:        auditor,
+		DisplayName: "Auditor",
+		Site: Permissions(map[string][]Action{
+			// Should be able to read all template details, even in orgs they
+			// are not in.
+			ResourceTemplate.Type: {ActionRead},
+			ResourceAuditLog.Type: {ActionRead},
+		}),
+		Org:  map[string][]Permission{},
+		User: []Permission{},
+	}.withCachedRegoValue()
+
+	templateAdminRole := Role{
+		Name:        templateAdmin,
+		DisplayName: "Template Admin",
+		Site: Permissions(map[string][]Action{
+			ResourceTemplate.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
+			// CRUD all files, even those they did not upload.
+			ResourceFile.Type:      {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
+			ResourceWorkspace.Type: {ActionRead},
+			// CRUD to provisioner daemons for now.
+			ResourceProvisionerDaemon.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
+			// Needs to read all organizations since
+			ResourceOrganization.Type: {ActionRead},
+		}),
+		Org:  map[string][]Permission{},
+		User: []Permission{},
+	}.withCachedRegoValue()
+
+	userAdminRole := Role{
+		Name:        userAdmin,
+		DisplayName: "User Admin",
+		Site: Permissions(map[string][]Action{
+			ResourceRoleAssignment.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
+			ResourceUser.Type:           {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
+			// Full perms to manage org members
+			ResourceOrganizationMember.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
+			ResourceGroup.Type:              {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
+		}),
+		Org:  map[string][]Permission{},
+		User: []Permission{},
+	}.withCachedRegoValue()
+
 	builtInRoles = map[string]func(orgID string) Role{
 		// admin grants all actions to all resources.
 		owner: func(_ string) Role {
-			return Role{
-				Name:        owner,
-				DisplayName: "Owner",
-				Site:        allPermsExcept(ownerAndAdminExceptions...),
-				Org:         map[string][]Permission{},
-				User:        []Permission{},
-			}
+			return ownerRole
 		},
 
 		// member grants all actions to all resources owned by the user
 		member: func(_ string) Role {
-			return Role{
-				Name:        member,
-				DisplayName: "",
-				Site: Permissions(map[string][]Action{
-					// All users can read all other users and know they exist.
-					ResourceUser.Type:           {ActionRead},
-					ResourceRoleAssignment.Type: {ActionRead},
-					// All users can see the provisioner daemons.
-					ResourceProvisionerDaemon.Type: {ActionRead},
-				}),
-				Org:  map[string][]Permission{},
-				User: allPermsExcept(),
-			}
+			return memberRole
 		},
 
 		// auditor provides all permissions required to effectively read and understand
 		// audit log events.
 		// TODO: Finish the auditor as we add resources.
 		auditor: func(_ string) Role {
-			return Role{
-				Name:        auditor,
-				DisplayName: "Auditor",
-				Site: Permissions(map[string][]Action{
-					// Should be able to read all template details, even in orgs they
-					// are not in.
-					ResourceTemplate.Type: {ActionRead},
-					ResourceAuditLog.Type: {ActionRead},
-				}),
-				Org:  map[string][]Permission{},
-				User: []Permission{},
-			}
+			return auditorRole
 		},
 
 		templateAdmin: func(_ string) Role {
-			return Role{
-				Name:        templateAdmin,
-				DisplayName: "Template Admin",
-				Site: Permissions(map[string][]Action{
-					ResourceTemplate.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
-					// CRUD all files, even those they did not upload.
-					ResourceFile.Type:      {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
-					ResourceWorkspace.Type: {ActionRead},
-					// CRUD to provisioner daemons for now.
-					ResourceProvisionerDaemon.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
-					// Needs to read all organizations since
-					ResourceOrganization.Type: {ActionRead},
-				}),
-				Org:  map[string][]Permission{},
-				User: []Permission{},
-			}
+			return templateAdminRole
 		},
 
 		userAdmin: func(_ string) Role {
-			return Role{
-				Name:        userAdmin,
-				DisplayName: "User Admin",
-				Site: Permissions(map[string][]Action{
-					ResourceRoleAssignment.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
-					ResourceUser.Type:           {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
-					// Full perms to manage org members
-					ResourceOrganizationMember.Type: {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
-					ResourceGroup.Type:              {ActionCreate, ActionRead, ActionUpdate, ActionDelete},
-				}),
-				Org:  map[string][]Permission{},
-				User: []Permission{},
-			}
+			return userAdminRole
 		},
 
 		// orgAdmin returns a role with all actions allows in a given
@@ -333,6 +348,10 @@ type Role struct {
 	// roles.
 	Org  map[string][]Permission `json:"org"`
 	User []Permission            `json:"user"`
+
+	// cachedRegoValue can be used to cache the rego value for this role.
+	// This is helpful for static roles that never change.
+	cachedRegoValue ast.Value
 }
 
 type Roles []Role

coderd/rbac/roles_internal_test.go

@@ -68,8 +68,19 @@ func BenchmarkRBACValueAllocation(b *testing.B) {
 func TestRegoInputValue(t *testing.T) {
 	t.Parallel()
 
+	// Expand all roles and make sure we have a good copy.
+	// This is because these tests modify the roles, and we don't want to
+	// modify the original roles.
+	roles, err := RoleNames{RoleOrgMember(uuid.New()), RoleOrgAdmin(uuid.New()), RoleMember()}.Expand()
+	require.NoError(t, err, "failed to expand roles")
+	for i := range roles {
+		// If all cached values are nil, then the role will not use
+		// the shared cached value.
+		roles[i].cachedRegoValue = nil
+	}
+
 	actor := Subject{
-		Roles:  RoleNames{RoleOrgMember(uuid.New()), RoleOrgAdmin(uuid.New()), RoleMember()},
+		Roles:  Roles(roles),
 		ID:     uuid.NewString(),
 		Scope:  ScopeAll,
 		Groups: []string{uuid.NewString(), uuid.NewString(), uuid.NewString()},