Skip to content

Commit 25da224

Browse files
mtojekmtojek
authored
Filter query: has-agent connecting, connected, disconnected, timeout (coder#5145)
* WIP * has-agent:connecting, connected * Fix * Fix * has-agent:disconnected, timeout * Fix: typo * Fix * TODOs * databasefake * Fix: typo * More TODOs * databasefake * Timeout tests * Address PR comments * Implement FIXMEs * Renamings * Address PR comments * Fix: readability * Fix: refactor CASE logic * CASE logic * Fix * Use CTE * Polishing * Comment * WIP * IS NOT NULL * Without CTE * One more optimization * 2nd optimization
1 parent 511bb46 commit 25da224

7 files changed

Lines changed: 336 additions & 27 deletions

File tree

‎coderd/database/databasefake/databasefake.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -871,6 +871,44 @@ func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
871871
}
872872
}
873873

874+
if arg.HasAgent != "" {
875+
build, err := q.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID)
876+
if err != nil {
877+
return nil, xerrors.Errorf("get latest build: %w", err)
878+
}
879+
880+
job, err := q.GetProvisionerJobByID(ctx, build.JobID)
881+
if err != nil {
882+
return nil, xerrors.Errorf("get provisioner job: %w", err)
883+
}
884+
885+
workspaceResources, err := q.GetWorkspaceResourcesByJobID(ctx, job.ID)
886+
if err != nil {
887+
return nil, xerrors.Errorf("get workspace resources: %w", err)
888+
}
889+
890+
var workspaceResourceIDs []uuid.UUID
891+
for _, wr := range workspaceResources {
892+
workspaceResourceIDs = append(workspaceResourceIDs, wr.ID)
893+
}
894+
895+
workspaceAgents, err := q.GetWorkspaceAgentsByResourceIDs(ctx, workspaceResourceIDs)
896+
if err != nil {
897+
return nil, xerrors.Errorf("get workspace agents: %w", err)
898+
}
899+
900+
var hasAgentMatched bool
901+
for _, wa := range workspaceAgents {
902+
if mapAgentStatus(wa, arg.AgentInactiveDisconnectTimeoutSeconds) == arg.HasAgent {
903+
hasAgentMatched = true
904+
}
905+
}
906+
907+
if !hasAgentMatched {
908+
continue
909+
}
910+
}
911+
874912
if len(arg.TemplateIds) > 0 {
875913
match := false
876914
for _, id := range arg.TemplateIds {
@@ -909,6 +947,40 @@ func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
909947
return convertToWorkspaceRows(workspaces, int64(beforePageCount)), nil
910948
}
911949

950+
// mapAgentStatus determines the agent status based on different timestamps like created_at, last_connected_at, disconnected_at, etc.
951+
// The function must be in sync with: coderd/workspaceagents.go:convertWorkspaceAgent.
952+
func mapAgentStatus(dbAgent database.WorkspaceAgent, agentInactiveDisconnectTimeoutSeconds int64) string {
953+
var status string
954+
connectionTimeout := time.Duration(dbAgent.ConnectionTimeoutSeconds) * time.Second
955+
switch {
956+
case !dbAgent.FirstConnectedAt.Valid:
957+
switch {
958+
case connectionTimeout > 0 && database.Now().Sub(dbAgent.CreatedAt) > connectionTimeout:
959+
// If the agent took too long to connect the first time,
960+
// mark it as timed out.
961+
status = "timeout"
962+
default:
963+
// If the agent never connected, it's waiting for the compute
964+
// to start up.
965+
status = "connecting"
966+
}
967+
case dbAgent.DisconnectedAt.Time.After(dbAgent.LastConnectedAt.Time):
968+
// If we've disconnected after our last connection, we know the
969+
// agent is no longer connected.
970+
status = "disconnected"
971+
case database.Now().Sub(dbAgent.LastConnectedAt.Time) > time.Duration(agentInactiveDisconnectTimeoutSeconds)*time.Second:
972+
// The connection died without updating the last connected.
973+
status = "disconnected"
974+
case dbAgent.LastConnectedAt.Valid:
975+
// The agent should be assumed connected if it's under inactivity timeouts
976+
// and last connected at has been properly set.
977+
status = "connected"
978+
default:
979+
panic("unknown agent status: " + status)
980+
}
981+
return status
982+
}
983+
912984
func convertToWorkspaceRows(workspaces []database.Workspace, count int64) []database.GetWorkspacesRow {
913985
rows := make([]database.GetWorkspacesRow, len(workspaces))
914986
for i, w := range workspaces {

‎coderd/database/modelqueries.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,8 @@ func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspa
132132
arg.TemplateName,
133133
pq.Array(arg.TemplateIds),
134134
arg.Name,
135+
arg.HasAgent,
136+
arg.AgentInactiveDisconnectTimeoutSeconds,
135137
arg.Offset,
136138
arg.Limit,
137139
)

‎coderd/database/queries.sql.go‎

Lines changed: 58 additions & 16 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaces.sql‎

Lines changed: 45 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ FROM
4545
LEFT JOIN LATERAL (
4646
SELECT
4747
workspace_builds.transition,
48+
provisioner_jobs.id AS provisioner_job_id,
4849
provisioner_jobs.started_at,
4950
provisioner_jobs.updated_at,
5051
provisioner_jobs.canceled_at,
@@ -146,7 +147,7 @@ WHERE
146147
-- Use the organization filter to restrict to 1 org if needed.
147148
AND CASE
148149
WHEN @template_name :: text != '' THEN
149-
template_id = ANY(SELECT id FROM templates WHERE lower(name) = lower(@template_name) AND deleted = false)
150+
template_id = ANY(SELECT id FROM templates WHERE lower(name) = lower(@template_name) AND deleted = false)
150151
ELSE true
151152
END
152153
-- Filter by template_ids
@@ -161,17 +162,54 @@ WHERE
161162
name ILIKE '%' || @name || '%'
162163
ELSE true
163164
END
165+
-- Filter by agent status
166+
-- has-agent: is only applicable for workspaces in "start" transition. Stopped and deleted workspaces don't have agents.
167+
AND CASE
168+
WHEN @has_agent :: text != '' THEN
169+
(
170+
SELECT COUNT(*)
171+
FROM
172+
workspace_resources
173+
JOIN
174+
workspace_agents
175+
ON
176+
workspace_agents.resource_id = workspace_resources.id
177+
WHERE
178+
workspace_resources.job_id = latest_build.provisioner_job_id AND
179+
latest_build.transition = 'start'::workspace_transition AND
180+
@has_agent = (
181+
CASE
182+
WHEN workspace_agents.first_connected_at IS NULL THEN
183+
CASE
184+
WHEN workspace_agents.connection_timeout_seconds > 0 AND NOW() - workspace_agents.created_at > workspace_agents.connection_timeout_seconds * INTERVAL '1 second' THEN
185+
'timeout'
186+
ELSE
187+
'connecting'
188+
END
189+
WHEN workspace_agents.disconnected_at > workspace_agents.last_connected_at THEN
190+
'disconnected'
191+
WHEN NOW() - workspace_agents.last_connected_at > INTERVAL '1 second' * @agent_inactive_disconnect_timeout_seconds :: bigint THEN
192+
'disconnected'
193+
WHEN workspace_agents.last_connected_at IS NOT NULL THEN
194+
'connected'
195+
ELSE
196+
NULL
197+
END
198+
)
199+
) > 0
200+
ELSE true
201+
END
164202
-- Authorize Filter clause will be injected below in GetAuthorizedWorkspaces
165203
-- @authorize_filter
166204
ORDER BY
167-
last_used_at DESC
205+
last_used_at DESC
168206
LIMIT
169-
CASE
170-
WHEN @limit_ :: integer > 0 THEN
171-
@limit_
172-
END
207+
CASE
208+
WHEN @limit_ :: integer > 0 THEN
209+
@limit_
210+
END
173211
OFFSET
174-
@offset_
212+
@offset_
175213
;
176214

177215
-- name: GetWorkspaceByOwnerIDAndName :one

‎coderd/workspaces.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) {
104104
}
105105

106106
queryStr := r.URL.Query().Get("q")
107-
filter, errs := workspaceSearchQuery(queryStr, page)
107+
filter, errs := workspaceSearchQuery(queryStr, page, api.AgentInactiveDisconnectTimeout)
108108
if len(errs) > 0 {
109109
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
110110
Message: "Invalid workspace search query.",
@@ -1091,8 +1091,10 @@ func validWorkspaceSchedule(s *string) (sql.NullString, error) {
10911091

10921092
// workspaceSearchQuery takes a query string and returns the workspace filter.
10931093
// It also can return the list of validation errors to return to the api.
1094-
func workspaceSearchQuery(query string, page codersdk.Pagination) (database.GetWorkspacesParams, []codersdk.ValidationError) {
1094+
func workspaceSearchQuery(query string, page codersdk.Pagination, agentInactiveDisconnectTimeout time.Duration) (database.GetWorkspacesParams, []codersdk.ValidationError) {
10951095
filter := database.GetWorkspacesParams{
1096+
AgentInactiveDisconnectTimeoutSeconds: int64(agentInactiveDisconnectTimeout.Seconds()),
1097+
10961098
Offset: int32(page.Offset),
10971099
Limit: int32(page.Limit),
10981100
}
@@ -1139,7 +1141,7 @@ func workspaceSearchQuery(query string, page codersdk.Pagination) (database.GetW
11391141
filter.TemplateName = parser.String(searchParams, "", "template")
11401142
filter.Name = parser.String(searchParams, "", "name")
11411143
filter.Status = parser.String(searchParams, "", "status")
1142-
1144+
filter.HasAgent = parser.String(searchParams, "", "has-agent")
11431145
return filter, parser.Errors
11441146
}
11451147

‎coderd/workspaces_internal_test.go‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"fmt"
55
"strings"
66
"testing"
7+
"time"
78

89
"github.com/coder/coder/coderd/database"
910
"github.com/coder/coder/codersdk"
@@ -136,7 +137,7 @@ func TestSearchWorkspace(t *testing.T) {
136137
c := c
137138
t.Run(c.Name, func(t *testing.T) {
138139
t.Parallel()
139-
values, errs := workspaceSearchQuery(c.Query, codersdk.Pagination{})
140+
values, errs := workspaceSearchQuery(c.Query, codersdk.Pagination{}, 0)
140141
if c.ExpectedErrorContains != "" {
141142
require.True(t, len(errs) > 0, "expect some errors")
142143
var s strings.Builder
@@ -150,4 +151,13 @@ func TestSearchWorkspace(t *testing.T) {
150151
}
151152
})
152153
}
154+
t.Run("AgentInactiveDisconnectTimeout", func(t *testing.T) {
155+
t.Parallel()
156+
157+
query := `foo:bar`
158+
timeout := 1337 * time.Second
159+
values, errs := workspaceSearchQuery(query, codersdk.Pagination{}, timeout)
160+
require.Empty(t, errs)
161+
require.Equal(t, int64(timeout.Seconds()), values.AgentInactiveDisconnectTimeoutSeconds)
162+
})
153163
}

0 commit comments

Comments
 (0)