Skip to content

Commit 8ed1c1d

Browse files
cstyancstyan
authored
perf: reduce calls to GetWorkspaceByAgentID in GetWorkspaceAgentByID (#21046)
This PR piggy backs on the agent API cached workspace added in an earlier PR to provide a fast path for avoiding `GetWorkspaceByAgentID` calls in dbauthz's `GetWorkspaceAgentByID`. This query is not the most expensive, but has a significant call volume at ~16 million calls per week. Signed-off-by: Callum Styan <callumstyan@gmail.com>
1 parent 8e460ca commit 8ed1c1d

File tree

8 files changed

+282
-28
lines changed

8 files changed

+282
-28
lines changed

‎coderd/agentapi/api.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ func New(opts Options, workspace database.Workspace) *API {
197197
AgentFn: api.agent,
198198
ConnectionLogger: opts.ConnectionLogger,
199199
Database: opts.Database,
200+
Workspace: api.cachedWorkspaceFields,
200201
Log: opts.Log,
201202
}
202203

‎coderd/agentapi/connectionlog.go‎

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,13 @@ import (
1414
"github.com/coder/coder/v2/coderd/connectionlog"
1515
"github.com/coder/coder/v2/coderd/database"
1616
"github.com/coder/coder/v2/coderd/database/db2sdk"
17+
"github.com/coder/coder/v2/coderd/database/dbauthz"
1718
)
1819

1920
type ConnLogAPI struct {
2021
AgentFn func(context.Context) (database.WorkspaceAgent, error)
2122
ConnectionLogger *atomic.Pointer[connectionlog.ConnectionLogger]
23+
Workspace *CachedWorkspaceFields
2224
Database database.Store
2325
Log slog.Logger
2426
}
@@ -51,14 +53,31 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
5153
}
5254
}
5355

56+
// Inject RBAC object into context for dbauthz fast path, avoid having to
57+
// call GetWorkspaceByAgentID on every metadata update.
58+
rbacCtx := ctx
59+
var ws database.WorkspaceIdentity
60+
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
61+
ws = dbws
62+
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
63+
if err != nil {
64+
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
65+
//nolint:gocritic
66+
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
67+
}
68+
}
69+
5470
// Fetch contextual data for this connection log event.
55-
workspaceAgent, err := a.AgentFn(ctx)
71+
workspaceAgent, err := a.AgentFn(rbacCtx)
5672
if err != nil {
5773
return nil, xerrors.Errorf("get agent: %w", err)
5874
}
59-
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
60-
if err != nil {
61-
return nil, xerrors.Errorf("get workspace by agent id: %w", err)
75+
if ws.Equal(database.WorkspaceIdentity{}) {
76+
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
77+
if err != nil {
78+
return nil, xerrors.Errorf("get workspace by agent id: %w", err)
79+
}
80+
ws = database.WorkspaceIdentityFromWorkspace(workspace)
6281
}
6382

6483
// Some older clients may incorrectly report "localhost" as the IP address.
@@ -74,10 +93,10 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
7493
err = connLogger.Upsert(ctx, database.UpsertConnectionLogParams{
7594
ID: uuid.New(),
7695
Time: req.GetConnection().GetTimestamp().AsTime(),
77-
OrganizationID: workspace.OrganizationID,
78-
WorkspaceOwnerID: workspace.OwnerID,
79-
WorkspaceID: workspace.ID,
80-
WorkspaceName: workspace.Name,
96+
OrganizationID: ws.OrganizationID,
97+
WorkspaceOwnerID: ws.OwnerID,
98+
WorkspaceID: ws.ID,
99+
WorkspaceName: ws.Name,
81100
AgentName: workspaceAgent.Name,
82101
Type: connectionType,
83102
Code: code,

‎coderd/agentapi/connectionlog_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ func TestConnectionLog(t *testing.T) {
117117
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
118118
return agent, nil
119119
},
120+
Workspace: &agentapi.CachedWorkspaceFields{},
120121
}
121122
api.ReportConnection(context.Background(), &agentproto.ReportConnectionRequest{
122123
Connection: &agentproto.Connection{

‎coderd/agentapi/metadata.go‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,20 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
4747
maxErrorLen = maxValueLen
4848
)
4949

50-
workspaceAgent, err := a.AgentFn(ctx)
50+
// Inject RBAC object into context for dbauthz fast path, avoid having to
51+
// call GetWorkspaceByAgentID on every metadata update.
52+
var err error
53+
rbacCtx := ctx
54+
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
55+
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
56+
if err != nil {
57+
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
58+
//nolint:gocritic
59+
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
60+
}
61+
}
62+
63+
workspaceAgent, err := a.AgentFn(rbacCtx)
5164
if err != nil {
5265
return nil, err
5366
}
@@ -109,18 +122,6 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
109122
)
110123
}
111124

112-
// Inject RBAC object into context for dbauthz fast path, avoid having to
113-
// call GetWorkspaceByAgentID on every metadata update.
114-
rbacCtx := ctx
115-
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
116-
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
117-
if err != nil {
118-
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
119-
//nolint:gocritic
120-
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
121-
}
122-
}
123-
124125
err = a.Database.UpdateWorkspaceAgentMetadata(rbacCtx, dbUpdate)
125126
if err != nil {
126127
return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)

