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

fix: limit calls to GetWorkspaceAgentByID in agentapi (#23015)

Browse Source 
We currently call GetWorkspaceAgentByID millions of times at scale
unnecessarily. This PR embeds immutable fields into the relevant
services instead of fetching for them every time.

resolves https://github.com/coder/scaletest/issues/84

Confirmed with a 10k scaletest that this changeset takes the query from
10M+ queries down to 39k
This commit is contained in:
Jon Ayers
2026-03-20 15:42:05 -05:00
committed by GitHub
parent 32021b3ac2
commit f135ffdb3a
23 changed files with 173 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/agentapi/apps_test.go
+28 -42
View File
@@ -67,12 +67,10 @@ func TestBatchUpdateAppHealths(t *testing.T) {
publishCalled := false
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: dbM,
Log: testutil.Logger(t),
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
PublishWorkspaceUpdateFn: func(ctx context.Context, _ uuid.UUID, kind wspubsub.WorkspaceEventKind) error {
publishCalled = true
return nil
},
@@ -105,12 +103,10 @@ func TestBatchUpdateAppHealths(t *testing.T) {
publishCalled := false
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: dbM,
Log: testutil.Logger(t),
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
PublishWorkspaceUpdateFn: func(ctx context.Context, _ uuid.UUID, kind wspubsub.WorkspaceEventKind) error {
publishCalled = true
return nil
},
@@ -144,12 +140,10 @@ func TestBatchUpdateAppHealths(t *testing.T) {
publishCalled := false
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: dbM,
Log: testutil.Logger(t),
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
PublishWorkspaceUpdateFn: func(ctx context.Context, _ uuid.UUID, kind wspubsub.WorkspaceEventKind) error {
publishCalled = true
return nil
},
@@ -180,9 +174,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app3}, nil)
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: dbM,
Log: testutil.Logger(t),
PublishWorkspaceUpdateFn: nil,
@@ -209,9 +201,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app1, app2}, nil)
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: dbM,
Log: testutil.Logger(t),
PublishWorkspaceUpdateFn: nil,
@@ -239,9 +229,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app1, app2}, nil)
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: dbM,
Log: testutil.Logger(t),
PublishWorkspaceUpdateFn: nil,
@@ -279,14 +267,26 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
}
workspaceUpdates := make(chan wspubsub.WorkspaceEventKind, 100)
workspace := database.Workspace{
ID: uuid.UUID{9},
TaskID: uuid.NullUUID{
Valid: true,
UUID: uuid.UUID{7},
},
}
cachedWs := &agentapi.CachedWorkspaceFields{}
cachedWs.UpdateValues(workspace)
api := &agentapi.AppsAPI{
AgentID: agent.ID,
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
Database: mDB,
Log: testutil.Logger(t),
PublishWorkspaceUpdateFn: func(_ context.Context, agnt *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
assert.Equal(t, *agnt, agent)
Database: mDB,
Log: testutil.Logger(t),
Workspace: cachedWs,
PublishWorkspaceUpdateFn: func(_ context.Context, agnt uuid.UUID, kind wspubsub.WorkspaceEventKind) error {
assert.Equal(t, agnt, agent.ID)
testutil.AssertSend(ctx, t, workspaceUpdates, kind)
return nil
},
@@ -309,14 +309,6 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
},
}
mDB.EXPECT().GetTaskByID(gomock.Any(), task.ID).Times(1).Return(task, nil)
workspace := database.Workspace{
ID: uuid.UUID{9},
TaskID: uuid.NullUUID{
Valid: true,
UUID: task.ID,
},
}
mDB.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agent.ID).Times(1).Return(workspace, nil)
appStatus := database.WorkspaceAppStatus{
ID: uuid.UUID{6},
}
@@ -363,9 +355,7 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
Return(database.WorkspaceApp{}, sql.ErrNoRows)
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: mDB,
Log: testutil.Logger(t),
}
@@ -392,9 +382,7 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
}
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: mDB,
Log: testutil.Logger(t),
}
@@ -422,9 +410,7 @@ func TestWorkspaceAgentAppStatus(t *testing.T) {
}
api := &agentapi.AppsAPI{
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
return agent, nil
},
AgentID: agent.ID,
Database: mDB,
Log: testutil.Logger(t),
}