Address PR comments

mtojek · mtojek · commit 5ca8c1750255 · 2022-11-23T13:05:26.000+01:00
diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -899,18 +899,8 @@ func (q *fakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
 
 			var hasAgentMatched bool
 			for _, wa := range workspaceAgents {
-				switch arg.HasAgent {
-				case "connected":
-					hasAgentMatched = wa.LastConnectedAt.Valid
-				case "connecting":
-					hasAgentMatched = !wa.FirstConnectedAt.Valid
-				case "disconnected":
-					hasAgentMatched = (wa.DisconnectedAt.Valid && wa.DisconnectedAt.Time.After(wa.LastConnectedAt.Time)) ||
-						(wa.LastConnectedAt.Valid && wa.LastConnectedAt.Time.Add(time.Duration(arg.AgentInactiveDisconnectTimeoutSeconds)*time.Second).Before(database.Now()))
-				case "timeout":
-					hasAgentMatched = !wa.FirstConnectedAt.Valid &&
-						wa.CreatedAt.Add(time.Duration(wa.ConnectionTimeoutSeconds)*time.Second).Before(database.Now())
-				}
+				mapped := mapAgentStatus(wa, arg.AgentInactiveDisconnectTimeoutSeconds)
+				hasAgentMatched = mapped == arg.HasAgent
 				break // only 1 agent is expected
 			}
 
@@ -957,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/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -187,7 +187,7 @@ WHERE
 				WHEN @has_agent = 'timeout' THEN
 					latest_build.agent_first_connected_at IS NULL AND (latest_build.agent_created_at + latest_build.agent_connection_timeout_seconds * INTERVAL '1 second' < NOW())
 				WHEN @has_agent = 'connecting' THEN
-					latest_build.agent_first_connected_at IS NULL
+					latest_build.agent_first_connected_at IS NULL AND (latest_build.agent_created_at + latest_build.agent_connection_timeout_seconds * INTERVAL '1 second' >= NOW())
 				WHEN @has_agent = 'disconnected' THEN
 					(
 						latest_build.agent_disconnected_at IS NOT NULL AND
@@ -197,7 +197,13 @@ WHERE
 						latest_build.agent_last_connected_at + INTERVAL '1 second' * @agent_inactive_disconnect_timeout_seconds :: bigint < NOW()
 					)
 				WHEN @has_agent = 'connected' THEN
-					latest_build.agent_last_connected_at IS NOT NULL
+					(
+						latest_build.agent_disconnected_at IS NOT NULL AND
+						latest_build.agent_disconnected_at <= latest_build.agent_last_connected_at
+					) OR (
+						latest_build.agent_last_connected_at IS NOT NULL AND
+						latest_build.agent_last_connected_at + INTERVAL '1 second' * @agent_inactive_disconnect_timeout_seconds :: bigint >= NOW()
+					)
 				ELSE true
 			END
 		ELSE true
