From f3de3c0e94373f61f392e8efcfd553b8dd7bfa4a Mon Sep 17 00:00:00 2001
From: Ivan Kozlovic <ivan@synadia.com>
Date: Tue, 7 May 2024 20:23:27 -0600
Subject: [PATCH] Made a function to close the websocket server

The function returns 1 if the server was running 0 otherwise, this
is used by the caller to increase the number of things to wait for
on shutdown.
Did capture the server under the lock and close outside - just in case.

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>
---
 server/server.go | 36 ++++++++++++++++++++----------------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/server/server.go b/server/server.go
index 3487701ad0..89e9c7de23 100644
--- a/server/server.go
+++ b/server/server.go
@@ -2478,14 +2478,7 @@ func (s *Server) Shutdown() {
 	}
 
 	// Kick websocket server
-	s.websocket.mu.Lock()
-	if s.websocket.server != nil {
-		doneExpected++
-		s.websocket.server.Close()
-		s.websocket.server = nil
-		s.websocket.listener = nil
-	}
-	s.websocket.mu.Unlock()
+	doneExpected += s.closeWebsocketServer()
 
 	// Kick MQTT accept loop
 	if s.mqtt.listener != nil {
@@ -2573,6 +2566,24 @@ func (s *Server) Shutdown() {
 	close(s.shutdownComplete)
 }
 
+// Close the websocket server if running. If so, returns 1, else 0.
+// Server lock held on entry.
+func (s *Server) closeWebsocketServer() int {
+	ws := &s.websocket
+	ws.mu.Lock()
+	hs := ws.server
+	if hs != nil {
+		ws.server = nil
+		ws.listener = nil
+	}
+	ws.mu.Unlock()
+	if hs != nil {
+		hs.Close()
+		return 1
+	}
+	return 0
+}
+
 // WaitForShutdown will block until the server has been fully shutdown.
 func (s *Server) WaitForShutdown() {
 	<-s.shutdownComplete
@@ -4092,14 +4103,7 @@ func (s *Server) lameDuckMode() {
 	expected := 1
 	s.listener.Close()
 	s.listener = nil
-	s.websocket.mu.Lock()
-	if s.websocket.server != nil {
-		expected++
-		s.websocket.server.Close()
-		s.websocket.server = nil
-		s.websocket.listener = nil
-	}
-	s.websocket.mu.Unlock()
+	expected += s.closeWebsocketServer()
 	s.ldmCh = make(chan bool, expected)
 	opts := s.getOpts()
 	gp := opts.LameDuckGracePeriod