Fix blocked streams blocking reconnects (#19017)

We have observed cases where a blocked stream will block for cancellations.

This happens when response channel is blocked and we want to push an error.
This will have the response mutex locked, which will prevent all other operations until upstream is unblocked.

Make this behavior non-blocking and if blocked spawn a goroutine that will send the response and close the output.

Still a lot of "dancing". Added a test for this and reviewed.

internal/grid/muxclient.go

@@ -50,6 +50,7 @@ type muxClient struct {
 	deadline         time.Duration
 	outBlock         chan struct{}
 	subroute         *subHandlerID
+	respErr          atomic.Pointer[error]
 }
 
 // Response is a response from the server.
@@ -250,25 +251,52 @@ func (m *muxClient) RequestStream(h HandlerID, payload []byte, requests chan []b
 
 	// Spawn simple disconnect
 	if requests == nil {
-		start := time.Now()
-		go m.handleOneWayStream(start, responseCh, responses)
-		return &Stream{responses: responseCh, Requests: nil, ctx: m.ctx, cancel: m.cancelFn}, nil
+		go m.handleOneWayStream(responseCh, responses)
+		return &Stream{responses: responseCh, Requests: nil, ctx: m.ctx, cancel: m.cancelFn, muxID: m.MuxID}, nil
 	}
 
 	// Deliver responses and send unblocks back to the server.
 	go m.handleTwowayResponses(responseCh, responses)
 	go m.handleTwowayRequests(responses, requests)
 
-	return &Stream{responses: responseCh, Requests: requests, ctx: m.ctx, cancel: m.cancelFn}, nil
+	return &Stream{responses: responseCh, Requests: requests, ctx: m.ctx, cancel: m.cancelFn, muxID: m.MuxID}, nil
 }
 
-func (m *muxClient) handleOneWayStream(start time.Time, respHandler chan<- Response, respServer <-chan Response) {
+func (m *muxClient) addErrorNonBlockingClose(respHandler chan<- Response, err error) {
+	m.respMu.Lock()
+	defer m.respMu.Unlock()
+	if !m.closed {
+		m.respErr.Store(&err)
+		// Do not block.
+		select {
+		case respHandler <- Response{Err: err}:
+			xioutil.SafeClose(respHandler)
+		default:
+			go func() {
+				respHandler <- Response{Err: err}
+				xioutil.SafeClose(respHandler)
+			}()
+		}
+		logger.LogIf(m.ctx, m.sendLocked(message{Op: OpDisconnectServerMux, MuxID: m.MuxID}))
+		m.closed = true
+	}
+}
+
+// respHandler
+func (m *muxClient) handleOneWayStream(respHandler chan<- Response, respServer <-chan Response) {
 	if debugPrint {
+		start := time.Now()
 		defer func() {
 			fmt.Println("Mux", m.MuxID, "Request took", time.Since(start).Round(time.Millisecond))
 		}()
 	}
-	defer xioutil.SafeClose(respHandler)
+	defer func() {
+		// addErrorNonBlockingClose will close the response channel
+		// - maybe async, so we shouldn't do it here.
+		if m.respErr.Load() == nil {
+			xioutil.SafeClose(respHandler)
+		}
+	}()
 	var pingTimer <-chan time.Time
 	if m.deadline == 0 || m.deadline > clientPingInterval {
 		ticker := time.NewTicker(clientPingInterval)
@@ -283,13 +311,7 @@ func (m *muxClient) handleOneWayStream(start time.Time, respHandler chan<- Respo
 			if debugPrint {
 				fmt.Println("Client sending disconnect to mux", m.MuxID)
 			}
-			m.respMu.Lock()
-			defer m.respMu.Unlock() // We always return in this path.
-			if !m.closed {
-				respHandler <- Response{Err: context.Cause(m.ctx)}
-				logger.LogIf(m.ctx, m.sendLocked(message{Op: OpDisconnectServerMux, MuxID: m.MuxID}))
-				m.closeLocked()
-			}
+			m.addErrorNonBlockingClose(respHandler, context.Cause(m.ctx))
 			return
 		case resp, ok := <-respServer:
 			if !ok {
@@ -308,13 +330,7 @@ func (m *muxClient) handleOneWayStream(start time.Time, respHandler chan<- Respo
 			}
 		case <-pingTimer:
 			if time.Since(time.Unix(atomic.LoadInt64(&m.LastPong), 0)) > clientPingInterval*2 {
-				m.respMu.Lock()
-				defer m.respMu.Unlock() // We always return in this path.
-				if !m.closed {
-					respHandler <- Response{Err: ErrDisconnected}
-					logger.LogIf(m.ctx, m.sendLocked(message{Op: OpDisconnectServerMux, MuxID: m.MuxID}))
-					m.closeLocked()
-				}
+				m.addErrorNonBlockingClose(respHandler, ErrDisconnected)
 				return
 			}
 			// Send new ping.
@@ -323,19 +339,21 @@ func (m *muxClient) handleOneWayStream(start time.Time, respHandler chan<- Respo
 	}
 }
 
-func (m *muxClient) handleTwowayResponses(responseCh chan Response, responses chan Response) {
+// responseCh is the channel to that goes to the requester.
+// internalResp is the channel that comes from the server.
+func (m *muxClient) handleTwowayResponses(responseCh chan<- Response, internalResp <-chan Response) {
 	defer m.parent.deleteMux(false, m.MuxID)
 	defer xioutil.SafeClose(responseCh)
-	for resp := range responses {
+	for resp := range internalResp {
 		responseCh <- resp
 		m.send(message{Op: OpUnblockSrvMux, MuxID: m.MuxID})
 	}
 }
 
-func (m *muxClient) handleTwowayRequests(responses chan<- Response, requests chan []byte) {
+func (m *muxClient) handleTwowayRequests(internalResp chan<- Response, requests <-chan []byte) {
 	var errState bool
-	start := time.Now()
 	if debugPrint {
+		start := time.Now()
 		defer func() {
 			fmt.Println("Mux", m.MuxID, "Request took", time.Since(start).Round(time.Millisecond))
 		}()
@@ -343,19 +361,22 @@ func (m *muxClient) handleTwowayRequests(responses chan<- Response, requests cha
 
 	// Listen for client messages.
 	for {
+		if errState {
+			go func() {
+				// Drain requests.
+				for range requests {
+				}
+			}()
+			return
+		}
 		select {
 		case <-m.ctx.Done():
 			if debugPrint {
 				fmt.Println("Client sending disconnect to mux", m.MuxID)
 			}
-			m.respMu.Lock()
-			defer m.respMu.Unlock()
-			logger.LogIf(m.ctx, m.sendLocked(message{Op: OpDisconnectServerMux, MuxID: m.MuxID}))
-			if !m.closed {
-				responses <- Response{Err: context.Cause(m.ctx)}
-				m.closeLocked()
-			}
-			return
+			m.addErrorNonBlockingClose(internalResp, context.Cause(m.ctx))
+			errState = true
+			continue
 		case req, ok := <-requests:
 			if !ok {
 				// Done send EOF
@@ -371,19 +392,14 @@ func (m *muxClient) handleTwowayRequests(responses chan<- Response, requests cha
 				msg.setZeroPayloadFlag()
 				err := m.send(msg)
 				if err != nil {
-					m.respMu.Lock()
-					responses <- Response{Err: err}
-					m.closeLocked()
-					m.respMu.Unlock()
+					m.addErrorNonBlockingClose(internalResp, err)
 				}
 				return
 			}
-			if errState {
-				continue
-			}
 			// Grab a send token.
 			select {
 			case <-m.ctx.Done():
+				m.addErrorNonBlockingClose(internalResp, context.Cause(m.ctx))
 				errState = true
 				continue
 			case <-m.outBlock:
@@ -398,8 +414,7 @@ func (m *muxClient) handleTwowayRequests(responses chan<- Response, requests cha
 			err := m.send(msg)
 			PutByteBuffer(req)
 			if err != nil {
-				responses <- Response{Err: err}
-				m.close()
+				m.addErrorNonBlockingClose(internalResp, err)
 				errState = true
 				continue
 			}
@@ -534,6 +549,7 @@ func (m *muxClient) closeLocked() {
 	if m.closed {
 		return
 	}
+	// We hold the lock, so nobody can modify m.respWait while we're closing.
 	if m.respWait != nil {
 		xioutil.SafeClose(m.respWait)
 		m.respWait = nil