fix: Race when shutting down and opening WebSockets (#576)

Adding to a WaitGroup while calling wait is a race condition. Surrounding
this in a mutex should solve the problem. Since context is used for
cancellation on all sockets, cleanup should occur properly.

See: https://github.com/coder/coder/runs/5701221057?check_suite_focus=true#step:10:98

coderd/coderd.go

@@ -176,13 +176,18 @@ func New(options *Options) (http.Handler, func()) {
 		})
 	})
 	r.NotFound(site.DefaultHandler().ServeHTTP)
-	return r, api.websocketWaitGroup.Wait
+	return r, func() {
+		api.websocketWaitMutex.Lock()
+		api.websocketWaitGroup.Wait()
+		api.websocketWaitMutex.Unlock()
+	}
 }
 
 // API contains all route handlers. Only HTTP handlers should
 // be added to this struct for code clarity.
 type api struct {
 	*Options
 
+	websocketWaitMutex sync.Mutex
 	websocketWaitGroup sync.WaitGroup
 }

coderd/provisionerdaemons.go

@@ -33,7 +33,9 @@ import (
 
 // Serves the provisioner daemon protobuf API over a WebSocket.
 func (api *api) provisionerDaemonsListen(rw http.ResponseWriter, r *http.Request) {
+	api.websocketWaitMutex.Lock()
 	api.websocketWaitGroup.Add(1)
+	api.websocketWaitMutex.Unlock()
 	defer api.websocketWaitGroup.Done()
 
 	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{

coderd/workspaceresources.go

@@ -64,7 +64,9 @@ func (api *api) workspaceResource(rw http.ResponseWriter, r *http.Request) {
 }
 
 func (api *api) workspaceResourceDial(rw http.ResponseWriter, r *http.Request) {
+	api.websocketWaitMutex.Lock()
 	api.websocketWaitGroup.Add(1)
+	api.websocketWaitMutex.Unlock()
 	defer api.websocketWaitGroup.Done()
 
 	resource := httpmw.WorkspaceResourceParam(r)
@@ -112,7 +114,9 @@ func (api *api) workspaceResourceDial(rw http.ResponseWriter, r *http.Request) {
 }
 
 func (api *api) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
+	api.websocketWaitMutex.Lock()
 	api.websocketWaitGroup.Add(1)
+	api.websocketWaitMutex.Unlock()
 	defer api.websocketWaitGroup.Done()
 
 	agent := httpmw.WorkspaceAgent(r)