Skip to content

Commit 3ea3859

Browse files
Callum Styanclaude
Callum Styan
and
claude
committed
perf(agentapi): cache agent ID and Name to reduce database queries
Previously, MetadataAPI, AppsAPI, StatsAPI, and ConnectionLogAPI called GetWorkspaceAgentByID on every request to retrieve the agent ID and Name. For deployments with 1000+ agents, this resulted in 4000+ queries/second. This change introduces agent field caching to eliminate these redundant database queries. Since agent ID and Name are static fields that never change during an agent connection lifetime, we can safely cache them when the agent API is initialized. Changes: - Add CachedAgentFields struct with thread-safe ID and Name getters - Update agentapi.New() to accept agent parameter and initialize cache - Modify MetadataAPI, AppsAPI, StatsAPI, and ConnectionLogAPI to use cached fields instead of calling AgentFn - Update all test files to initialize agent cache properly - Add fallback to AgentFn if cache is not populated (safety) This reduces database query load by ~4000 queries/second for deployments with 1000 agents, while maintaining backward compatibility with existing code paths. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
1 parent e4a06f8 commit 3ea3859

File tree

11 files changed

+194
-40
lines changed

11 files changed

+194
-40
lines changed

‎coderd/agentapi/api.go‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ type API struct {
5757
*tailnet.DRPCService
5858

5959
cachedWorkspaceFields *CachedWorkspaceFields
60+
cachedAgentFields *CachedAgentFields
6061

6162
mu sync.Mutex
6263
}
@@ -96,7 +97,7 @@ type Options struct {
9697
UpdateAgentMetricsFn func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
9798
}
9899