‎coderd/agentapi/metadata_test.go‎

Lines changed: 130 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
295295
now = dbtime.Now()
296296
// Set up consistent IDs that represent a valid workspace->agent relationship
297297
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
298+
templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
298299
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
299300
orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111")
300301
agentID = uuid.MustParse("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
@@ -358,8 +359,48 @@ func TestBatchUpdateMetadata(t *testing.T) {
358359
OrganizationID: orgID,
359360
})
360361

361-
// Create context with system actor so authorization passes
362-
ctx := dbauthz.AsSystemRestricted(context.Background())
362+
// Create roles with workspace permissions
363+
userRoles := rbac.Roles([]rbac.Role{
364+
{
365+
Identifier: rbac.RoleMember(),
366+
User: []rbac.Permission{
367+
{
368+
Negate: false,
369+
ResourceType: rbac.ResourceWorkspace.Type,
370+
Action: policy.WildcardSymbol,
371+
},
372+
},
373+
ByOrgID: map[string]rbac.OrgPermissions{
374+
orgID.String(): {
375+
Member: []rbac.Permission{
376+
{
377+
Negate: false,
378+
ResourceType: rbac.ResourceWorkspace.Type,
379+
Action: policy.WildcardSymbol,
380+
},
381+
},
382+
},
383+
},
384+
},
385+
})
386+
387+
agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
388+
WorkspaceID: workspaceID,
389+
OwnerID: ownerID,
390+
TemplateID: templateID,
391+
VersionID: uuid.New(),
392+
})
393+
394+
ctx := dbauthz.As(context.Background(), rbac.Subject{
395+
Type: rbac.SubjectTypeUser,
396+
FriendlyName: "testuser",
397+
Email: "testuser@example.com",
398+
ID: ownerID.String(),
399+
Roles: userRoles,
400+
Groups: []string{orgID.String()},
401+
Scope: agentScope,
402+
}.WithCachedASTValue())
403+
363404
resp, err := api.BatchUpdateMetadata(ctx, req)
364405
require.NoError(t, err)
365406
require.NotNil(t, resp)
@@ -376,6 +417,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
376417
pub = &fakePublisher{}
377418
now = dbtime.Now()
378419
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
420+
templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
379421
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
380422
orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111")
381423
agentID = uuid.MustParse("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb")
@@ -445,12 +487,53 @@ func TestBatchUpdateMetadata(t *testing.T) {
445487
OrganizationID: uuid.Nil, // Invalid: fails dbauthz fast path validation
446488
})
447489

448-
// Create context with system actor so authorization passes
449-
ctx := dbauthz.AsSystemRestricted(context.Background())
490+
// Create roles with workspace permissions
491+
userRoles := rbac.Roles([]rbac.Role{
492+
{
493+
Identifier: rbac.RoleMember(),
494+
User: []rbac.Permission{
495+
{
496+
Negate: false,
497+
ResourceType: rbac.ResourceWorkspace.Type,
498+
Action: policy.WildcardSymbol,
499+
},
500+
},
501+
ByOrgID: map[string]rbac.OrgPermissions{
502+
orgID.String(): {
503+
Member: []rbac.Permission{
504+
{
505+
Negate: false,
506+
ResourceType: rbac.ResourceWorkspace.Type,
507+
Action: policy.WildcardSymbol,
508+
},
509+
},
510+
},
511+
},
512+
},
513+
})
514+
515+
agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
516+
WorkspaceID: workspaceID,
517+
OwnerID: ownerID,
518+
TemplateID: templateID,
519+
VersionID: uuid.New(),
520+
})
521+
522+
ctx := dbauthz.As(context.Background(), rbac.Subject{
523+
Type: rbac.SubjectTypeUser,
524+
FriendlyName: "testuser",
525+
Email: "testuser@example.com",
526+
ID: ownerID.String(),
527+
Roles: userRoles,
528+
Groups: []string{orgID.String()},
529+
Scope: agentScope,
530+
}.WithCachedASTValue())
531+
450532
resp, err := api.BatchUpdateMetadata(ctx, req)
451533
require.NoError(t, err)
452534
require.NotNil(t, resp)
453535
})
536+
454537
// Test RBAC slow path - no RBAC object in context
455538
// This test verifies that when no RBAC object is present in context, the dbauthz layer
456539
// falls back to the slow path and calls GetWorkspaceByAgentID.
@@ -463,6 +546,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
463546
pub = &fakePublisher{}
464547
now = dbtime.Now()
465548
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
549+
templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
466550
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
467551
orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111")
468552
agentID = uuid.MustParse("dddddddd-dddd-dddd-dddd-dddddddddddd")
@@ -523,8 +607,48 @@ func TestBatchUpdateMetadata(t *testing.T) {
523607
},
524608
}
525609

526-
// Create context with system actor so authorization passes
527-
ctx := dbauthz.AsSystemRestricted(context.Background())
610+
// Create roles with workspace permissions
611+
userRoles := rbac.Roles([]rbac.Role{
612+
{
613+
Identifier: rbac.RoleMember(),
614+
User: []rbac.Permission{
615+
{
616+
Negate: false,
617+
ResourceType: rbac.ResourceWorkspace.Type,
618+
Action: policy.WildcardSymbol,
619+
},
620+
},
621+
ByOrgID: map[string]rbac.OrgPermissions{
622+
orgID.String(): {
623+
Member: []rbac.Permission{
624+
{
625+
Negate: false,
626+
ResourceType: rbac.ResourceWorkspace.Type,
627+
Action: policy.WildcardSymbol,
628+
},
629+
},
630+
},
631+
},
632+
},
633+
})
634+
635+
agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
636+
WorkspaceID: workspaceID,
637+
OwnerID: ownerID,
638+
TemplateID: templateID,
639+
VersionID: uuid.New(),
640+
})
641+
642+
ctx := dbauthz.As(context.Background(), rbac.Subject{
643+
Type: rbac.SubjectTypeUser,
644+
FriendlyName: "testuser",
645+
Email: "testuser@example.com",
646+
ID: ownerID.String(),
647+
Roles: userRoles,
648+
Groups: []string{orgID.String()},
649+
Scope: agentScope,
650+
}.WithCachedASTValue())
651+
528652
resp, err := api.BatchUpdateMetadata(ctx, req)
529653
require.NoError(t, err)
530654
require.NotNil(t, resp)

‎coderd/agentapi/stats.go‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"cdr.dev/slog"
1111
agentproto "github.com/coder/coder/v2/agent/proto"
1212
"github.com/coder/coder/v2/coderd/database"
13+
"github.com/coder/coder/v2/coderd/database/dbauthz"
1314
"github.com/coder/coder/v2/coderd/database/dbtime"
1415
"github.com/coder/coder/v2/coderd/workspacestats"
1516
"github.com/coder/coder/v2/codersdk"
@@ -43,7 +44,21 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
4344
return res, nil
4445
}
4546

46-
workspaceAgent, err := a.AgentFn(ctx)
47+
// Inject RBAC object into context for dbauthz fast path, avoid having to
48+
// call GetWorkspaceAgentByID on every stats update.
49+
50+
rbacCtx := ctx
51+
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
52+
var err error
53+
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
54+
if err != nil {
55+
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
56+
//nolint:gocritic
57+
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
58+
}
59+
}
60+
61+
workspaceAgent, err := a.AgentFn(rbacCtx)
4762
if err != nil {
4863
return nil, err
4964
}

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3560,6 +3560,21 @@ func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context
35603560
}
35613561

35623562
func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) {
3563+
// Fast path: Check if we have a workspace RBAC object in context.
3564+
// In the agent API this is set at agent connection time to avoid the expensive
3565+
// GetWorkspaceByAgentID query for every agent operation.
3566+
// NOTE: The cached RBAC object is refreshed every 5 minutes in agentapi/api.go.
3567+
if rbacObj, ok := WorkspaceRBACFromContext(ctx); ok {
3568+
// Errors here will result in falling back to GetWorkspaceByAgentID,
3569+
// in case the cached data is stale.
3570+
if err := q.authorizeContext(ctx, policy.ActionRead, rbacObj); err == nil {
3571+
return q.db.GetWorkspaceAgentByID(ctx, id)
3572+
}
3573+
q.log.Debug(ctx, "fast path authorization failed for GetWorkspaceAgentByID, using slow path",
3574+
slog.F("agent_id", id))
3575+
}
3576+
3577+
// Slow path: Fallback to fetching the workspace for authorization
35633578
if _, err := q.GetWorkspaceByAgentID(ctx, id); err != nil {
35643579
return database.WorkspaceAgent{}, err
35653580
}

0 commit comments

Comments
 (0)