From b34d0a4a6d3b9092f59481eff97d924e2ebf8e2d Mon Sep 17 00:00:00 2001
From: ankur22 <ankur.agarwal@grafana.com>
Date: Thu, 27 Jul 2023 15:34:17 +0100
Subject: [PATCH] Add case to redirect error responses early

If we receive an error response with an id, but not sessionId, there's
a slim chance that this response should be redirected to a session
but will instead be redirected to the connection's event loop. This
means that there is a chance that we could end up with a handler that
is waiting for a response indefinitely (or until a timeout occurs).

Since we cannot know exactly which session the response should be
redirected to, we're sending it to all live sessions as well as the
connection's event loop.

Most handlers should ignore the extra error message, and the handler
waiting for the msg with the unique msgId will action on it. If no
handlers are alive or the session has already been closed that this
response msg should go to, then no handlers will action on it either.
---
 common/connection.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/common/connection.go b/common/connection.go
index b0bc951f6..32455455f 100644
--- a/common/connection.go
+++ b/common/connection.go
@@ -368,6 +368,36 @@ func (c *Connection) recvLoop() {
 				return
 			}
 
+		// In some cases the error response will only have the id and error,
+		// but no sessionId. In these cases we can't guarantee the origin of
+		// the request and so where the msg should be redirected to. To ensure
+		// the msg gets to the correct handler (which is potentially blocking
+		// a test iteration) we need to send it to all sessions and the
+		// connection's event loop.
+		case msg.ID != 0 && msg.Error != nil && msg.SessionID == "":
+			c.sessionsMu.RLock()
+			for _, s := range c.sessions {
+				select {
+				case s.readCh <- &msg:
+				case code := <-c.closeCh:
+					c.logger.Debugf(
+						"Connection:recvLoop:msg.ID:msg.Error:<-c.closeCh",
+						"sid:%v tid:%v wsURL:%v crashed:%t",
+						s.id, s.targetID, c.wsURL, s.crashed,
+					)
+					_ = c.close(code)
+				case <-c.done:
+					c.logger.Debugf(
+						"Connection:recvLoop:msg.ID:msg.Error:<-c.done",
+						"sid:%v tid:%v wsURL:%v crashed:%t",
+						s.id, s.targetID, c.wsURL, s.crashed,
+					)
+					return
+				}
+			}
+			c.sessionsMu.RUnlock()
+			c.emit("", &msg)
+
 		case msg.Method != "":
 			c.logger.Debugf("Connection:recvLoop:msg.Method:emit", "sid:%v method:%q", msg.SessionID, msg.Method)
 			ev, err := cdproto.UnmarshalMessage(&msg)