From 4f7bd48bece1a4b4a634819913452c852d8602fb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Thu, 9 Sep 2021 23:55:53 +0200
Subject: [PATCH] ws close connections on shutdown and add wg.

---
 internal/websocket/manager.go | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go
index 5a347425..3d68d74c 100644
--- a/internal/websocket/manager.go
+++ b/internal/websocket/manager.go
@@ -3,6 +3,7 @@ package websocket
 import (
 	"encoding/json"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -28,6 +29,7 @@ func New(
 
 	return &WebSocketManagerCtx{
 		logger:   logger,
+		shutdown: make(chan interface{}),
 		sessions: sessions,
 		desktop:  desktop,
 		handler:  handler.New(sessions, desktop, capture, webrtc),
@@ -37,6 +39,8 @@ func New(
 
 type WebSocketManagerCtx struct {
 	logger   zerolog.Logger
+	wg       sync.WaitGroup
+	shutdown chan interface{}
 	sessions types.SessionManager
 	desktop  types.DesktopManager
 	handler  *handler.MessageHandlerCtx
@@ -132,7 +136,8 @@ func (manager *WebSocketManagerCtx) Start() {
 
 func (manager *WebSocketManagerCtx) Shutdown() error {
 	manager.logger.Info().Msg("shutdown")
-	// TODO: Kill all connections and add waitgroup for gorutines.
+	close(manager.shutdown)
+	manager.wg.Wait()
 	return nil
 }
 
@@ -218,7 +223,10 @@ func (manager *WebSocketManagerCtx) handle(connection *websocket.Conn, session t
 	ticker := time.NewTicker(pingPeriod)
 	defer ticker.Stop()
 
+	manager.wg.Add(1)
 	go func() {
+		defer manager.wg.Done()
+
 		for {
 			_, raw, err := connection.ReadMessage()
 			if err != nil {
@@ -267,6 +275,9 @@ func (manager *WebSocketManagerCtx) handle(connection *websocket.Conn, session t
 			}
 		case <-cancel:
 			return
+		case <-manager.shutdown:
+			connection.Close()
+			return
 		case <-ticker.C:
 			if err := connection.WriteMessage(websocket.PingMessage, nil); err != nil {
 				logger.Err(err).Msg("ping message has failed")