99-
func New(opts Options, workspace database.Workspace) *API {
100+
func New(opts Options, workspace database.Workspace, agent database.WorkspaceAgent) *API {
100101
if opts.Clock == nil {
101102
opts.Clock = quartz.NewReal()
102103
}
@@ -125,6 +126,11 @@ func New(opts Options, workspace database.Workspace) *API {
125126
api.cachedWorkspaceFields.UpdateValues(workspace)
126127
}
127128

129+
// Initialize agent cache with static fields.
130+
// These fields never change during an agent connection lifetime.
131+
api.cachedAgentFields = &CachedAgentFields{}
132+
api.cachedAgentFields.UpdateValues(agent.ID, agent.Name)
133+
128134
api.AnnouncementBannerAPI = &AnnouncementBannerAPI{
129135
appearanceFetcher: opts.AppearanceFetcher,
130136
}
@@ -150,6 +156,7 @@ func New(opts Options, workspace database.Workspace) *API {
150156

151157
api.StatsAPI = &StatsAPI{
152158
AgentFn: api.agent,
159+
Agent: api.cachedAgentFields,
153160
Workspace: api.cachedWorkspaceFields,
154161
Database: opts.Database,
155162
Log: opts.Log,
@@ -168,13 +175,15 @@ func New(opts Options, workspace database.Workspace) *API {
168175

169176
api.AppsAPI = &AppsAPI{
170177
AgentFn: api.agent,
178+
Agent: api.cachedAgentFields,
171179
Database: opts.Database,
172180
Log: opts.Log,
173181
PublishWorkspaceUpdateFn: api.publishWorkspaceUpdate,
174182
}
175183

176184
api.MetadataAPI = &MetadataAPI{
177185
AgentFn: api.agent,
186+
Agent: api.cachedAgentFields,
178187
Workspace: api.cachedWorkspaceFields,
179188
Database: opts.Database,
180189
Pubsub: opts.Pubsub,
@@ -195,6 +204,7 @@ func New(opts Options, workspace database.Workspace) *API {
195204

196205
api.ConnLogAPI = &ConnLogAPI{
197206
AgentFn: api.agent,
207+
Agent: api.cachedAgentFields,
198208
ConnectionLogger: opts.ConnectionLogger,
199209
Database: opts.Database,
200210
Workspace: api.cachedWorkspaceFields,

‎coderd/agentapi/apps.go‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,29 +14,36 @@ import (
1414

1515
type AppsAPI struct {
1616
AgentFn func(context.Context) (database.WorkspaceAgent, error)
17+
Agent *CachedAgentFields
1718
Database database.Store
1819
Log slog.Logger
1920
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent, wspubsub.WorkspaceEventKind) error
2021
}
2122

2223
func (a *AppsAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) {
23-
workspaceAgent, err := a.AgentFn(ctx)
24-
if err != nil {
25-
return nil, err
24+
// Use cached agent ID if available to avoid database query.
25+
agentID := a.Agent.ID()
26+
if agentID == uuid.Nil {
27+
// Fallback to querying the agent if cache is not populated.
28+
workspaceAgent, err := a.AgentFn(ctx)
29+
if err != nil {
30+
return nil, err
31+
}
32+
agentID = workspaceAgent.ID
2633
}
2734

2835
a.Log.Debug(ctx, "got batch app health update",
29-
slog.F("agent_id", workspaceAgent.ID.String()),
36+
slog.F("agent_id", agentID.String()),
3037
slog.F("updates", req.Updates),
3138
)
3239

3340
if len(req.Updates) == 0 {
3441
return &agentproto.BatchUpdateAppHealthResponse{}, nil
3542
}
3643

37-
apps, err := a.Database.GetWorkspaceAppsByAgentID(ctx, workspaceAgent.ID)
44+
apps, err := a.Database.GetWorkspaceAppsByAgentID(ctx, agentID)
3845
if err != nil {
39-
return nil, xerrors.Errorf("get workspace apps by agent ID %q: %w", workspaceAgent.ID, err)
46+
return nil, xerrors.Errorf("get workspace apps by agent ID %q: %w", agentID, err)
4047
}
4148

4249
var newApps []database.WorkspaceApp
@@ -97,7 +104,12 @@ func (a *AppsAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.Bat
97104
}
98105

99106
if a.PublishWorkspaceUpdateFn != nil && len(newApps) > 0 {
100-
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAppHealthUpdate)
107+
// Create minimal agent record with cached ID for publishing.
108+
// PublishWorkspaceUpdateFn only needs ID from the agent.
109+
minimalAgent := database.WorkspaceAgent{
110+
ID: agentID,
111+
}
112+
err = a.PublishWorkspaceUpdateFn(ctx, &minimalAgent, wspubsub.WorkspaceEventKindAppHealthUpdate)
101113
if err != nil {
102114
return nil, xerrors.Errorf("publish workspace update: %w", err)
103115
}

‎coderd/agentapi/apps_test.go‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,11 @@ func TestBatchUpdateAppHealths(t *testing.T) {
2121

2222
var (
2323
agent = database.WorkspaceAgent{
24-
ID: uuid.New(),
24+
ID: uuid.New(),
25+
Name: "test-agent",
2526
}
26-
app1 = database.WorkspaceApp{
27+
agentAsCacheFields = agentapi.CachedAgentFields{}
28+
app1 = database.WorkspaceApp{
2729
ID: uuid.New(),
2830
AgentID: agent.ID,
2931
Slug: "code-server-1",
@@ -43,6 +45,8 @@ func TestBatchUpdateAppHealths(t *testing.T) {
4345
}
4446
)
4547

48+
agentAsCacheFields.UpdateValues(agent.ID, agent.Name)
49+
4650
t.Run("OK", func(t *testing.T) {
4751
t.Parallel()
4852

@@ -62,6 +66,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
6266
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
6367
return agent, nil
6468
},
69+
Agent: &agentAsCacheFields,
6570
Database: dbM,
6671
Log: testutil.Logger(t),
6772
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
@@ -100,6 +105,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
100105
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
101106
return agent, nil
102107
},
108+
Agent: &agentAsCacheFields,
103109
Database: dbM,
104110
Log: testutil.Logger(t),
105111
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
@@ -139,6 +145,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
139145
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
140146
return agent, nil
141147
},
148+
Agent: &agentAsCacheFields,
142149
Database: dbM,
143150
Log: testutil.Logger(t),
144151
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
@@ -175,6 +182,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
175182
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
176183
return agent, nil
177184
},
185+
Agent: &agentAsCacheFields,
178186
Database: dbM,
179187
Log: testutil.Logger(t),
180188
PublishWorkspaceUpdateFn: nil,
@@ -204,6 +212,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
204212
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
205213
return agent, nil
206214
},
215+
Agent: &agentAsCacheFields,
207216
Database: dbM,
208217
Log: testutil.Logger(t),
209218
PublishWorkspaceUpdateFn: nil,
@@ -234,6 +243,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
234243
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
235244
return agent, nil
236245
},
246+
Agent: &agentAsCacheFields,
237247
Database: dbM,
238248
Log: testutil.Logger(t),
239249
PublishWorkspaceUpdateFn: nil,

‎coderd/agentapi/cached_agent.go‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
package agentapi
2+
3+
import (
4+
"sync"
5+
6+
"github.com/google/uuid"
7+
)
8+
9+
// CachedAgentFields contains agent data that is safe to cache for the
10+
// duration of an agent connection. These fields are used to reduce database calls
11+
// in high-frequency operations like metadata updates, stats reporting, and connection logging.
12+
//
13+
// IMPORTANT: Only static fields that never change during an agent's lifetime should be cached here.
14+
// Dynamic fields (like StartedAt, ReadyAt, LogsOverflowed) should NOT be cached as they can be
15+
// modified by API calls or external processes.
16+
type CachedAgentFields struct {
17+
lock sync.RWMutex
18+
19+
// Static fields that never change during agent connection
20+
id uuid.UUID
21+
name string
22+
}
23+
24+
// UpdateValues sets the cached agent fields. This should only be called once
25+
// at agent connection initialization.
26+
func (caf *CachedAgentFields) UpdateValues(id uuid.UUID, name string) {
27+
caf.lock.Lock()
28+
defer caf.lock.Unlock()
29+
caf.id = id
30+
caf.name = name
31+
}
32+
33+
// ID returns the cached agent ID.
34+
func (caf *CachedAgentFields) ID() uuid.UUID {
35+
caf.lock.RLock()
36+
defer caf.lock.RUnlock()
37+
return caf.id
38+
}
39+
40+
// Name returns the cached agent name.
41+
func (caf *CachedAgentFields) Name() string {
42+
caf.lock.RLock()
43+
defer caf.lock.RUnlock()
44+
return caf.name
45+
}
46+
47+
// IsPopulated returns true if the cache has been initialized with values.
48+
func (caf *CachedAgentFields) IsPopulated() bool {
49+
caf.lock.RLock()
50+
defer caf.lock.RUnlock()
51+
return caf.id != uuid.Nil
52+
}

‎coderd/agentapi/connectionlog.go‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919

2020
type ConnLogAPI struct {
2121
AgentFn func(context.Context) (database.WorkspaceAgent, error)
22+
Agent *CachedAgentFields
2223
ConnectionLogger *atomic.Pointer[connectionlog.ConnectionLogger]
2324
Workspace *CachedWorkspaceFields
2425
Database database.Store
@@ -67,13 +68,20 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
6768
}
6869
}
6970

70-
// Fetch contextual data for this connection log event.
71-
workspaceAgent, err := a.AgentFn(rbacCtx)
72-
if err != nil {
73-
return nil, xerrors.Errorf("get agent: %w", err)
71+
// Use cached agent fields if available to avoid database query.
72+
agentID := a.Agent.ID()
73+
agentName := a.Agent.Name()
74+
if agentID == uuid.Nil {
75+
// Fallback to querying the agent if cache is not populated.
76+
workspaceAgent, err := a.AgentFn(rbacCtx)
77+
if err != nil {
78+
return nil, xerrors.Errorf("get agent: %w", err)
79+
}
80+
agentID = workspaceAgent.ID
81+
agentName = workspaceAgent.Name
7482
}
7583
if ws.Equal(database.WorkspaceIdentity{}) {
76-
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
84+
workspace, err := a.Database.GetWorkspaceByAgentID(ctx, agentID)
7785
if err != nil {
7886
return nil, xerrors.Errorf("get workspace by agent id: %w", err)
7987
}
@@ -97,7 +105,7 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
97105
WorkspaceOwnerID: ws.OwnerID,
98106
WorkspaceID: ws.ID,
99107
WorkspaceName: ws.Name,
100-
AgentName: workspaceAgent.Name,
108+
AgentName: agentName,
101109
Type: connectionType,
102110
Code: code,
103111
Ip: logIP,

‎coderd/agentapi/connectionlog_test.go‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,16 @@ func TestConnectionLog(t *testing.T) {
3636
Name: "cool-workspace",
3737
}
3838
agent = database.WorkspaceAgent{
39-
ID: uuid.New(),
39+
ID: uuid.New(),
40+
Name: "test-agent",
4041
}
42+
agentAsCacheFields = agentapi.CachedAgentFields{}
43+
workspaceAsCacheFields = agentapi.CachedWorkspaceFields{}
4144
)
4245

46+
agentAsCacheFields.UpdateValues(agent.ID, agent.Name)
47+
workspaceAsCacheFields.UpdateValues(workspace)
48+
4349
tests := []struct {
4450
name string
4551
id uuid.UUID
@@ -109,15 +115,16 @@ func TestConnectionLog(t *testing.T) {
109115
connLogger := connectionlog.NewFake()
110116

111117
mDB := dbmock.NewMockStore(gomock.NewController(t))
112-
mDB.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Return(workspace, nil)
118+
// With agent/workspace caching, GetWorkspaceByAgentID is not called
113119

114120
api := &agentapi.ConnLogAPI{
115121
ConnectionLogger: asAtomicPointer[connectionlog.ConnectionLogger](connLogger),
116122
Database: mDB,
117123
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
118124
return agent, nil
119125
},
120-
Workspace: &agentapi.CachedWorkspaceFields{},
126+
Agent: &agentAsCacheFields,
127+
Workspace: &workspaceAsCacheFields,
121128
}
122129
api.ReportConnection(context.Background(), &agentproto.ReportConnectionRequest{
123130
Connection: &agentproto.Connection{

‎coderd/agentapi/metadata.go‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919

2020
type MetadataAPI struct {
2121
AgentFn func(context.Context) (database.WorkspaceAgent, error)
22+
Agent *CachedAgentFields
2223
Workspace *CachedWorkspaceFields
2324
Database database.Store
2425
Pubsub pubsub.Pubsub
@@ -60,16 +61,23 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
6061
}
6162
}
6263

63-
workspaceAgent, err := a.AgentFn(rbacCtx)
64-
if err != nil {
65-
return nil, err
64+
// Use cached agent ID if available to avoid database query.
65+
// The agent ID never changes during the connection lifetime, so it's safe to cache.
66+
agentID := a.Agent.ID()
67+
if agentID == uuid.Nil {
68+
// Fallback to querying the agent if cache is not populated (shouldn't happen in normal operation).
69+
workspaceAgent, err := a.AgentFn(rbacCtx)
70+
if err != nil {
71+
return nil, err
72+
}
73+
agentID = workspaceAgent.ID
6674
}
6775

6876
var (
6977
collectedAt = a.now()
7078
allKeysLen = 0
7179
dbUpdate = database.UpdateWorkspaceAgentMetadataParams{
72-
WorkspaceAgentID: workspaceAgent.ID,
80+
WorkspaceAgentID: agentID,
7381
// These need to be `make(x, 0, len(req.Metadata))` instead of
7482
// `make(x, len(req.Metadata))` because we may not insert all
7583
// metadata if the keys are large.
@@ -134,7 +142,7 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
134142
if err != nil {
135143
return nil, xerrors.Errorf("marshal workspace agent metadata channel payload: %w", err)
136144
}
137-
err = a.Pubsub.Publish(WatchWorkspaceAgentMetadataChannel(workspaceAgent.ID), payload)
145+
err = a.Pubsub.Publish(WatchWorkspaceAgentMetadataChannel(agentID), payload)
138146
if err != nil {
139147
return nil, xerrors.Errorf("publish workspace agent metadata: %w", err)
140148
}

0 commit comments

Comments
 (0)