From b6e668bc53bfd2b8e28dad3bc5bec2975157beb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Sat, 11 Sep 2021 18:21:28 +0200
Subject: [PATCH] send disconnect on server shutdown.

---
 README.md                              |  1 +
 server/internal/websocket/websocket.go | 24 ++++++++++++++++++++++--
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index f972437..8318ede 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 d006368..f5222af 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: