fix: remove deadlock on .stop re-entrance

Author: 0x416e746f6e

proxy/websocket_pump.go

@@ -41,7 +41,7 @@ func newWebsocketPump(
 		frontend: frontend,
 		backend:  backend,
 		logger:   logger,
-		done:     make(chan struct{}),
+		done:     make(chan struct{}, 1),
 	}
 }
 
@@ -93,10 +93,19 @@ func (p *websocketPump) run() error {
 }
 
 func (p *websocketPump) stop() error {
-	p.done <- struct{}{}
+	if !p.active.Load() {
+		return nil // already stopped
+	}
 
 	errs := make([]error, 0)
 
+	select {
+	case p.done <- struct{}{}:
+		// no-op
+	default:
+		errs = append(errs, errors.New("double-closing the pump"))
+	}
+
 	countdown := 60
 	for p.active.Load() && countdown > 0 {
 		time.Sleep(time.Second)
@@ -130,8 +139,19 @@ func (p *websocketPump) pumpMessages(
 	)
 
 	messages := make(chan *websocketMessage, 16)
-	doneReads := make(chan struct{})
-	doneWrites := make(chan struct{})
+	doneReads := make(chan struct{}, 1)
+	doneWrites := make(chan struct{}, 1)
+
+	notifyOnFailure := func(err error) {
+		select {
+		case failure <- err:
+			// no-op
+		default:
+			l.Warn("Dropping websocket pump failure b/c the channel is full",
+				zap.Error(err),
+			)
+		}
+	}
 
 	go func() { // read
 		for {
@@ -141,13 +161,13 @@ func (p *websocketPump) pumpMessages(
 
 			default:
 				if err := from.SetReadDeadline(utils.Deadline(timeout)); err != nil {
-					failure <- err
+					notifyOnFailure(err)
 					continue
 				}
 
 				msgType, bytes, err := from.ReadMessage()
 				if err != nil {
-					failure <- err
+					notifyOnFailure(err)
 					continue
 				}
 
@@ -213,12 +233,12 @@ func (p *websocketPump) pumpMessages(
 				}
 
 				if err := to.SetWriteDeadline(utils.Deadline(timeout)); err != nil {
-					failure <- err
+					notifyOnFailure(err)
 					continue
 				}
 
 				if err := to.WriteMessage(m.msgType, m.bytes); err != nil {
-					failure <- err
+					notifyOnFailure(err)
 					continue
 				}