chore: Refactor accepting websocket connections to track for close

coder · Emyrk · Apr 4, 2023 · Apr 4, 2023 · Apr 4, 2023 · Apr 4, 2023
commit 884c71b0c6e7f02a06d8dd77d9755872dfc36b03
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -784,8 +784,8 @@ type API struct {
 
 	siteHandler http.Handler
 
-	WebsocketWaitMutex sync.Mutex
 	WebsocketWaitGroup sync.WaitGroup
+	WebsocketWatch     *ActiveWebsockets
 	derpCloseFunc      func()
 
 	metricsCache          *metricscache.Cache
@@ -805,9 +805,7 @@ func (api *API) Close() error {
 	api.cancel()
 	api.derpCloseFunc()
 
-	api.WebsocketWaitMutex.Lock()
 	api.WebsocketWaitGroup.Wait()
-	api.WebsocketWaitMutex.Unlock()
 
 	api.metricsCache.Close()
 	if api.updateChecker != nil {

diff --git a/coderd/provisionerjobs.go b/coderd/provisionerjobs.go
@@ -113,9 +113,7 @@ func (api *API) provisionerJobLogs(rw http.ResponseWriter, r *http.Request, job
 		logs = []database.ProvisionerJobLog{}
 	}
 
-	api.WebsocketWaitMutex.Lock()
 	api.WebsocketWaitGroup.Add(1)
-	api.WebsocketWaitMutex.Unlock()
 	defer api.WebsocketWaitGroup.Done()
 	conn, err := websocket.Accept(rw, r, nil)
 	if err != nil {

diff --git a/coderd/sockets.go b/coderd/sockets.go
@@ -0,0 +1,81 @@
+package coderd
+
+import (
+	"context"
+	"net/http"
+	"sync"
+
+	"nhooyr.io/websocket"
+
+	"github.com/coder/coder/coderd/httpapi"
+	"github.com/coder/coder/codersdk"
+)
+
+// ActiveWebsockets is a helper struct that can be used to track active
+// websocket connections.
+type ActiveWebsockets struct {
+	ctx    context.Context
+	cancel func()
+
+	wg sync.WaitGroup
+}
+
+func NewActiveWebsockets(ctx context.Context, cancel func()) *ActiveWebsockets {
+	return &ActiveWebsockets{
+		ctx:    ctx,
+		cancel: cancel,
+	}
+}
+
+// Accept accepts a websocket connection and calls f with the connection.
+// The function will be tracked by the ActiveWebsockets struct and will be
+// closed when the parent context is canceled.
+func (a *ActiveWebsockets) Accept(rw http.ResponseWriter, r *http.Request, options *websocket.AcceptOptions, f func(conn *websocket.Conn)) {
+	// Ensure we are still accepting websocket connections, and not shutting down.
+	if err := a.ctx.Err(); err != nil {
+		httpapi.Write(context.Background(), rw, http.StatusBadRequest, codersdk.Response{
+			Message: "No longer accepting websocket requests.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	// Ensure we decrement the wait group when we are done.
+	a.wg.Add(1)
+	defer a.wg.Done()
+
+	// Accept the websocket connection
+	conn, err := websocket.Accept(rw, r, options)
+	if err != nil {
+		httpapi.Write(context.Background(), rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Failed to accept websocket.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	// Always track the connection before allowing the caller to handle it.
+	// This ensures the connection is closed when the parent context is canceled.
+	// This new context will end if the parent context is cancelled or if
+	// the connection is closed.
+	ctx, cancel := context.WithCancel(a.ctx)
+	defer cancel()
+	a.track(ctx, conn)
+
+	// Handle the websocket connection
+	f(conn)
+}
+
+// Track runs a go routine that will close a given websocket connection when
+// the parent context is canceled.
+func (a *ActiveWebsockets) track(ctx context.Context, conn *websocket.Conn) {
+	go func() {
+		select {
+		case <-ctx.Done():
+			_ = conn.Close(websocket.StatusNormalClosure, "")
+		}
+	}()
+}
+
+func (a *ActiveWebsockets) Close() {
+	a.cancel()
+	a.wg.Wait()
+}
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -435,9 +435,7 @@ func (api *API) workspaceAgentStartupLogs(rw http.ResponseWriter, r *http.Reques
 		return
 	}
 
-	api.WebsocketWaitMutex.Lock()
 	api.WebsocketWaitGroup.Add(1)
-	api.WebsocketWaitMutex.Unlock()
 	defer api.WebsocketWaitGroup.Done()
 	conn, err := websocket.Accept(rw, r, nil)
 	if err != nil {
@@ -559,9 +557,7 @@ func (api *API) workspaceAgentStartupLogs(rw http.ResponseWriter, r *http.Reques
 func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
-	api.WebsocketWaitMutex.Lock()
 	api.WebsocketWaitGroup.Add(1)
-	api.WebsocketWaitMutex.Unlock()
 	defer api.WebsocketWaitGroup.Done()
 
 	appToken, ok := workspaceapps.ResolveRequest(api.Logger, api.AccessURL, api.WorkspaceAppsProvider, rw, r, workspaceapps.Request{
@@ -816,9 +812,7 @@ func (api *API) workspaceAgentConnection(rw http.ResponseWriter, r *http.Request
 func (api *API) workspaceAgentCoordinate(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
-	api.WebsocketWaitMutex.Lock()
 	api.WebsocketWaitGroup.Add(1)
-	api.WebsocketWaitMutex.Unlock()
 	defer api.WebsocketWaitGroup.Done()
 	workspaceAgent := httpmw.WorkspaceAgent(r)
 	resource, err := api.Database.GetWorkspaceResourceByID(ctx, workspaceAgent.ResourceID)
@@ -1096,31 +1090,25 @@ func (api *API) workspaceAgentClientCoordinate(rw http.ResponseWriter, r *http.R
 		}
 	}
 
-	api.WebsocketWaitMutex.Lock()
-	api.WebsocketWaitGroup.Add(1)
-	api.WebsocketWaitMutex.Unlock()
-	defer api.WebsocketWaitGroup.Done()
 	workspaceAgent := httpmw.WorkspaceAgentParam(r)
 
-	conn, err := websocket.Accept(rw, r, nil)
-	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
-			Message: "Failed to accept websocket.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-	ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
-	defer wsNetConn.Close()
+	api.WebsocketWatch.Accept(rw, r, nil, func(conn *websocket.Conn) {
+		ctx, wsNetConn := websocketNetConn(ctx, conn, websocket.MessageBinary)
+		defer wsNetConn.Close()
 
-	go httpapi.Heartbeat(ctx, conn)
+		// Track for graceful shutdown.
+		api.WebsocketWatch.Add(wsNetConn)
+		defer api.WebsocketWatch.Done()
 
-	defer conn.Close(websocket.StatusNormalClosure, "")
-	err = (*api.TailnetCoordinator.Load()).ServeClient(wsNetConn, uuid.New(), workspaceAgent.ID)
-	if err != nil {
-		_ = conn.Close(websocket.StatusInternalError, err.Error())
-		return
-	}
+		go httpapi.Heartbeat(ctx, conn)
+
+		defer conn.Close(websocket.StatusNormalClosure, "")
+		err := (*api.TailnetCoordinator.Load()).ServeClient(wsNetConn, uuid.New(), workspaceAgent.ID)
+		if err != nil {
+			_ = conn.Close(websocket.StatusInternalError, err.Error())
+			return
+		}
+	})
 }
 
 func convertApps(dbApps []database.WorkspaceApp) []codersdk.WorkspaceApp {

diff --git a/enterprise/coderd/provisionerdaemons.go b/enterprise/coderd/provisionerdaemons.go
@@ -185,9 +185,7 @@ func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	api.AGPL.WebsocketWaitMutex.Lock()
 	api.AGPL.WebsocketWaitGroup.Add(1)
-	api.AGPL.WebsocketWaitMutex.Unlock()
 	defer api.AGPL.WebsocketWaitGroup.Done()
 
 	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{