diff --git a/README.md b/README.md index f972437b..8318edec 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,7 @@ For n.eko room management software, visit https://github.com/m1k1o/neko-rooms. - Switched to the latest Firefox version instead of esr. - Fixed very fast scroll speed on macOS. - Broadcast pipeline errors are reported to the user. +- On stopping server all websocket connections are going to be gracefully disconnected. # Getting started & FAQ diff --git a/server/internal/websocket/websocket.go b/server/internal/websocket/websocket.go index d006368d..f5222afd 100644 --- a/server/internal/websocket/websocket.go +++ b/server/internal/websocket/websocket.go @@ -3,6 +3,7 @@ package websocket import ( "fmt" "net/http" + "sync" "sync/atomic" "time" @@ -22,6 +23,7 @@ func New(sessions types.SessionManager, remote types.RemoteManager, broadcast ty return &WebSocketHandler{ logger: logger, + shutdown: make(chan interface{}), conf: conf, sessions: sessions, remote: remote, @@ -48,13 +50,14 @@ const pingPeriod = 60 * time.Second type WebSocketHandler struct { logger zerolog.Logger + wg sync.WaitGroup + shutdown chan interface{} upgrader websocket.Upgrader sessions types.SessionManager remote types.RemoteManager conf *config.WebSocket handler *MessageHandler conns uint32 - shutdown chan bool } func (ws *WebSocketHandler) Start() error { @@ -82,9 +85,11 @@ func (ws *WebSocketHandler) Start() error { } }) + ws.wg.Add(1) go func() { defer func() { ws.logger.Info().Msg("shutdown") + ws.wg.Done() }() current := ws.remote.ReadClipboard() @@ -116,7 +121,8 @@ func (ws *WebSocketHandler) Start() error { } func (ws *WebSocketHandler) Shutdown() error { - ws.shutdown <- true + close(ws.shutdown) + ws.wg.Wait() return nil } @@ -250,11 +256,13 @@ func (ws *WebSocketHandler) handle(connection *websocket.Conn, id string) { cancel := make(chan struct{}) ticker := time.NewTicker(pingPeriod) + ws.wg.Add(1) go func() { defer func() { ticker.Stop() ws.logger.Debug().Str("address", connection.RemoteAddr().String()).Msg("handle socket ending") ws.handler.Disconnected(id) + ws.wg.Done() }() for { @@ -283,6 +291,18 @@ func (ws *WebSocketHandler) handle(connection *websocket.Conn, id string) { if err := ws.handler.Message(id, raw); err != nil { ws.logger.Error().Err(err).Msg("message handler has failed") } + case <-ws.shutdown: + if err := connection.WriteJSON(message.SystemMessage{ + Event: event.SYSTEM_DISCONNECT, + Message: "server_shutdown", + }); err != nil { + ws.logger.Err(err).Msg("failed to send disconnect") + } + + if err := connection.Close(); err != nil { + ws.logger.Err(err).Msg("connection closed with an error") + } + return case <-cancel: return case <-ticker.C: