perf: reduce DB calls to `GetWorkspaceByAgentID` via caching workspace info #20662

cstyan · 2025-11-04T00:07:12Z

This query is still relatively expensive because of its call volume, at least 1 million times per day over the last week. After #20430 the two main remaining call paths are via the Workspace Agents Stats and Metadata APIs.

Disclaimer: Blink helped write tests, identify usage of database.Workspace fields in downstream functions, identify existing pattern for and write the ticker function, and write comments.

This PR adds caching of the Workspace information on the AgentAPI struct for usage in both child APIs, Notably, for the Metadata API we need to also implement a "fast path" for the section that currently calls GetWorkspaceByAgentID since it's nested within a dbauthz call. We do this by attaching the relevant RBAC object from the Workspace object to the request context, then checking in UpdateWorkspaceAgentMetadata if we have the RBAC object in the context and if it is correct for authorization purposes. If we do not have the object or it is invalid we fall back to still calling GetWorkspaceByAgentID.

Note that while the workspace agent API struct is constructed at startup and (IIUC) is persistent for the lifetime of the agent/workspace, and makes a call to GetWorkspaceByAgentID during this process, we also need to handle updating of some of the cached fields since the entries within the database can change as part of the workspace prebuild process. We could potentially due this via subscribing, via pubsub, to updates about prebuild workspace claims. However, for now this has simply been solved with a 5m ticker to make a call to GetWorkspaceByAgentID.

This should mean we on average call GetWorkspaceByAgentID once every 5 minutes per agent, as opposed to multiple times per-minute per agent for the Stats API and again for the Metadata API. Even in the worst case scenario of a workspace prebuild being claimed at t=N and the next ticker not occurring until t=N+5minutes we should have at most 5 minutes of the current behaviour for each agent in a workspace.

NOTE: there is a remaining issue here; caching of the workspace for stats as it is currently implemented is actually invalid, at least for the ~5 minute period between when a prebuild is claimed and the next update of the ticker the cached Owner information will be incorrect and we'll update some incorrect metrics + skip the entire if !workspace.IsPrebuild() block.

We can either:

revert this optimization here and still just always call GetWorkspaceByAgentID for every stats update
similar to option 1, but only always call GetWorkspaceByAgentID if the workspace is a prebuild, and the first time it is not we could grab the lock and update the cached workspace information
do the right thing for both stats and metadata, and hook into pubsub to get the prebuild claim update ASAP while still avoiding extra calls to GetWorkspaceByAgentID

Signed-off-by: Callum Styan <callumstyan@gmail.com>

GetWorkspaceByAgentID for agent metadata updates Signed-off-by: Callum Styan <callumstyan@gmail.com>

Emyrk · 2025-11-14T15:44:40Z

coderd/agentapi/api.go

+	AutostartSchedule sql.NullString
+	DormantAt         sql.NullTime


These can change in flight. Unsure if we need these to be up to date with any change

we actually need the AutostartSchedule for the statsAPI, and while it in theory can change after the workspace has started it seems like it will change infrequently enough that the 5m refresh time is okay here

Emyrk · 2025-11-14T15:46:32Z

coderd/agentapi/api.go

+	if err != nil {
+		a.opts.Log.Warn(ctx, "failed to refresh cached workspace fields", slog.Error(err))
+		return
+	}


This error should probably nuke the cache? If the workspace can not be fetched, this could be a permission issue for example. Then the cache is stale with the last authorized called.

I guess it depends on whether we want to always nuke it, or only nuke for permissions issues? There could be transient DB connection problems that could cause an error here, but those are likely to impact the follow up calls in the Stats/Metadata APIs anyways.

Just for security reasons, it feels safer to nuke it. If there is transient db issues, things are going to fail in other places too. Having a stale workspace is the least of our issues.

Is it easy to nuke it?

easy, nukes the cache now via call to Clear()

coderd/agentapi/api.go

coderd/agentapi/metadata.go

coderd/database/dbauthz/dbauthz.go

Emyrk · 2025-11-14T16:01:10Z

coderd/database/dbauthz/dbauthz.go

+			act, ok := ActorFromContext(ctx)
+			if !ok {
+				return ErrNoActor
+			}
+			err := q.auth.Authorize(ctx, act, policy.ActionUpdate, rbacObj)
+			if err == nil {
+				return q.db.UpdateWorkspaceAgentMetadata(ctx, arg)
+			}


Suggested change

act, ok := ActorFromContext(ctx)

if !ok {

return ErrNoActor

}

err := q.auth.Authorize(ctx, act, policy.ActionUpdate, rbacObj)

if err == nil {

return q.db.UpdateWorkspaceAgentMetadata(ctx, arg)

}

err := q.authorizeContext(ctx, policy.ActionUpdate, rbacObj)

