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) }