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 <[email protected]>

server/server.go

@@ -2521,14 +2521,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 {
@@ -2616,6 +2609,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
@@ -4135,14 +4146,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