fix: close server pty connections on client disconnect (#15201)

f0ssel · web-flow · commit 81e99bec6b92 · 2024-10-24T15:12:35.000-04:00
diff --git a/coderd/httpapi/websocket.go b/coderd/httpapi/websocket.go
@@ -2,8 +2,10 @@ package httpapi
 
 import (
 	"context"
+	"errors"
 	"time"
 
+	"golang.org/x/xerrors"
 	"nhooyr.io/websocket"
 
 	"cdr.dev/slog"
@@ -31,7 +33,8 @@ func Heartbeat(ctx context.Context, conn *websocket.Conn) {
 // Heartbeat loops to ping a WebSocket to keep it alive. It calls `exit` on ping
 // failure.
 func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn *websocket.Conn) {
-	ticker := time.NewTicker(15 * time.Second)
+	interval := 15 * time.Second
+	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
 
 	for {
@@ -40,12 +43,26 @@ func HeartbeatClose(ctx context.Context, logger slog.Logger, exit func(), conn *
 			return
 		case <-ticker.C:
 		}
-		err := conn.Ping(ctx)
+		err := pingWithTimeout(ctx, conn, interval)
 		if err != nil {
+			// context.DeadlineExceeded is expected when the client disconnects without sending a close frame
+			if !errors.Is(err, context.DeadlineExceeded) {
+				logger.Error(ctx, "failed to heartbeat ping", slog.Error(err))
+			}
 			_ = conn.Close(websocket.StatusGoingAway, "Ping failed")
-			logger.Info(ctx, "failed to heartbeat ping", slog.Error(err))
 			exit()
 			return
 		}
 	}
 }
+
+func pingWithTimeout(ctx context.Context, conn *websocket.Conn, timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+	err := conn.Ping(ctx)
+	if err != nil {
+		return xerrors.Errorf("failed to ping: %w", err)
+	}
+
+	return nil
+}
diff --git a/coderd/workspaceapps/proxy.go b/coderd/workspaceapps/proxy.go
@@ -593,7 +593,6 @@ func (s *Server) proxyWorkspaceApp(rw http.ResponseWriter, r *http.Request, appT
 	tracing.EndHTTPSpan(r, http.StatusOK, trace.SpanFromContext(ctx))
 
 	report := newStatsReportFromSignedToken(appToken)
-	s.collectStats(report)
 	defer func() {
 		// We must use defer here because ServeHTTP may panic.
 		report.SessionEndedAt = dbtime.Now()
@@ -614,7 +613,8 @@ func (s *Server) proxyWorkspaceApp(rw http.ResponseWriter, r *http.Request, appT
 // @Success 101
 // @Router /workspaceagents/{workspaceagent}/pty [get]
 func (s *Server) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
-	ctx := r.Context()
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
 
 	s.websocketWaitMutex.Lock()
 	s.websocketWaitGroup.Add(1)
@@ -670,12 +670,11 @@ func (s *Server) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 		})
 		return
 	}
+	go httpapi.HeartbeatClose(ctx, s.Logger, cancel, conn)
 
 	ctx, wsNetConn := WebsocketNetConn(ctx, conn, websocket.MessageBinary)
 	defer wsNetConn.Close() // Also closes conn.
 
-	go httpapi.Heartbeat(ctx, conn)
-
 	agentConn, release, err := s.AgentProvider.AgentConn(ctx, appToken.AgentID)
 	if err != nil {
 		log.Debug(ctx, "dial workspace agent", slog.Error(err))
