coder · mtojek · Nov 24, 2022 · Nov 21, 2022 · Nov 22, 2022 · Nov 22, 2022
diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -871,6 +871,44 @@ func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
 			}
 		}
 
+		if arg.HasAgent != "" {
+			build, err := q.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID)
+			if err != nil {
+				return nil, xerrors.Errorf("get latest build: %w", err)
+			}
+
+			job, err := q.GetProvisionerJobByID(ctx, build.JobID)
+			if err != nil {
+				return nil, xerrors.Errorf("get provisioner job: %w", err)
+			}
+
+			workspaceResources, err := q.GetWorkspaceResourcesByJobID(ctx, job.ID)
+			if err != nil {
+				return nil, xerrors.Errorf("get workspace resources: %w", err)
+			}
+
+			var workspaceResourceIDs []uuid.UUID
+			for _, wr := range workspaceResources {
+				workspaceResourceIDs = append(workspaceResourceIDs, wr.ID)
+			}
+
+			workspaceAgents, err := q.GetWorkspaceAgentsByResourceIDs(ctx, workspaceResourceIDs)
+			if err != nil {
+				return nil, xerrors.Errorf("get workspace agents: %w", err)
+			}
+
+			var hasAgentMatched bool
+			for _, wa := range workspaceAgents {
+				if mapAgentStatus(wa, arg.AgentInactiveDisconnectTimeoutSeconds) == arg.HasAgent {
+					hasAgentMatched = true
+				}
+			}
+
+			if !hasAgentMatched {
+				continue
+			}
+		}
+
 		if len(arg.TemplateIds) > 0 {
 			match := false
 			for _, id := range arg.TemplateIds {
@@ -909,6 +947,38 @@ func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
 	return convertToWorkspaceRows(workspaces, int64(beforePageCount)), nil
 }
 
+func mapAgentStatus(dbAgent database.WorkspaceAgent, agentInactiveDisconnectTimeoutSeconds int64) string {
+	var status string
+	connectionTimeout := time.Duration(dbAgent.ConnectionTimeoutSeconds) * time.Second
+	switch {
+	case !dbAgent.FirstConnectedAt.Valid:
+		switch {
+		case connectionTimeout > 0 && database.Now().Sub(dbAgent.CreatedAt) > connectionTimeout:
+			// If the agent took too long to connect the first time,
+			// mark it as timed out.
+			status = "timeout"
+		default:
+			// If the agent never connected, it's waiting for the compute
+			// to start up.
+			status = "connecting"
+		}
+	case dbAgent.DisconnectedAt.Time.After(dbAgent.LastConnectedAt.Time):
+		// If we've disconnected after our last connection, we know the
+		// agent is no longer connected.
+		status = "disconnected"
+	case database.Now().Sub(dbAgent.LastConnectedAt.Time) > time.Duration(agentInactiveDisconnectTimeoutSeconds)*time.Second:
+		// The connection died without updating the last connected.
+		status = "disconnected"
+	case dbAgent.LastConnectedAt.Valid:
+		// The agent should be assumed connected if it's under inactivity timeouts
+		// and last connected at has been properly set.
+		status = "connected"
+	default:
+		panic("unknown agent status: " + status)
+	}
+	return status
+}
+
 func convertToWorkspaceRows(workspaces []database.Workspace, count int64) []database.GetWorkspacesRow {
 	rows := make([]database.GetWorkspacesRow, len(workspaces))
 	for i, w := range workspaces {

diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go
@@ -132,8 +132,10 @@ func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspa
 		arg.TemplateName,
 		pq.Array(arg.TemplateIds),
 		arg.Name,
+		arg.HasAgent,
 		arg.Offset,
 		arg.Limit,
+		arg.AgentInactiveDisconnectTimeoutSeconds,
 	)
 	if err != nil {
 		return nil, xerrors.Errorf("get authorized workspaces: %w", err)

diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -38,12 +38,55 @@ WHERE
 	);
 
 -- name: GetWorkspaces :many
+WITH workspace_builds_agents AS (
+	SELECT
+		workspace_builds.workspace_id AS workspace_id,
+		workspace_builds.build_number AS build_number,
+		workspace_agents.id AS agent_id,
+		(
+			CASE
+				WHEN workspace_agents.first_connected_at IS NULL THEN
+					CASE
+						WHEN workspace_agents.connection_timeout_seconds > 0 AND NOW() - workspace_agents.created_at > workspace_agents.connection_timeout_seconds * INTERVAL '1 second' THEN
+							'timeout'
+						ELSE
+							'connecting'
+					END
+				WHEN workspace_agents.disconnected_at > workspace_agents.last_connected_at THEN
+					'disconnected'
+				WHEN NOW() - workspace_agents.last_connected_at > INTERVAL '1 second' * @agent_inactive_disconnect_timeout_seconds :: bigint THEN
+					'disconnected'
+				WHEN workspace_agents.last_connected_at IS NOT NULL THEN
+					'connected'
+				ELSE
+					NULL
+			END
+		) AS agent_status
+	FROM
+		workspace_builds
+	LEFT JOIN
+		provisioner_jobs
+	ON
+		provisioner_jobs.id = workspace_builds.job_id
+	LEFT JOIN
+		workspace_resources
+	ON
+		workspace_resources.job_id = provisioner_jobs.id
+	LEFT JOIN
+		workspace_agents
+	ON
+		workspace_agents.resource_id = workspace_resources.id
+	WHERE
+		workspace_builds.transition = 'start'::workspace_transition AND
+		workspace_agents.id IS NOT NULL
+)
 SELECT
 	workspaces.*, COUNT(*) OVER () as count
 FROM
 	workspaces
 LEFT JOIN LATERAL (
 	SELECT
+		workspace_builds.build_number,
 		workspace_builds.transition,
 		provisioner_jobs.started_at,
 		provisioner_jobs.updated_at,
@@ -146,7 +189,7 @@ WHERE
 	-- Use the organization filter to restrict to 1 org if needed.
 	AND CASE
 		WHEN @template_name :: text != '' THEN
-			template_id = ANY(SELECT id FROM templates WHERE lower(name) = lower(@template_name)  AND deleted = false)
+			template_id = ANY(SELECT id FROM templates WHERE lower(name) = lower(@template_name) AND deleted = false)
 		ELSE true
 	END
 	-- Filter by template_ids
@@ -161,17 +204,30 @@ WHERE
 			name ILIKE '%' || @name || '%'
 		ELSE true
 	END
+	-- Filter by agent status
+	-- has-agent: is only applicable for workspaces in "start" transition. Stopped and deleted workspaces don't have agents.
+	AND CASE
+		WHEN @has_agent :: text != '' THEN
+			(
+				SELECT COUNT(*) FROM workspace_builds_agents
+				WHERE
+					workspace_builds_agents.workspace_id = workspaces.id AND
+					workspace_builds_agents.build_number = latest_build.build_number AND
+					agent_status = @has_agent
+			) > 0
+		ELSE true
+	END
 	-- Authorize Filter clause will be injected below in GetAuthorizedWorkspaces
 	-- @authorize_filter
 ORDER BY
-    last_used_at DESC
+	last_used_at DESC
 LIMIT
-    CASE
-        WHEN @limit_ :: integer > 0 THEN
-            @limit_
-    END
+	CASE
+		WHEN @limit_ :: integer > 0 THEN
+			@limit_
+	END
 OFFSET
-    @offset_
+	@offset_
 ;
 
 -- name: GetWorkspaceByOwnerIDAndName :one

diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -104,7 +104,7 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	queryStr := r.URL.Query().Get("q")
-	filter, errs := workspaceSearchQuery(queryStr, page)
+	filter, errs := workspaceSearchQuery(queryStr, page, api.AgentInactiveDisconnectTimeout)
 	if len(errs) > 0 {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
 			Message:     "Invalid workspace search query.",
@@ -1098,8 +1098,10 @@ func validWorkspaceSchedule(s *string) (sql.NullString, error) {
 
 // workspaceSearchQuery takes a query string and returns the workspace filter.
 // It also can return the list of validation errors to return to the api.
-func workspaceSearchQuery(query string, page codersdk.Pagination) (database.GetWorkspacesParams, []codersdk.ValidationError) {
+func workspaceSearchQuery(query string, page codersdk.Pagination, agentInactiveDisconnectTimeout time.Duration) (database.GetWorkspacesParams, []codersdk.ValidationError) {
 	filter := database.GetWorkspacesParams{
+		AgentInactiveDisconnectTimeoutSeconds: int64(agentInactiveDisconnectTimeout.Seconds()),
+
 		Offset: int32(page.Offset),
 		Limit:  int32(page.Limit),
 	}
@@ -1146,7 +1148,7 @@ func workspaceSearchQuery(query string, page codersdk.Pagination) (database.GetW
 	filter.TemplateName = parser.String(searchParams, "", "template")
 	filter.Name = parser.String(searchParams, "", "name")
 	filter.Status = parser.String(searchParams, "", "status")
-
+	filter.HasAgent = parser.String(searchParams, "", "has-agent")
 	return filter, parser.Errors
 }
 

diff --git a/coderd/workspaces_internal_test.go b/coderd/workspaces_internal_test.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/codersdk"
@@ -136,7 +137,7 @@ func TestSearchWorkspace(t *testing.T) {
 		c := c
 		t.Run(c.Name, func(t *testing.T) {
 			t.Parallel()
-			values, errs := workspaceSearchQuery(c.Query, codersdk.Pagination{})
+			values, errs := workspaceSearchQuery(c.Query, codersdk.Pagination{}, 0)
 			if c.ExpectedErrorContains != "" {
 				require.True(t, len(errs) > 0, "expect some errors")
 				var s strings.Builder
@@ -150,4 +151,13 @@ func TestSearchWorkspace(t *testing.T) {
 			}
 		})
 	}
+	t.Run("AgentInactiveDisconnectTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		query := `foo:bar`
+		timeout := 1337 * time.Second
+		values, errs := workspaceSearchQuery(query, codersdk.Pagination{}, timeout)
+		require.Empty(t, errs)
+		require.Equal(t, int64(timeout.Seconds()), values.AgentInactiveDisconnectTimeoutSeconds)
+	})
 }