lib/protocol: Prioritize close msg and add close timeout (#5746)

lib/protocol/protocol.go

@@ -184,6 +184,7 @@ type rawConnection struct {
 
 	inbox                 chan message
 	outbox                chan asyncMessage
+	closeBox              chan asyncMessage
 	clusterConfigBox      chan *ClusterConfig
 	dispatcherLoopStopped chan struct{}
 	closed                chan struct{}
@@ -218,6 +219,11 @@ const (
 	ReceiveTimeout = 300 * time.Second
 )
 
+// CloseTimeout is the longest we'll wait when trying to send the close
+// message before just closing the connection.
+// Should not be modified in production code, just for testing.
+var CloseTimeout = 10 * time.Second
+
 func NewConnection(deviceID DeviceID, reader io.Reader, writer io.Writer, receiver Model, name string, compress Compression) Connection {
 	cr := &countingReader{Reader: reader}
 	cw := &countingWriter{Writer: writer}
@@ -231,6 +237,7 @@ func NewConnection(deviceID DeviceID, reader io.Reader, writer io.Writer, receiv
 		awaiting:              make(map[int32]chan asyncResult),
 		inbox:                 make(chan message),
 		outbox:                make(chan asyncMessage),
+		closeBox:              make(chan asyncMessage),
 		clusterConfigBox:      make(chan *ClusterConfig),
 		dispatcherLoopStopped: make(chan struct{}),
 		closed:                make(chan struct{}),
@@ -671,6 +678,10 @@ func (c *rawConnection) writerLoop() {
 			c.internalClose(err)
 			return
 		}
+	case hm := <-c.closeBox:
+		_ = c.writeMessage(hm.msg)
+		close(hm.done)
+		return
 	case <-c.closed:
 		return
 	}
@@ -686,6 +697,11 @@ func (c *rawConnection) writerLoop() {
 				return
 			}
 
+		case hm := <-c.closeBox:
+			_ = c.writeMessage(hm.msg)
+			close(hm.done)
+			return
+
 		case <-c.closed:
 			return
 		}
@@ -853,17 +869,20 @@ func (c *rawConnection) shouldCompressMessage(msg message) bool {
 func (c *rawConnection) Close(err error) {
 	c.sendCloseOnce.Do(func() {
 		done := make(chan struct{})
-		c.send(&Close{err.Error()}, done)
+		timeout := time.NewTimer(CloseTimeout)
 		select {
-		case <-done:
+		case c.closeBox <- asyncMessage{&Close{err.Error()}, done}:
+			select {
+			case <-done:
+			case <-timeout.C:
+			case <-c.closed:
+			}
+		case <-timeout.C:
 		case <-c.closed:
 		}
 	})
 
-	// No more sends are necessary, therefore further steps to close the
-	// connection outside of this package can proceed immediately.
-	// And this prevents a potential deadlock due to calling c.receiver.Closed
-	go c.internalClose(err)
+	c.internalClose(err)
 }
 
 // internalClose is called if there is an unexpected error during normal operation.