opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf: reduce calls to GetWorkspaceByAgentID in GetWorkspaceAgentByID (#21046)

Browse Source 
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>
This commit is contained in:
Callum Styan
2025-12-10 14:03:24 -08:00
committed by GitHub
parent 8e460ca865
commit 8ed1c1d372
8 changed files with 282 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/agentapi/api.go
+1
View File
@@ -197,6 +197,7 @@ func New(opts Options, workspace database.Workspace) *API {
AgentFn: api.agent,
ConnectionLogger: opts.ConnectionLogger,
Database: opts.Database,
Workspace: api.cachedWorkspaceFields,
Log: opts.Log,
}
coderd/agentapi/connectionlog.go
+27 -8
View File
@@ -14,11 +14,13 @@ import (
"github.com/coder/coder/v2/coderd/connectionlog"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/db2sdk"
"github.com/coder/coder/v2/coderd/database/dbauthz"
)
type ConnLogAPI struct {
AgentFn func(context.Context) (database.WorkspaceAgent, error)
ConnectionLogger *atomic.Pointer[connectionlog.ConnectionLogger]
Workspace *CachedWorkspaceFields
Database database.Store
Log slog.Logger
}
@@ -51,14 +53,31 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
}
}
// Inject RBAC object into context for dbauthz fast path, avoid having to
// call GetWorkspaceByAgentID on every metadata update.
rbacCtx := ctx
var ws database.WorkspaceIdentity
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
ws = dbws
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
if err != nil {
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
//nolint:gocritic
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
}
}
// Fetch contextual data for this connection log event.
workspaceAgent, err := a.AgentFn(ctx)
workspaceAgent, err := a.AgentFn(rbacCtx)
if err != nil {
return nil, xerrors.Errorf("get agent: %w", err)
}
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
if err != nil {
return nil, xerrors.Errorf("get workspace by agent id: %w", err)
if ws.Equal(database.WorkspaceIdentity{}) {
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
if err != nil {
return nil, xerrors.Errorf("get workspace by agent id: %w", err)
}
ws = database.WorkspaceIdentityFromWorkspace(workspace)
}
// 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
err = connLogger.Upsert(ctx, database.UpsertConnectionLogParams{
ID: uuid.New(),
Time: req.GetConnection().GetTimestamp().AsTime(),
OrganizationID: workspace.OrganizationID,
WorkspaceOwnerID: workspace.OwnerID,
WorkspaceID: workspace.ID,
WorkspaceName: workspace.Name,
OrganizationID: ws.OrganizationID,
WorkspaceOwnerID: ws.OwnerID,
WorkspaceID: ws.ID,
WorkspaceName: ws.Name,
AgentName: workspaceAgent.Name,
Type: connectionType,
Code: code,
coderd/agentapi/connectionlog_test.go
+1
View File
@@ -117,6 +117,7 @@ func TestConnectionLog(t *testing.T) {
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
Workspace: &agentapi.CachedWorkspaceFields{},
}
api.ReportConnection(context.Background(), &agentproto.ReportConnectionRequest{
Connection: &agentproto.Connection{
coderd/agentapi/metadata.go
+14 -13
View File
@@ -47,7 +47,20 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
maxErrorLen = maxValueLen
)
workspaceAgent, err := a.AgentFn(ctx)
// Inject RBAC object into context for dbauthz fast path, avoid having to
// call GetWorkspaceByAgentID on every metadata update.
var err error
rbacCtx := ctx
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
if err != nil {
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
//nolint:gocritic
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
}
}
workspaceAgent, err := a.AgentFn(rbacCtx)
if err != nil {
return nil, err
}
@@ -109,18 +122,6 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
)
}
// Inject RBAC object into context for dbauthz fast path, avoid having to
// call GetWorkspaceByAgentID on every metadata update.
rbacCtx := ctx
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
if err != nil {
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
//nolint:gocritic
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
}
}
err = a.Database.UpdateWorkspaceAgentMetadata(rbacCtx, dbUpdate)
if err != nil {
return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)
coderd/agentapi/metadata_test.go
+130 -6
View File
@@ -295,6 +295,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
now = dbtime.Now()
// Set up consistent IDs that represent a valid workspace->agent relationship
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111")
agentID = uuid.MustParse("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
@@ -358,8 +359,48 @@ func TestBatchUpdateMetadata(t *testing.T) {
OrganizationID: orgID,
})
// Create context with system actor so authorization passes
ctx := dbauthz.AsSystemRestricted(context.Background())
// Create roles with workspace permissions
userRoles := rbac.Roles([]rbac.Role{
{
Identifier: rbac.RoleMember(),
User: []rbac.Permission{
{
Negate: false,
ResourceType: rbac.ResourceWorkspace.Type,
Action: policy.WildcardSymbol,
},
},
ByOrgID: map[string]rbac.OrgPermissions{
orgID.String(): {
Member: []rbac.Permission{
{
Negate: false,
ResourceType: rbac.ResourceWorkspace.Type,
Action: policy.WildcardSymbol,
},
},
},
},
},
})
agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
WorkspaceID: workspaceID,
OwnerID: ownerID,
TemplateID: templateID,
VersionID: uuid.New(),
})
ctx := dbauthz.As(context.Background(), rbac.Subject{
Type: rbac.SubjectTypeUser,
FriendlyName: "testuser",
Email: "testuser@example.com",
ID: ownerID.String(),
Roles: userRoles,
Groups: []string{orgID.String()},
Scope: agentScope,
}.WithCachedASTValue())
resp, err := api.BatchUpdateMetadata(ctx, req)
require.NoError(t, err)
require.NotNil(t, resp)
@@ -376,6 +417,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
pub = &fakePublisher{}
now = dbtime.Now()
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111")
agentID = uuid.MustParse("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb")
@@ -445,12 +487,53 @@ func TestBatchUpdateMetadata(t *testing.T) {
OrganizationID: uuid.Nil, // Invalid: fails dbauthz fast path validation
})
// Create context with system actor so authorization passes
ctx := dbauthz.AsSystemRestricted(context.Background())
// Create roles with workspace permissions
userRoles := rbac.Roles([]rbac.Role{
{
Identifier: rbac.RoleMember(),
User: []rbac.Permission{
{
Negate: false,
ResourceType: rbac.ResourceWorkspace.Type,
Action: policy.WildcardSymbol,
},
},
ByOrgID: map[string]rbac.OrgPermissions{
orgID.String(): {
Member: []rbac.Permission{
{
Negate: false,
ResourceType: rbac.ResourceWorkspace.Type,
Action: policy.WildcardSymbol,
},
},
},
},
},
})
agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
WorkspaceID: workspaceID,
OwnerID: ownerID,
TemplateID: templateID,
VersionID: uuid.New(),
})
ctx := dbauthz.As(context.Background(), rbac.Subject{
Type: rbac.SubjectTypeUser,
FriendlyName: "testuser",
Email: "testuser@example.com",
ID: ownerID.String(),
Roles: userRoles,
Groups: []string{orgID.String()},
Scope: agentScope,
}.WithCachedASTValue())
resp, err := api.BatchUpdateMetadata(ctx, req)
require.NoError(t, err)
require.NotNil(t, resp)
})
// Test RBAC slow path - no RBAC object in context
// This test verifies that when no RBAC object is present in context, the dbauthz layer
// falls back to the slow path and calls GetWorkspaceByAgentID.
@@ -463,6 +546,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
pub = &fakePublisher{}
now = dbtime.Now()
workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
templateID = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
ownerID = uuid.MustParse("87654321-4321-4321-4321-210987654321")
orgID = uuid.MustParse("11111111-1111-1111-1111-111111111111")
agentID = uuid.MustParse("dddddddd-dddd-dddd-dddd-dddddddddddd")
@@ -523,8 +607,48 @@ func TestBatchUpdateMetadata(t *testing.T) {
},
}
// Create context with system actor so authorization passes
ctx := dbauthz.AsSystemRestricted(context.Background())
// Create roles with workspace permissions
userRoles := rbac.Roles([]rbac.Role{
{
Identifier: rbac.RoleMember(),
User: []rbac.Permission{
{
Negate: false,
ResourceType: rbac.ResourceWorkspace.Type,
Action: policy.WildcardSymbol,
},
},
ByOrgID: map[string]rbac.OrgPermissions{
orgID.String(): {
Member: []rbac.Permission{
{
Negate: false,
ResourceType: rbac.ResourceWorkspace.Type,
Action: policy.WildcardSymbol,
},
},
},
},
},
})
agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
WorkspaceID: workspaceID,
OwnerID: ownerID,
TemplateID: templateID,
VersionID: uuid.New(),
})
ctx := dbauthz.As(context.Background(), rbac.Subject{
Type: rbac.SubjectTypeUser,
FriendlyName: "testuser",
Email: "testuser@example.com",
ID: ownerID.String(),
Roles: userRoles,
Groups: []string{orgID.String()},
Scope: agentScope,
}.WithCachedASTValue())
resp, err := api.BatchUpdateMetadata(ctx, req)
require.NoError(t, err)
require.NotNil(t, resp)
coderd/agentapi/stats.go
+16 -1
View File
@@ -10,6 +10,7 @@ import (
"cdr.dev/slog"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/workspacestats"
"github.com/coder/coder/v2/codersdk"
@@ -43,7 +44,21 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
return res, nil
}
workspaceAgent, err := a.AgentFn(ctx)
// Inject RBAC object into context for dbauthz fast path, avoid having to
// call GetWorkspaceAgentByID on every stats update.
rbacCtx := ctx
if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
var err error
rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject())
if err != nil {
// Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID.
//nolint:gocritic
a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err))
}
}
workspaceAgent, err := a.AgentFn(rbacCtx)
if err != nil {
return nil, err
}