Don't leak goroutines everywhere in the ws code

js/modules/k6/ws/ws.go

@@ -247,7 +247,7 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP
 	readErrChan := make(chan error)
 
 	// Wraps a couple of channels around conn.ReadMessage
-	go readPump(conn, readDataChan, readErrChan, readCloseChan)
+	go socket.readPump(readDataChan, readErrChan, readCloseChan)
 
 	// This is the main control loop. All JS code (including error handlers)
 	// should only be executed by this thread to avoid race conditions
@@ -384,7 +384,11 @@ func (s *Socket) SetTimeout(fn goja.Callable, timeoutMs int) {
 	go func() {
 		select {
 		case <-time.After(time.Duration(timeoutMs) * time.Millisecond):
-			s.scheduled <- fn
+			select {
+			case s.scheduled <- fn:
+			case <-s.done:
+				return
+			}
 
 		case <-s.done:
 			return
@@ -402,7 +406,11 @@ func (s *Socket) SetInterval(fn goja.Callable, intervalMs int) {
 		for {
 			select {
 			case <-ticker.C:
-				s.scheduled <- fn
+				select {
+				case s.scheduled <- fn:
+				case <-s.done:
+					return
+				}
 
 			case <-s.done:
 				return
@@ -450,24 +458,35 @@ func (s *Socket) closeConnection(code int) error {
 }
 
 // Wraps conn.ReadMessage in a channel
-func readPump(conn *websocket.Conn, readChan chan []byte, errorChan chan error, closeChan chan int) {
+func (s *Socket) readPump(readChan chan []byte, errorChan chan error, closeChan chan int) {
 	for {
-		_, message, err := conn.ReadMessage()
+		_, message, err := s.conn.ReadMessage()
 		if err != nil {
 			if websocket.IsUnexpectedCloseError(
 				err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
 				// Report an unexpected closure
-				errorChan <- err
+				select {
+				case errorChan <- err:
+				case <-s.done:
+					return
+				}
 			}
 			code := websocket.CloseGoingAway
 			if e, ok := err.(*websocket.CloseError); ok {
 				code = e.Code
 			}
-			closeChan <- code
+			select {
+			case closeChan <- code:
+			case <-s.done:
+			}
 			return
 		}
 
-		readChan <- message
+		select {
+		case readChan <- message:
+		case <-s.done:
+			return
+		}
 	}
 }
 

js/modules/k6/ws/ws_test.go

@@ -581,7 +581,8 @@ func TestReadPump(t *testing.T) {
 			msgChan := make(chan []byte)
 			errChan := make(chan error)
 			closeChan := make(chan int)
-			go readPump(conn, msgChan, errChan, closeChan)
+			s := &Socket{conn: conn}
+			go s.readPump(msgChan, errChan, closeChan)
 
 		readChans:
 			for {