if err == nil {

return q.db.UpdateWorkspaceAgentMetadata(ctx, arg)

}

we don't need the actor here?

The q.authorizeContext handles the actor for you. See

coder/coderd/database/dbauthz/dbauthz.go

Lines 138 to 150 in afe37c9

// authorizeContext is a helper function to authorize an action on an object.

func (q *querier) authorizeContext(ctx context.Context, action policy.Action, object rbac.Objecter) error {

act, ok := ActorFromContext(ctx)

if !ok {

return ErrNoActor

}

err := q.auth.Authorize(ctx, act, action, object.RBACObject())

if err != nil {

return logNotAuthorizedError(ctx, q.log, err)

}

return nil

}

coderd/database/dbauthz/dbauthz.go

Signed-off-by: Callum Styan <callumstyan@gmail.com>

and invalidate the cached workspace Signed-off-by: Callum Styan <callumstyan@gmail.com>

Signed-off-by: Callum Styan <callumstyan@gmail.com>

context Signed-off-by: Callum Styan <callumstyan@gmail.com>

Signed-off-by: Callum Styan <callumstyan@gmail.com>

but also not log at error level Signed-off-by: Callum Styan <callumstyan@gmail.com>

Signed-off-by: Callum Styan <callumstyan@gmail.com>

auto cache in the API constructor if the workspace is a prebuild Signed-off-by: Callum Styan <callumstyan@gmail.com>

Signed-off-by: Callum Styan <callumstyan@gmail.com>

struct, add clear + update functions and very basic tests Signed-off-by: Callum Styan <callumstyan@gmail.com>

Signed-off-by: Callum Styan <callumstyan@gmail.com>

coderd/agentapi/api.go

Emyrk · 2025-11-18T22:59:30Z

coderd/agentapi/metadata.go

+	// Inject RBAC object into context for dbauthz fast path, avoid having to
+	// call GetWorkspaceByAgentID on every metadata update.
+	rbacCtx, err := dbauthz.WithWorkspaceRBAC(ctx, a.Workspace.AsDatabaseWorkspace().RBACObject())
+	if err != nil {
+		// Don't use error level here; that will fail the call but in reality an invalid object here is okay
+		// since we will then just fallback to the actual GetWorkspaceByAgentID call in dbauthz.
+		a.Log.Warn(ctx, "cached RBAC Workspace context wrapping was invalid", slog.Error(err))
+	}


Isn't this empty for prebuilds, so WithWorkspaceRBAC will always throw a log.Warn?

We should not log a warning when the right behavior is occurring (that being a prebuild workspace is not being cached)

this now only logs if the workspace was cached and the rbac object was still somehow invalid

Emyrk · 2025-11-18T23:00:14Z

coderd/agentapi/stats.go

-		return nil, xerrors.Errorf("get workspace by agent ID %q: %w", workspaceAgent.ID, err)
+
+	// If cache is empty (prebuild or invalid), fall back to DB
+	ws := a.Workspace.AsDatabaseWorkspace()


What if AsDatabaseWorkspace returned a *database.Workspace. Or database.Workspace, ok?

Knowing to check workspace.ID is not ideal.

Emyrk · 2025-11-18T23:00:46Z

coderd/database/dbauthz/dbauthz.go

+	// NOTE: The cached RBAC object is refreshed every 5 minutes. After a prebuild
+	// claim, there may be up to a 5-minute window where this uses stale owner_id.


The owner is no longer stale because prebuilds are not cached.

Emyrk · 2025-11-18T23:01:54Z

coderd/rbac/object.go

+func (z Object) IsEmpty() bool {
+	return z.Owner == "" || z.Owner == uuid.Nil.String() || z.OrgID == "" || z.OrgID == uuid.Nil.String()
+}


This is not correct for all objects. Some objects have no owner, some have no org. Some have neither (like deployment config)

Emyrk · 2025-11-18T23:03:04Z

coderd/database/dbauthz/dbauthz.go

