coder · deansheather · Mar 30, 2023 · Mar 16, 2023 · Mar 17, 2023 · Mar 28, 2023
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -287,6 +287,7 @@ func New(options *Options) *API {
 			options.Database,
 			options.DeploymentValues,
 			oauthConfigs,
+			options.AgentInactiveDisconnectTimeout,
 			options.AppSigningKey,
 		),
 		metricsCache:          metricsCache,
@@ -612,14 +613,16 @@ func New(options *Options) *API {
 				r.Post("/report-stats", api.workspaceAgentReportStats)
 				r.Post("/report-lifecycle", api.workspaceAgentReportLifecycle)
 			})
+			// No middleware on the PTY endpoint since it uses workspace
+			// application auth and tickets.
+			r.Get("/{workspaceagent}/pty", api.workspaceAgentPTY)
 			r.Route("/{workspaceagent}", func(r chi.Router) {
 				r.Use(
 					apiKeyMiddleware,
 					httpmw.ExtractWorkspaceAgentParam(options.Database),
 					httpmw.ExtractWorkspaceParam(options.Database),
 				)
 				r.Get("/", api.workspaceAgent)
-				r.Get("/pty", api.workspaceAgentPTY)
 				r.Get("/startup-logs", api.workspaceAgentStartupLogs)
 				r.Get("/listening-ports", api.workspaceAgentListeningPorts)
 				r.Get("/connection", api.workspaceAgentConnection)

diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -638,10 +638,17 @@ func AwaitWorkspaceBuildJob(t *testing.T, client *codersdk.Client, build uuid.UU
 	return workspaceBuild
 }
 
-// AwaitWorkspaceAgents waits for all resources with agents to be connected.
-func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, workspaceID uuid.UUID) []codersdk.WorkspaceResource {
+// AwaitWorkspaceAgents waits for all resources with agents to be connected. If
+// specific agents are provided, it will wait for those agents to be connected
+// but will not fail if other agents are not connected.
+func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource {
 	t.Helper()
 
+	agentNamesMap := make(map[string]struct{}, len(agentNames))
+	for _, name := range agentNames {
+		agentNamesMap[name] = struct{}{}
+	}
+
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
@@ -659,6 +666,12 @@ func AwaitWorkspaceAgents(t *testing.T, client *codersdk.Client, workspaceID uui
 
 		for _, resource := range workspace.LatestBuild.Resources {
 			for _, agent := range resource.Agents {
+				if len(agentNames) > 0 {
+					if _, ok := agentNamesMap[agent.Name]; !ok {
+						continue
+					}
+				}
+
 				if agent.Status != codersdk.WorkspaceAgentConnected {
 					t.Logf("agent %s not connected yet", agent.Name)
 					return false

diff --git a/coderd/database/dbauthz/querier.go b/coderd/database/dbauthz/querier.go
@@ -1292,6 +1292,15 @@ func (q *querier) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanc
 	return agent, nil
 }
 
+func (q *querier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]database.WorkspaceAgent, error) {
+	workspace, err := q.GetWorkspaceByID(ctx, workspaceID)
+	if err != nil {
+		return nil, err
+	}
+
+	return q.db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, workspace.ID)
+}
+
 func (q *querier) UpdateWorkspaceAgentLifecycleStateByID(ctx context.Context, arg database.UpdateWorkspaceAgentLifecycleStateByIDParams) error {
 	agent, err := q.db.GetWorkspaceAgentByID(ctx, arg.ID)
 	if err != nil {

diff --git a/coderd/database/dbfake/databasefake.go b/coderd/database/dbfake/databasefake.go
@@ -314,6 +314,38 @@ func (*fakeQuerier) DeleteOldWorkspaceAgentStats(_ context.Context) error {
 	return nil
 }
 
+func (q *fakeQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]database.WorkspaceAgent, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	// Get latest build for workspace.
+	workspaceBuild, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, workspaceID)
+	if err != nil {
+		return nil, xerrors.Errorf("get latest workspace build: %w", err)
+	}
+
+	// Get resources for build.
+	resources, err := q.GetWorkspaceResourcesByJobID(ctx, workspaceBuild.JobID)
+	if err != nil {
+		return nil, xerrors.Errorf("get workspace resources: %w", err)
+	}
+	if len(resources) == 0 {
+		return []database.WorkspaceAgent{}, nil
+	}
+
+	resourceIDs := make([]uuid.UUID, len(resources))
+	for i, resource := range resources {
+		resourceIDs[i] = resource.ID
+	}
+
+	agents, err := q.GetWorkspaceAgentsByResourceIDs(ctx, resourceIDs)
+	if err != nil {
+		return nil, xerrors.Errorf("get workspace agents: %w", err)
+	}
+
+	return agents, nil
+}
+
 func (q *fakeQuerier) GetDeploymentWorkspaceAgentStats(_ context.Context, createdAfter time.Time) (database.GetDeploymentWorkspaceAgentStatsRow, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go
@@ -3,6 +3,7 @@ package database
 import (
 	"sort"
 	"strconv"
+	"time"
 
 	"golang.org/x/exp/maps"
 
@@ -36,6 +37,26 @@ func (s WorkspaceStatus) Valid() bool {
 	}
 }
 
+type WorkspaceAgentStatus string
+
+// This is also in codersdk/workspaceagents.go and should be kept in sync.
+const (
+	WorkspaceAgentStatusConnecting   WorkspaceAgentStatus = "connecting"
+	WorkspaceAgentStatusConnected    WorkspaceAgentStatus = "connected"
+	WorkspaceAgentStatusDisconnected WorkspaceAgentStatus = "disconnected"
+	WorkspaceAgentStatusTimeout      WorkspaceAgentStatus = "timeout"
+)
+
+func (s WorkspaceAgentStatus) Valid() bool {
+	switch s {
+	case WorkspaceAgentStatusConnecting, WorkspaceAgentStatusConnected,
+		WorkspaceAgentStatusDisconnected, WorkspaceAgentStatusTimeout:
+		return true
+	default:
+		return false
+	}
+}
+
 type AuditableGroup struct {
 	Group
 	Members []GroupMember `json:"members"`
@@ -199,6 +220,61 @@ func (l License) RBACObject() rbac.Object {
 	return rbac.ResourceLicense.WithIDString(strconv.FormatInt(int64(l.ID), 10))
 }
 
+type WorkspaceAgentConnectionStatus struct {
+	Status           WorkspaceAgentStatus `json:"status"`
+	FirstConnectedAt *time.Time           `json:"first_connected_at"`
+	LastConnectedAt  *time.Time           `json:"last_connected_at"`
+	DisconnectedAt   *time.Time           `json:"disconnected_at"`
+}
+
+func (a WorkspaceAgent) Status(inactiveTimeout time.Duration) WorkspaceAgentConnectionStatus {
+	connectionTimeout := time.Duration(a.ConnectionTimeoutSeconds) * time.Second
+
+	status := WorkspaceAgentConnectionStatus{
+		Status: WorkspaceAgentStatusDisconnected,
+	}
+	if a.FirstConnectedAt.Valid {
+		status.FirstConnectedAt = &a.FirstConnectedAt.Time
+	}
+	if a.LastConnectedAt.Valid {
+		status.LastConnectedAt = &a.LastConnectedAt.Time
+	}
+	if a.DisconnectedAt.Valid {
+		status.DisconnectedAt = &a.DisconnectedAt.Time
+	}
+
+	switch {
+	case !a.FirstConnectedAt.Valid:
+		switch {
+		case connectionTimeout > 0 && Now().Sub(a.CreatedAt) > connectionTimeout:
+			// If the agent took too long to connect the first time,
+			// mark it as timed out.
+			status.Status = WorkspaceAgentStatusTimeout
+		default:
+			// If the agent never connected, it's waiting for the compute
+			// to start up.
+			status.Status = WorkspaceAgentStatusConnecting
+		}
+	// We check before instead of after because last connected at and
+	// disconnected at can be equal timestamps in tight-timed tests.
+	case !a.DisconnectedAt.Time.Before(a.LastConnectedAt.Time):
+		// If we've disconnected after our last connection, we know the
+		// agent is no longer connected.
+		status.Status = WorkspaceAgentStatusDisconnected
+	case Now().Sub(a.LastConnectedAt.Time) > inactiveTimeout:
+		// The connection died without updating the last connected.
+		status.Status = WorkspaceAgentStatusDisconnected
+		// Client code needs an accurate disconnected at if the agent has been inactive.
+		status.DisconnectedAt = &a.LastConnectedAt.Time
+	case a.LastConnectedAt.Valid:
+		// The agent should be assumed connected if it's under inactivity timeouts
+		// and last connected at has been properly set.
+		status.Status = WorkspaceAgentStatusConnected
+	}
+
+	return status
+}
+
 func ConvertUserRows(rows []GetUsersRow) []User {
 	users := make([]User, len(rows))
 	for i, r := range rows {

diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql
@@ -132,3 +132,23 @@ INSERT INTO
 DELETE FROM workspace_agent_startup_logs WHERE agent_id IN
 	(SELECT id FROM workspace_agents WHERE last_connected_at IS NOT NULL
 		AND last_connected_at < NOW() - INTERVAL '7 day');
+
+-- name: GetWorkspaceAgentsInLatestBuildByWorkspaceID :many
+SELECT
+	workspace_agents.*
+FROM
+	workspace_agents
+JOIN
+	workspace_resources ON workspace_agents.resource_id = workspace_resources.id
+JOIN
+	workspace_builds ON workspace_resources.job_id = workspace_builds.job_id
+WHERE
+	workspace_builds.workspace_id = @workspace_id :: uuid AND
+	workspace_builds.build_number = (
+    	SELECT
+			MAX(build_number)
+    	FROM
+			workspace_builds AS wb
+    	WHERE
+			wb.workspace_id = @workspace_id :: uuid
+	);