From e12b621ff0ce4d4da052e87c7d83d4be88df4a97 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Fri, 5 Sep 2025 14:26:46 +0100 Subject: [PATCH] fix(coderd): ensure agent WebSocket conn is cleaned up (#19711) When clients disconnected from the /containers/watch endpoint, the WebSocket connection between coderd and the agent stayed open. This caused heartbeat traffic every 15s that was incorrectly counted as workspace activity, extending workspace lifetimes indefinitely. Now properly cancels the agent connection context when the client disconnects. --- coderd/workspaceagents.go | 13 ++- coderd/workspaceagents_internal_test.go | 137 +++++++++++++++++++++++- 2 files changed, 144 insertions(+), 6 deletions(-) diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index f2ee1ac18e..ddab39ed8a 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -817,12 +817,13 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re var ( ctx = r.Context() workspaceAgent = httpmw.WorkspaceAgentParam(r) + logger = api.Logger.Named("agent_container_watcher").With(slog.F("agent_id", workspaceAgent.ID)) ) // If the agent is unreachable, the request will hang. Assume that if we // don't get a response after 30s that the agent is unreachable. - dialCtx, cancel := context.WithTimeout(ctx, 30*time.Second) - defer cancel() + dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second) + defer dialCancel() apiAgent, err := db2sdk.WorkspaceAgent( api.DERPMap(), *api.TailnetCoordinator.Load(), @@ -857,8 +858,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re } defer release() - watcherLogger := api.Logger.Named("agent_container_watcher").With(slog.F("agent_id", workspaceAgent.ID)) - containersCh, closer, err := agentConn.WatchContainers(ctx, watcherLogger) + containersCh, closer, err := agentConn.WatchContainers(ctx, logger) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error watching agent's containers.", @@ -877,6 +877,9 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re return } + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + // Here we close the websocket for reading, so that the websocket library will handle pings and // close frames. _ = conn.CloseRead(context.Background()) @@ -884,7 +887,7 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re ctx, wsNetConn := codersdk.WebsocketNetConn(ctx, conn, websocket.MessageText) defer wsNetConn.Close() - go httpapi.Heartbeat(ctx, conn) + go httpapi.HeartbeatClose(ctx, logger, cancel, conn) encoder := json.NewEncoder(wsNetConn) diff --git a/coderd/workspaceagents_internal_test.go b/coderd/workspaceagents_internal_test.go index c7520f05ab..90f5d2ab70 100644 --- a/coderd/workspaceagents_internal_test.go +++ b/coderd/workspaceagents_internal_test.go @@ -59,10 +59,145 @@ func (fakeAgentProvider) Close() error { return nil } +type channelCloser struct { + closeFn func() +} + +func (c *channelCloser) Close() error { + c.closeFn() + return nil +} + func TestWatchAgentContainers(t *testing.T) { t.Parallel() - t.Run("WebSocketClosesProperly", func(t *testing.T) { + t.Run("CoderdWebSocketCanHandleClientClosing", func(t *testing.T) { + t.Parallel() + + // This test ensures that the agent containers `/watch` websocket can gracefully + // handle the client websocket closing. This test was created in + // response to this issue: https://github.com/coder/coder/issues/19449 + + var ( + ctx = testutil.Context(t, testutil.WaitLong) + logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug).Named("coderd") + + mCtrl = gomock.NewController(t) + mDB = dbmock.NewMockStore(mCtrl) + mCoordinator = tailnettest.NewMockCoordinator(mCtrl) + mAgentConn = agentconnmock.NewMockAgentConn(mCtrl) + + fAgentProvider = fakeAgentProvider{ + agentConn: func(ctx context.Context, agentID uuid.UUID) (_ workspacesdk.AgentConn, release func(), _ error) { + return mAgentConn, func() {}, nil + }, + } + + workspaceID = uuid.New() + agentID = uuid.New() + resourceID = uuid.New() + jobID = uuid.New() + buildID = uuid.New() + + containersCh = make(chan codersdk.WorkspaceAgentListContainersResponse) + + r = chi.NewMux() + + api = API{ + ctx: ctx, + Options: &Options{ + AgentInactiveDisconnectTimeout: testutil.WaitShort, + Database: mDB, + Logger: logger, + DeploymentValues: &codersdk.DeploymentValues{}, + TailnetCoordinator: tailnettest.NewFakeCoordinator(), + }, + } + ) + + var tailnetCoordinator tailnet.Coordinator = mCoordinator + api.TailnetCoordinator.Store(&tailnetCoordinator) + api.agentProvider = fAgentProvider + + // Setup: Allow `ExtractWorkspaceAgentParams` to complete. + mDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(database.WorkspaceAgent{ + ID: agentID, + ResourceID: resourceID, + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + FirstConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, + LastConnectedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, + }, nil) + mDB.EXPECT().GetWorkspaceResourceByID(gomock.Any(), resourceID).Return(database.WorkspaceResource{ + ID: resourceID, + JobID: jobID, + }, nil) + mDB.EXPECT().GetProvisionerJobByID(gomock.Any(), jobID).Return(database.ProvisionerJob{ + ID: jobID, + Type: database.ProvisionerJobTypeWorkspaceBuild, + }, nil) + mDB.EXPECT().GetWorkspaceBuildByJobID(gomock.Any(), jobID).Return(database.WorkspaceBuild{ + WorkspaceID: workspaceID, + ID: buildID, + }, nil) + + // And: Allow `db2dsk.WorkspaceAgent` to complete. + mCoordinator.EXPECT().Node(gomock.Any()).Return(nil) + + // And: Allow `WatchContainers` to be called, returing our `containersCh` channel. + mAgentConn.EXPECT().WatchContainers(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ slog.Logger) (<-chan codersdk.WorkspaceAgentListContainersResponse, io.Closer, error) { + return containersCh, &channelCloser{closeFn: func() { + close(containersCh) + }}, nil + }) + + // And: We mount the HTTP Handler + r.With(httpmw.ExtractWorkspaceAgentParam(mDB)). + Get("/workspaceagents/{workspaceagent}/containers/watch", api.watchWorkspaceAgentContainers) + + // Given: We create the HTTP server + srv := httptest.NewServer(r) + defer srv.Close() + + // And: Dial the WebSocket + wsURL := strings.Replace(srv.URL, "http://", "ws://", 1) + conn, resp, err := websocket.Dial(ctx, fmt.Sprintf("%s/workspaceagents/%s/containers/watch", wsURL, agentID), nil) + require.NoError(t, err) + if resp.Body != nil { + defer resp.Body.Close() + } + + // And: Create a streaming decoder + decoder := wsjson.NewDecoder[codersdk.WorkspaceAgentListContainersResponse](conn, websocket.MessageText, logger) + defer decoder.Close() + decodeCh := decoder.Chan() + + // And: We can successfully send through the channel. + testutil.RequireSend(ctx, t, containersCh, codersdk.WorkspaceAgentListContainersResponse{ + Containers: []codersdk.WorkspaceAgentContainer{{ + ID: "test-container-id", + }}, + }) + + // And: Receive the data. + containerResp := testutil.RequireReceive(ctx, t, decodeCh) + require.Len(t, containerResp.Containers, 1) + require.Equal(t, "test-container-id", containerResp.Containers[0].ID) + + // When: We close the WebSocket + conn.Close(websocket.StatusNormalClosure, "test closing connection") + + // Then: We expect `containersCh` to be closed. + select { + case <-ctx.Done(): + t.Fail() + + case _, ok := <-containersCh: + require.False(t, ok, "channel is expected to be closed") + } + }) + + t.Run("CoderdWebSocketCanHandleAgentClosing", func(t *testing.T) { t.Parallel() // This test ensures that the agent containers `/watch` websocket can gracefully