fix: Ensure WebSockets routinely transfer data

kylecarbs · kylecarbs · commit 44564c56ae75 · 2022-10-04T21:39:11.000Z
Fixes #4351.
diff --git a/coderd/httpapi/websocket.go b/coderd/httpapi/websocket.go
@@ -0,0 +1,27 @@
+package httpapi
+
+import (
+	"context"
+	"time"
+
+	"nhooyr.io/websocket"
+)
+
+// Heartbeat loops to ping a WebSocket to keep it alive.
+// Default idle connection timeouts are typically 60 seconds.
+// See: https://docs.aws.amazon.com/elasticloadbalancing/latest/application/application-load-balancers.html#connection-idle-timeout
+func Heartbeat(ctx context.Context, conn *websocket.Conn) {
+	ticker := time.NewTicker(15 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+		}
+		err := conn.Ping(ctx)
+		if err != nil {
+			return
+		}
+	}
+}
diff --git a/coderd/provisionerjobs.go b/coderd/provisionerjobs.go
@@ -151,6 +151,7 @@ func (api *API) provisionerJobLogs(rw http.ResponseWriter, r *http.Request, job
 		})
 		return
 	}
+	go httpapi.Heartbeat(ctx, conn)
 
 	ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageText)
 	defer wsNetConn.Close() // Also closes conn.
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -195,6 +195,7 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 		})
 		return
 	}
+	go httpapi.Heartbeat(ctx, conn)
 
 	_, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
 	defer wsNetConn.Close() // Also closes conn.
@@ -356,6 +357,8 @@ func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request
 		})
 		return
 	}
+	go httpapi.Heartbeat(ctx, conn)
+
 	ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
 	defer wsNetConn.Close()
 
@@ -477,6 +480,8 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
 		})
 		return
 	}
+	go httpapi.Heartbeat(ctx, conn)
+
 	defer conn.Close(websocket.StatusNormalClosure, "")
 	err = api.TailnetCoordinator.ServeClient(websocket.NetConn(ctx, conn, websocket.MessageBinary), uuid.New(), workspaceAgent.ID)
 	if err != nil {
@@ -582,6 +587,7 @@ func convertWorkspaceAgent(derpMap *tailcfg.DERPMap, coordinator *tailnet.Coordi
 
 	return workspaceAgent, nil
 }
+
 func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
@@ -628,6 +634,8 @@ func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Reques
 		})
 		return
 	}
+	go httpapi.Heartbeat(ctx, conn)
+
 	defer conn.Close(websocket.StatusGoingAway, "")
 
 	var lastReport codersdk.AgentStatsReportResponse

Original file line number	Diff line number	Diff line change
`@@ -151,6 +151,7 @@ func (api API) provisionerJobLogs(rw http.ResponseWriter, r http.Request, job`
`151`	`151`	`})`
`152`	`152`	`return`
`153`	`153`	`}`
	`154`	`+ go httpapi.Heartbeat(ctx, conn)`
`154`	`155`
`155`	`156`	`ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageText)`
`156`	`157`	`defer wsNetConn.Close() // Also closes conn.`