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: reduce agentfake manager startup time (#25669)

Browse Source 
Signed-off-by: Callum Styan <callumstyan@gmail.com>
Co-authored-by: Mux <noreply@coder.com>
This commit is contained in:
Callum Styan
2026-06-04 15:28:13 -07:00
committed by GitHub
co-authored by Mux
parent 6dedae4858
commit 4627b01415
14 changed files with 918 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+14
View File
@@ -3457,6 +3457,20 @@ func (q *querier) GetEnabledMCPServerConfigs(ctx context.Context) ([]database.MC
return q.db.GetEnabledMCPServerConfigs(ctx)
}
// GetExternalAgentTokensByTemplateID is used for scaletesting purposes; the
// scaletest agentfake path calls this query directly via a connection to the
// database. There is no production code path that uses this method, and it is
// deliberately not exposed over HTTP. The query filters for running
// workspaces only (latest build has transition=start and job_status=succeeded).
func (q *querier) GetExternalAgentTokensByTemplateID(ctx context.Context, arg database.GetExternalAgentTokensByTemplateIDParams) ([]database.GetExternalAgentTokensByTemplateIDRow, error) {
// ResourceSystem is used because the query spans multiple workspaces
// with no single RBAC object to check.
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
return nil, err
}
return q.db.GetExternalAgentTokensByTemplateID(ctx, arg)
}
func (q *querier) GetExternalAuthLink(ctx context.Context, arg database.GetExternalAuthLinkParams) (database.ExternalAuthLink, error) {
return fetchWithAction(q.log, q.auth, policy.ActionReadPersonal, q.db.GetExternalAuthLink)(ctx, arg)
}
coderd/database/dbauthz/dbauthz_test.go
+6
View File
@@ -3137,6 +3137,12 @@ func (s *MethodTestSuite) TestUser() {
dbm.EXPECT().UpdateGitSSHKey(gomock.Any(), arg).Return(key, nil).AnyTimes()
check.Args(arg).Asserts(key, policy.ActionUpdatePersonal).Returns(key)
}))
s.Run("GetExternalAgentTokensByTemplateID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
arg := database.GetExternalAgentTokensByTemplateIDParams{TemplateID: uuid.New(), OwnerID: uuid.Nil}
row := testutil.Fake(s.T(), faker, database.GetExternalAgentTokensByTemplateIDRow{})
dbm.EXPECT().GetExternalAgentTokensByTemplateID(gomock.Any(), arg).Return([]database.GetExternalAgentTokensByTemplateIDRow{row}, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceSystem, policy.ActionRead).Returns(slice.New(row))
}))
s.Run("GetExternalAuthLink", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
link := testutil.Fake(s.T(), faker, database.ExternalAuthLink{})
arg := database.GetExternalAuthLinkParams{ProviderID: link.ProviderID, UserID: link.UserID}
coderd/database/dbmetrics/querymetrics.go
Generated
+8
View File
@@ -1841,6 +1841,14 @@ func (m queryMetricsStore) GetEnabledMCPServerConfigs(ctx context.Context) ([]da
return r0, r1
}
func (m queryMetricsStore) GetExternalAgentTokensByTemplateID(ctx context.Context, arg database.GetExternalAgentTokensByTemplateIDParams) ([]database.GetExternalAgentTokensByTemplateIDRow, error) {
start := time.Now()
r0, r1 := m.s.GetExternalAgentTokensByTemplateID(ctx, arg)
m.queryLatencies.WithLabelValues("GetExternalAgentTokensByTemplateID").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetExternalAgentTokensByTemplateID").Inc()
return r0, r1
}
func (m queryMetricsStore) GetExternalAuthLink(ctx context.Context, arg database.GetExternalAuthLinkParams) (database.ExternalAuthLink, error) {
start := time.Now()
r0, r1 := m.s.GetExternalAuthLink(ctx, arg)
coderd/database/dbmock/dbmock.go
Generated
+15
View File
@@ -3420,6 +3420,21 @@ func (mr *MockStoreMockRecorder) GetEnabledMCPServerConfigs(ctx any) *gomock.Cal
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnabledMCPServerConfigs", reflect.TypeOf((*MockStore)(nil).GetEnabledMCPServerConfigs), ctx)
}
// GetExternalAgentTokensByTemplateID mocks base method.
func (m *MockStore) GetExternalAgentTokensByTemplateID(ctx context.Context, arg database.GetExternalAgentTokensByTemplateIDParams) ([]database.GetExternalAgentTokensByTemplateIDRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetExternalAgentTokensByTemplateID", ctx, arg)
ret0, _ := ret[0].([]database.GetExternalAgentTokensByTemplateIDRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetExternalAgentTokensByTemplateID indicates an expected call of GetExternalAgentTokensByTemplateID.
func (mr *MockStoreMockRecorder) GetExternalAgentTokensByTemplateID(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetExternalAgentTokensByTemplateID", reflect.TypeOf((*MockStore)(nil).GetExternalAgentTokensByTemplateID), ctx, arg)
}
// GetExternalAuthLink mocks base method.
func (m *MockStore) GetExternalAuthLink(ctx context.Context, arg database.GetExternalAuthLinkParams) (database.ExternalAuthLink, error) {
m.ctrl.T.Helper()
coderd/database/querier.go
Generated
+9
View File
@@ -459,6 +459,15 @@ type sqlcQuerier interface {
GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error)
GetEnabledChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error)
GetEnabledMCPServerConfigs(ctx context.Context) ([]MCPServerConfig, error)
// GetExternalAgentTokensByTemplateID returns the auth tokens for all
// non-deleted external agents on the latest build of every running workspace
// of the given template. "Running" means the latest build has
// transition=start and job_status=succeeded (matches the workspace-status
// definition used by coderd/database/queries/workspaces.sql).
// An owner_id of '00000000-0000-0000-0000-000000000000' (uuid.Nil) means
// "all owners"; any other value restricts results to workspaces owned by
// that user.
GetExternalAgentTokensByTemplateID(ctx context.Context, arg GetExternalAgentTokensByTemplateIDParams) ([]GetExternalAgentTokensByTemplateIDRow, error)
GetExternalAuthLink(ctx context.Context, arg GetExternalAuthLinkParams) (ExternalAuthLink, error)
GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]ExternalAuthLink, error)
GetFailedWorkspaceBuildsByTemplateID(ctx context.Context, arg GetFailedWorkspaceBuildsByTemplateIDParams) ([]GetFailedWorkspaceBuildsByTemplateIDRow, error)
coderd/database/queries.sql.go
Generated
+96
View File
@@ -30362,6 +30362,102 @@ func (q *sqlQuerier) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx conte
return i, err
}
const getExternalAgentTokensByTemplateID = `-- name: GetExternalAgentTokensByTemplateID :many
SELECT
workspaces.id AS workspace_id,
workspaces.name AS workspace_name,
workspace_agents.id AS agent_id,
workspace_agents.name AS agent_name,
workspace_agents.auth_token AS agent_token
FROM
workspaces
JOIN (
-- latest build per workspace
SELECT DISTINCT ON (workspace_id)
id, workspace_id, job_id, transition, has_external_agent
FROM
workspace_builds
ORDER BY
workspace_id, build_number DESC
) AS latest_builds
ON
latest_builds.workspace_id = workspaces.id
JOIN
provisioner_jobs
ON
provisioner_jobs.id = latest_builds.job_id
JOIN
workspace_resources
ON
workspace_resources.job_id = latest_builds.job_id
JOIN
workspace_agents
ON
workspace_agents.resource_id = workspace_resources.id
WHERE
workspaces.template_id = $1
AND (
$2 :: uuid = '00000000-0000-0000-0000-000000000000' :: uuid
OR workspaces.owner_id = $2
)
AND workspaces.deleted = FALSE
AND latest_builds.has_external_agent = TRUE
AND latest_builds.transition = 'start' :: workspace_transition
AND provisioner_jobs.job_status = 'succeeded' :: provisioner_job_status
AND workspace_agents.deleted = FALSE
AND workspace_agents.auth_instance_id IS NULL
`
type GetExternalAgentTokensByTemplateIDParams struct {
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
OwnerID uuid.UUID `db:"owner_id" json:"owner_id"`
}
type GetExternalAgentTokensByTemplateIDRow struct {
WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"`
WorkspaceName string `db:"workspace_name" json:"workspace_name"`
AgentID uuid.UUID `db:"agent_id" json:"agent_id"`
AgentName string `db:"agent_name" json:"agent_name"`
AgentToken uuid.UUID `db:"agent_token" json:"agent_token"`
}
// GetExternalAgentTokensByTemplateID returns the auth tokens for all
// non-deleted external agents on the latest build of every running workspace
// of the given template. "Running" means the latest build has
// transition=start and job_status=succeeded (matches the workspace-status
// definition used by coderd/database/queries/workspaces.sql).
// An owner_id of '00000000-0000-0000-0000-000000000000' (uuid.Nil) means
// "all owners"; any other value restricts results to workspaces owned by
// that user.
func (q *sqlQuerier) GetExternalAgentTokensByTemplateID(ctx context.Context, arg GetExternalAgentTokensByTemplateIDParams) ([]GetExternalAgentTokensByTemplateIDRow, error) {
rows, err := q.db.QueryContext(ctx, getExternalAgentTokensByTemplateID, arg.TemplateID, arg.OwnerID)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetExternalAgentTokensByTemplateIDRow
for rows.Next() {
var i GetExternalAgentTokensByTemplateIDRow
if err := rows.Scan(
&i.WorkspaceID,
&i.WorkspaceName,
&i.AgentID,
&i.AgentName,
&i.AgentToken,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const getWorkspaceAgentAndWorkspaceByID = `-- name: GetWorkspaceAgentAndWorkspaceByID :one
SELECT
workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted,
coderd/database/queries/workspaceagents.sql
+53
View File
@@ -352,6 +352,59 @@ WHERE
-- Filter out deleted sub agents.
AND workspace_agents.deleted = FALSE;
-- name: GetExternalAgentTokensByTemplateID :many
-- GetExternalAgentTokensByTemplateID returns the auth tokens for all
-- non-deleted external agents on the latest build of every running workspace
-- of the given template. "Running" means the latest build has
-- transition=start and job_status=succeeded (matches the workspace-status
-- definition used by coderd/database/queries/workspaces.sql).
-- An owner_id of '00000000-0000-0000-0000-000000000000' (uuid.Nil) means
-- "all owners"; any other value restricts results to workspaces owned by
-- that user.
SELECT
workspaces.id AS workspace_id,
workspaces.name AS workspace_name,
workspace_agents.id AS agent_id,
workspace_agents.name AS agent_name,
workspace_agents.auth_token AS agent_token
FROM
workspaces
JOIN (
-- latest build per workspace
SELECT DISTINCT ON (workspace_id)
id, workspace_id, job_id, transition, has_external_agent
FROM
workspace_builds
ORDER BY
workspace_id, build_number DESC
) AS latest_builds
ON
latest_builds.workspace_id = workspaces.id
JOIN
provisioner_jobs
ON
provisioner_jobs.id = latest_builds.job_id
JOIN
workspace_resources
ON
workspace_resources.job_id = latest_builds.job_id
JOIN
workspace_agents
ON
workspace_agents.resource_id = workspace_resources.id
WHERE
workspaces.template_id = @template_id
AND (
@owner_id :: uuid = '00000000-0000-0000-0000-000000000000' :: uuid
OR workspaces.owner_id = @owner_id
)
AND workspaces.deleted = FALSE
AND latest_builds.has_external_agent = TRUE
AND latest_builds.transition = 'start' :: workspace_transition
AND provisioner_jobs.job_status = 'succeeded' :: provisioner_job_status
AND workspace_agents.deleted = FALSE
AND workspace_agents.auth_instance_id IS NULL;
-- GetAuthenticatedWorkspaceAgentAndBuildByAuthToken returns an authenticated
-- workspace agent and its associated build. During normal operation, this is
-- the latest build. During shutdown, this may be the previous START build while