From a6874b21072c5056cc399aa734a23d85dabf4f8e Mon Sep 17 00:00:00 2001
From: Ivan Kozlovic <ivan@synadia.com>
Date: Mon, 10 Feb 2020 17:35:40 -0700
Subject: [PATCH] [FIXED] Server that is sent the lame duck mode signal does
 not exit

Resolves #1275

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>
---
 main.go          |  3 +--
 server/server.go | 13 ++++++++++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/main.go b/main.go
index 9b8d1483..b59a26f6 100644
--- a/main.go
+++ b/main.go
@@ -17,7 +17,6 @@ import (
 	"flag"
 	"fmt"
 	"os"
-	"runtime"
 
 	"github.com/nats-io/nats-server/v2/server"
 )
@@ -110,5 +109,5 @@ func main() {
 	if err := server.Run(s); err != nil {
 		server.PrintAndDie(err.Error())
 	}
-	runtime.Goexit()
+	s.WaitForShutdown()
 }
diff --git a/server/server.go b/server/server.go
index dd7f9ac4..ad2fb476 100644
--- a/server/server.go
+++ b/server/server.go
@@ -139,7 +139,8 @@ type Server struct {
 		dialTimeout time.Duration
 	}
 
-	quitCh chan struct{}
+	quitCh           chan struct{}
+	shutdownComplete chan struct{}
 
 	// Tracking Go routines
 	grMu         sync.Mutex
@@ -323,6 +324,9 @@ func NewServer(opts *Options) (*Server, error) {
 	// Used to kick out all go routines possibly waiting on server
 	// to shutdown.
 	s.quitCh = make(chan struct{})
+	// Closed when Shutdown() is complete. Allows WaitForShutdown() to block
+	// waiting for complete shutdown.
+	s.shutdownComplete = make(chan struct{})
 
 	// For tracking accounts
 	if err := s.configureAccounts(); err != nil {
@@ -1334,6 +1338,13 @@ func (s *Server) Shutdown() {
 			l.Close()
 		}
 	}
+	// Notify that the shutdown is complete
+	close(s.shutdownComplete)
+}
+
+// WaitForShutdown will block until the server has been fully shutdown.
+func (s *Server) WaitForShutdown() {
+	<-s.shutdownComplete
 }
 
 // AcceptLoop is exported for easier testing.