From d075cff11fce61748a6307c3b024bcfa606ddb50 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Wed, 5 Jul 2023 15:14:42 +0300
Subject: [PATCH] Add warning log if websocket closing is stuck

---
 appservice/websocket.go | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/appservice/websocket.go b/appservice/websocket.go
index 671222b8..08957c25 100644
--- a/appservice/websocket.go
+++ b/appservice/websocket.go
@@ -374,7 +374,15 @@ func (as *AppService) StartWebsocket(baseURL string, onConnect func()) error {
 	closeChanOnce := sync.Once{}
 	stopFunc := func(err error) {
 		closeChanOnce.Do(func() {
-			closeChan <- err
+			select {
+			case closeChan <- err:
+			default:
+				as.Log.Warn().
+					AnErr("close_error", err).
+					Msg("Nothing is reading on close channel")
+				closeChan <- err
+				as.Log.Warn().Msg("Websocket close completed after being stuck")
+			}
 		})
 	}
 	as.ws = ws
@@ -384,11 +392,20 @@ func (as *AppService) StartWebsocket(baseURL string, onConnect func()) error {
 
 	go as.consumeWebsocket(stopFunc, ws)
 
+	var onConnectDone atomic.Bool
 	if onConnect != nil {
-		onConnect()
+		go func() {
+			onConnect()
+			onConnectDone.Store(true)
+		}()
+	} else {
+		onConnectDone.Store(true)
 	}
 
 	closeErr := <-closeChan
+	if onConnectDone.Load() {
+		as.Log.Warn().Msg("Websocket closed before onConnect returned, things may explode")
+	}
 
 	if as.ws == ws {
 		as.clearWebsocketResponseWaiters()