From 0c5ddac9f55f1b0e1e22f205845454ed91f7653d Mon Sep 17 00:00:00 2001 From: mountain Date: Thu, 15 Feb 2024 12:06:31 +0900 Subject: [PATCH] =?UTF-8?q?=ED=94=84=EB=A1=9C=EC=84=B8=EC=8A=A4=20?= =?UTF-8?q?=EC=A2=85=EB=A3=8C=20=EC=95=88=EB=90=98=EB=8A=94=20=EB=AC=B8?= =?UTF-8?q?=EC=A0=9C=20=EC=88=98=EC=A0=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- wshandler/wshandler.go | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/wshandler/wshandler.go b/wshandler/wshandler.go index 17c6756..887daf0 100644 --- a/wshandler/wshandler.go +++ b/wshandler/wshandler.go @@ -50,8 +50,9 @@ type DownstreamMessage struct { type commandType string const ( - commandType_EnterRoom = commandType("enter_room") - commandType_LeaveRoom = commandType("leave_room") + commandType_EnterRoom = commandType("enter_room") + commandType_LeaveRoom = commandType("leave_room") + ForceShutdownCloseMessage = "force_shutdown" ) type commandMessage struct { @@ -302,6 +303,7 @@ func (ws *WebsocketHandler) mainLoop(ctx context.Context) { logger.Println(r) } }() + c.closeMessage = ForceShutdownCloseMessage ws.ClientDisconnected(c) c.Close() } @@ -495,6 +497,11 @@ func upgrade_core(ws *WebsocketHandler, conn *websocket.Conn, accid primitive.Ob ws.connWaitGroup.Add(1) go func(c *wsconn, accid primitive.ObjectID, deliveryChan chan<- any) { + defer func() { + recover() + ws.connWaitGroup.Done() + }() + for { messageType, r, err := c.NextReader() if err != nil { @@ -519,10 +526,10 @@ func upgrade_core(ws *WebsocketHandler, conn *websocket.Conn, accid primitive.Ob ws.Call(newconn.sender, string(cmd), r) } } - ws.connWaitGroup.Done() - c.Conn = nil - ws.connInOutChan <- c + if c.closeMessage != ForceShutdownCloseMessage { + ws.connInOutChan <- c + } }(newconn, accid, ws.deliveryChan) }