#155 Minimise timing issues with websocket close.

internal/api/websocket.go

@@ -186,15 +186,9 @@ type rawToCodeOceanWriter struct {
 // Write implements the io.Writer interface.
 // The passed data is forwarded to the WebSocket to CodeOcean.
 func (rc *rawToCodeOceanWriter) Write(p []byte) (int, error) {
-	select {
-	case <-rc.proxy.webSocketCtx.Done():
-		return 0, nil
-	default:
-		log.Info("Passed WriteToCodeOceanCheck")
 	err := rc.proxy.sendToClient(dto.WebSocketMessage{Type: rc.outputType, Data: string(p)})
 	return len(p), err
 }
-}
 
 // webSocketProxy is an encapsulation of logic for forwarding between Runners and CodeOcean.
 type webSocketProxy struct {
@@ -308,6 +302,10 @@ func (wp *webSocketProxy) sendToClient(message dto.WebSocketMessage) error {
 		return fmt.Errorf("error marshaling WebSocket message: %w", err)
 	}
 	log.WithField("message", message).Trace("Sending message to client")
+	select {
+	case <-wp.webSocketCtx.Done():
+	default:
+		log.Info("Passed WriteToCodeOceanCheck")
 		err = wp.writeMessage(websocket.TextMessage, encodedMessage)
 		if err != nil {
 			errorMessage := "Error writing the message"
@@ -315,6 +313,7 @@ func (wp *webSocketProxy) sendToClient(message dto.WebSocketMessage) error {
 			wp.closeWithError(errorMessage)
 			return fmt.Errorf("error writing WebSocket message: %w", err)
 		}
+	}
 	return nil
 }