Clean up websocket close handling a bit (#4508)

* Be more tolerant of errors in transport Close() Signed-off-by: Eric Chlebek <[email protected]> * Don't double send close message Signed-off-by: Eric Chlebek <[email protected]>
sensu · Nov 19, 2021 · 62a9274 · 62a9274
1 parent 3f13ae0
commit 62a9274
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 15 deletions.
diff --git a/backend/agentd/session.go b/backend/agentd/session.go
@@ -555,15 +555,6 @@ func (s *Session) stop() {
 		}
 	}()
 
-	// Send a close message to ensure the agent closes its connection if the
-	// connection is not already closed
-	if !s.conn.Closed() {
-		if err := s.conn.SendCloseMessage(); err != nil {
-			websocketErrorCounter.WithLabelValues("send", "SendCloseMessage").Inc()
-			logger.Warning("unexpected error while sending a close message to the agent")
-		}
-	}
-
 	sessionCounter.WithLabelValues(s.cfg.Namespace).Dec()
 
 	// Gracefully wait for the send and receiver to exit, but force the websocket

diff --git a/transport/transport.go b/transport/transport.go
@@ -146,10 +146,8 @@ func NewMessage(msgType string, payload []byte) *Message {
 	}
 }
 
-// Close attempts to send a "going away" message over the websocket connection.
-// This will cause a Write over the websocket transport, which can cause a
-// panic. We rescue potential panics and consider the connection closed,
-// returning nil, because the connection _will_ be closed. Hay!
+// Close closes the WebsocketTransport. Before closing, a closing message will
+// be sent to the other side.
 func (t *WebSocketTransport) Close() (err error) {
 	if t.Closed() {
 		return nil
@@ -162,12 +160,16 @@ func (t *WebSocketTransport) Close() (err error) {
 
 	defer func() {
 		cerr := t.Connection.Close()
-		if err == nil {
+		if err == nil && cerr != websocket.ErrCloseSent {
 			err = cerr
 		}
 	}()
 
-	return t.Connection.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, "bye"))
+	if err := t.Connection.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, "bye")); err != websocket.ErrCloseSent {
+		return err
+	}
+
+	return nil
 }
 
 // Closed returns true if the underlying websocket connection has been closed.