perf: reduce number of queries made by /api/v2/workspaceagents/{id} (#21522)

Relates to https://github.com/coder/internal/issues/1214

The `ExtractWorkspaceAgentParam` middleware ends up making 4 database
queries to follow the chain of `WorkspaceAgent` -> `WorkspaceResource`
-> `ProvisionerJob` -> `WorkspaceBuild` -- but then dropping all that
hard work on the floor. The `api.workspaceAgent` handler that references
this middleware then has to do all of that work again, plus one more
query to get the related `User` so we can get the username. This pattern
is also mirrored in `getDatabaseTerminal` but without the middleware.

This PR:
* Adds a new query `GetWorkspaceAgentAndWorkspaceByID` to fetch all
this information at once to avoid the multiple round-trips,
* Updates the existing usage of `GetWorkspaceAgentByID` to this new
query instead,
* Updates `ExtractWorkspaceAgentParam` to also store the workspace in
the request context

Dalibo: [0.63ms](https://explain.dalibo.com/plan/40bb597f3539gc6c)
This commit is contained in:
Cian Johnston
2026-01-19 12:36:33 +00:00
committed by GitHub
parent d176714f90
commit 08343a7a9f
17 changed files with 259 additions and 230 deletions
+24 -42
View File
@@ -96,8 +96,6 @@ func TestWatchAgentContainers(t *testing.T) {
workspaceID = uuid.New()
agentID = uuid.New()
resourceID = uuid.New()
jobID = uuid.New()
buildID = uuid.New()
containersCh = make(chan codersdk.WorkspaceAgentListContainersResponse)
@@ -120,24 +118,17 @@ func TestWatchAgentContainers(t *testing.T) {
api.agentProvider = fAgentProvider
// Setup: Allow `ExtractWorkspaceAgentParams` to complete.
mDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(database.WorkspaceAgent{
ID: agentID,
ResourceID: resourceID,
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
}, nil)
mDB.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID).Return(database.WorkspaceResource{
ID: resourceID,
JobID: jobID,
}, nil)
mDB.EXPECT().GetProvisionerJobByID(gomock.Any(), jobID).Return(database.ProvisionerJob{
ID: jobID,
Type: database.ProvisionerJobTypeWorkspaceBuild,
}, nil)
mDB.EXPECT().GetWorkspaceBuildByJobID(gomock.Any(), jobID).Return(database.WorkspaceBuild{
WorkspaceID: workspaceID,
ID: buildID,
mDB.EXPECT().GetWorkspaceAgentAndWorkspaceByID(gomock.Any(), agentID).Return(database.GetWorkspaceAgentAndWorkspaceByIDRow{
WorkspaceAgent: database.WorkspaceAgent{
ID: agentID,
ResourceID: resourceID,
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
},
WorkspaceTable: database.WorkspaceTable{
ID: workspaceID,
},
}, nil)
// And: Allow `db2dsk.WorkspaceAgent` to complete.
@@ -152,7 +143,7 @@ func TestWatchAgentContainers(t *testing.T) {
})
// And: We mount the HTTP Handler
r.With(httpmw.ExtractWorkspaceAgentParam(mDB)).
r.With(httpmw.ExtractWorkspaceAgentAndWorkspaceParam(mDB)).
Get("/workspaceagents/{workspaceagent}/containers/watch", api.watchWorkspaceAgentContainers)
// Given: We create the HTTP server
@@ -222,8 +213,6 @@ func TestWatchAgentContainers(t *testing.T) {
workspaceID = uuid.New()
agentID = uuid.New()
resourceID = uuid.New()
jobID = uuid.New()
buildID = uuid.New()
containersCh = make(chan codersdk.WorkspaceAgentListContainersResponse)
@@ -246,24 +235,17 @@ func TestWatchAgentContainers(t *testing.T) {
api.agentProvider = fAgentProvider
// Setup: Allow `ExtractWorkspaceAgentParams` to complete.
mDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(database.WorkspaceAgent{
ID: agentID,
ResourceID: resourceID,
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
}, nil)
mDB.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID).Return(database.WorkspaceResource{
ID: resourceID,
JobID: jobID,
}, nil)
mDB.EXPECT().GetProvisionerJobByID(gomock.Any(), jobID).Return(database.ProvisionerJob{
ID: jobID,
Type: database.ProvisionerJobTypeWorkspaceBuild,
}, nil)
mDB.EXPECT().GetWorkspaceBuildByJobID(gomock.Any(), jobID).Return(database.WorkspaceBuild{
WorkspaceID: workspaceID,
ID: buildID,
mDB.EXPECT().GetWorkspaceAgentAndWorkspaceByID(gomock.Any(), agentID).Return(database.GetWorkspaceAgentAndWorkspaceByIDRow{
WorkspaceAgent: database.WorkspaceAgent{
ID: agentID,
ResourceID: resourceID,
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
},
WorkspaceTable: database.WorkspaceTable{
ID: workspaceID,
},
}, nil)
// And: Allow `db2dsk.WorkspaceAgent` to complete.
@@ -274,7 +256,7 @@ func TestWatchAgentContainers(t *testing.T) {
Return(containersCh, io.NopCloser(&bytes.Buffer{}), nil)
// And: We mount the HTTP Handler
r.With(httpmw.ExtractWorkspaceAgentParam(mDB)).
r.With(httpmw.ExtractWorkspaceAgentAndWorkspaceParam(mDB)).
Get("/workspaceagents/{workspaceagent}/containers/watch", api.watchWorkspaceAgentContainers)
// Given: We create the HTTP server