+	if rbacObj, ok := WorkspaceRBACFromContext(ctx); ok {
+		// Verify the RBAC object is valid (has required fields).
+		if !rbacObj.IsEmpty() && rbacObj.Type == rbac.ResourceWorkspace.Type {


WorkspaceRBACFromContext should never return an invalid RBACObject imo. The WithWorkspaceRBAC protects against these errors like the wrong type.

So we can trust it and not guard against things if ok == true

Signed-off-by: Callum Styan <callumstyan@gmail.com>

valid/present checks, plus ensure proper fallback (less noisy logging) when a cached workspace is not found/valid in dbauthz Signed-off-by: Callum Styan <callumstyan@gmail.com>

object types there Signed-off-by: Callum Styan <callumstyan@gmail.com>

Signed-off-by: Callum Styan <callumstyan@gmail.com>

Emyrk · 2025-11-24T18:07:28Z

coderd/agentapi/api.go

+	}
+
+	if ws.IsPrebuild() {
+		a.opts.Log.Debug(ctx, "workspace is a prebuild, not caching in AgentAPI")


We probably don't need this log, your call as the logs are infrequent

Emyrk · 2025-11-24T18:08:48Z

coderd/agentapi/api.go

+	a.opts.Log.Debug(ctx, "refreshed cached workspace fields",
+		slog.F("workspace_id", ws.ID),
+		slog.F("owner_id", ws.OwnerID),
+		slog.F("name", ws.Name))


Should we only log if the update yields a different workspace then what we had?

Not a huge deal as the logs are infrequent. Just trying to keep the logs down.

Emyrk · 2025-11-24T18:17:24Z

coderd/agentapi/stats.go

 		ctx,
 		a.now(),
-		workspace,
+		ws,


It is unfortunate we are passing a possibly incomplete workspace to this function. The function does not know some workspace fields might not be present.

Can we add a concrete type, maybe to database/modelmethods.go of this subset type of a workspace? Then pass around this subset instead of database.Workspace. It can implement RBACObject like:

func (w WorkspaceIdentity) RBACObject() rbac.Object { return Workspace{ // .. Copy fields }.RBACObject() }

So the rbac stuff is kept in check, and the ReportAgentStats does not have the possibility of using unset fields.

Clear() becomes a bit easier,

cws.WorkspaceIdentity = database.WorkspaceIdentity{}

workspace info. can't misuse database.Workspace and access fields which in theory may not be set Signed-off-by: Callum Styan <callumstyan@gmail.com>

Emyrk

The WorkspaceIdentity does add the type safety. 👍

Signed-off-by: Callum Styan <callumstyan@gmail.com>

test that fields are updated via the 5m update Signed-off-by: Callum Styan <callumstyan@gmail.com>

Signed-off-by: Callum Styan <callumstyan@gmail.com>

cstyan added 2 commits October 29, 2025 22:06

remove unnecessary calls to GetWorkspaceAgentByID for stats api calls

77f9c6a

Signed-off-by: Callum Styan <callumstyan@gmail.com>

modify workspace caching, significantly reduce calls to

0bf08af

GetWorkspaceByAgentID for agent metadata updates Signed-off-by: Callum Styan <callumstyan@gmail.com>

cstyan changed the title ~~perf: significantly reduce DB calls to GetWorkspaceByAgentID via caching workspace information in Agent API~~ WIP perf: significantly reduce DB calls to GetWorkspaceByAgentID via caching workspace information in Agent API Nov 4, 2025

cstyan requested a review from Emyrk November 4, 2025 00:07

github-actions bot assigned cstyan Nov 4, 2025

cstyan mentioned this pull request Nov 6, 2025

[Epic] Most Expensive Queries coder/internal#729

Open

16 tasks

github-actions bot added the stale This issue is like stale bread. label Nov 12, 2025

Emyrk reviewed Nov 14, 2025

View reviewed changes

Add an IsEmpty check for rbac object.

f447927

Signed-off-by: Callum Styan <callumstyan@gmail.com>

github-actions bot removed the stale This issue is like stale bread. label Nov 15, 2025

cstyan added 7 commits November 17, 2025 19:18

don't need to wait on ctx after tickerfunc, plus log when we've exited

9561447

and invalidate the cached workspace Signed-off-by: Callum Styan <callumstyan@gmail.com>

only wrap context for the one database call

c5f0c60

Signed-off-by: Callum Styan <callumstyan@gmail.com>

formatting

f77df04

Signed-off-by: Callum Styan <callumstyan@gmail.com>

guard against the rbac object being invalid when adding it to the

c39830b

context Signed-off-by: Callum Styan <callumstyan@gmail.com>

lint

8ae5993

Signed-off-by: Callum Styan <callumstyan@gmail.com>

fix tests never finishing; we need to wait on the tickerFunc finishing

e99a8c5

but also not log at error level Signed-off-by: Callum Styan <callumstyan@gmail.com>

don't cache prebuilds

b8d3604

Signed-off-by: Callum Styan <callumstyan@gmail.com>

cstyan force-pushed the callum/workspace-agent-db-call-pt2 branch from 3ea6865 to b8d3604 Compare November 18, 2025 00:45

cstyan added 6 commits November 18, 2025 01:04

use the opts workspaceID in refreshCachedWorkspace now that we don't

31432ed

auto cache in the API constructor if the workspace is a prebuild Signed-off-by: Callum Styan <callumstyan@gmail.com>

these shouldn't have been removed

40e423a

Signed-off-by: Callum Styan <callumstyan@gmail.com>

lint and fmt + move workspace rbac context stuff to a separate file

dac1117

Signed-off-by: Callum Styan <callumstyan@gmail.com>

store pointer to cached workspace instead of functions

bdf076a

Signed-off-by: Callum Styan <callumstyan@gmail.com>

we need to cache the AutostartSchedule for the stats api

a77b310

Signed-off-by: Callum Styan <callumstyan@gmail.com>

move cached workspace struct to its own file, include mutex on the

02d84bd

struct, add clear + update functions and very basic tests Signed-off-by: Callum Styan <callumstyan@gmail.com>

cstyan changed the title ~~WIP perf: significantly reduce DB calls to GetWorkspaceByAgentID via caching workspace information in Agent API~~ perf: significantly reduce DB calls to GetWorkspaceByAgentID via caching workspace information in Agent API Nov 18, 2025

make fmt and linter happy

afe37c9

Signed-off-by: Callum Styan <callumstyan@gmail.com>

cstyan changed the title ~~perf: significantly reduce DB calls to GetWorkspaceByAgentID via caching workspace information in Agent API~~ perf: reduce DB calls to GetWorkspaceByAgentID via caching workspace info Nov 18, 2025

cstyan mentioned this pull request Nov 18, 2025

bug: GetWorkspaceByAgentID is called over a million times a day on dogfood coder/internal#728

Open

Emyrk reviewed Nov 18, 2025

View reviewed changes

coderd/agentapi/api.go Outdated Show resolved Hide resolved

Emyrk reviewed Nov 18, 2025

View reviewed changes

cstyan added 3 commits November 21, 2025 17:40

make cached workspace a pointer on the API struct

15081e8

Signed-off-by: Callum Styan <callumstyan@gmail.com>

clear the cached workspace on DB error

089a5ef

Signed-off-by: Callum Styan <callumstyan@gmail.com>

address more review feedback; attempt to simplify cached workspace is

a2ee19d

valid/present checks, plus ensure proper fallback (less noisy logging) when a cached workspace is not found/valid in dbauthz Signed-off-by: Callum Styan <callumstyan@gmail.com>

cstyan force-pushed the callum/workspace-agent-db-call-pt2 branch from 55ee6d3 to a2ee19d Compare November 21, 2025 18:29

cstyan added 2 commits November 21, 2025 18:43

remove IsEmpty from rbac object, we'd need to handle too many rbac

1d48500

object types there Signed-off-by: Callum Styan <callumstyan@gmail.com>

typo in function name

96eb4c0

Signed-off-by: Callum Styan <callumstyan@gmail.com>

Emyrk reviewed Nov 24, 2025

View reviewed changes

refactor so that functions which will now make use of the cached

6811090

workspace info. can't misuse database.Workspace and access fields which in theory may not be set Signed-off-by: Callum Styan <callumstyan@gmail.com>

Emyrk approved these changes Nov 25, 2025

View reviewed changes

cstyan added 3 commits November 25, 2025 18:15

fmt and lint again

81bc27a

Signed-off-by: Callum Styan <callumstyan@gmail.com>

cleanup/refactor tests, we no longer cache prebuilds but still want to

8304864

test that fields are updated via the 5m update Signed-off-by: Callum Styan <callumstyan@gmail.com>

reduce logging for expected paths

740348c

Signed-off-by: Callum Styan <callumstyan@gmail.com>

cstyan merged commit b0e8384 into main Nov 25, 2025
30 checks passed

cstyan deleted the callum/workspace-agent-db-call-pt2 branch November 25, 2025 22:45

github-actions bot locked and limited conversation to collaborators Nov 25, 2025

	// authorizeContext is a helper function to authorize an action on an object.
	func (q *querier) authorizeContext(ctx context.Context, action policy.Action, object rbac.Objecter) error {
	act, ok := ActorFromContext(ctx)
	if !ok {
	return ErrNoActor
	}

	err := q.auth.Authorize(ctx, act, action, object.RBACObject())
	if err != nil {
	return logNotAuthorizedError(ctx, q.log, err)
	}
	return nil
	}

		// NOTE: The cached RBAC object is refreshed every 5 minutes. After a prebuild
		// claim, there may be up to a 5-minute window where this uses stale owner_id.

perf: reduce DB calls to GetWorkspaceByAgentID via caching workspace info #20662

perf: reduce DB calls to GetWorkspaceByAgentID via caching workspace info #20662

Uh oh!

Conversation

cstyan commented Nov 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Emyrk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

perf: reduce DB calls to `GetWorkspaceByAgentID` via caching workspace info #20662

perf: reduce DB calls to `GetWorkspaceByAgentID` via caching workspace info #20662

cstyan commented Nov 4, 2025 •

edited

Loading