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