From 461c30969ebf67bbb5bbf8a98772b2af4049a5ce Mon Sep 17 00:00:00 2001
From: vishalnayak <vishalnayakv@gmail.com>
Date: Sun, 31 Jul 2016 10:09:16 -0400
Subject: [PATCH] Sharing shutdown message with physical consul backend

---
 command/server.go    | 6 +++---
 physical/consul.go   | 9 +++++----
 physical/physical.go | 2 +-
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/command/server.go b/command/server.go
index 13fc449cec..1e1d39b6ea 100644
--- a/command/server.go
+++ b/command/server.go
@@ -375,6 +375,8 @@ func (c *ServerCommand) Run(args []string) int {
 
 	// Instantiate the wait group
 	c.WaitGroup = &sync.WaitGroup{}
+	// Wait for shutdown
+	shutdownTriggered := false
 
 	// If the backend supports service discovery, run service discovery
 	if coreConfig.HAPhysical != nil && coreConfig.HAPhysical.HAEnabled() {
@@ -394,7 +396,7 @@ func (c *ServerCommand) Run(args []string) int {
 				return true
 			}
 
-			if err := sd.RunServiceDiscovery(c.WaitGroup, c.ShutdownCh, coreConfig.AdvertiseAddr, activeFunc, sealedFunc); err != nil {
+			if err := sd.RunServiceDiscovery(&shutdownTriggered, c.WaitGroup, c.ShutdownCh, coreConfig.AdvertiseAddr, activeFunc, sealedFunc); err != nil {
 				c.Ui.Error(fmt.Sprintf("Error initializing service discovery: %v", err))
 				return 1
 			}
@@ -419,8 +421,6 @@ func (c *ServerCommand) Run(args []string) int {
 	// Release the log gate.
 	logGate.Flush()
 
-	// Wait for shutdown
-	shutdownTriggered := false
 	for !shutdownTriggered {
 		select {
 		case <-c.ShutdownCh:
diff --git a/physical/consul.go b/physical/consul.go
index c379dd64de..754a314c3b 100644
--- a/physical/consul.go
+++ b/physical/consul.go
@@ -416,19 +416,19 @@ func (c *ConsulBackend) checkDuration() time.Duration {
 	return lib.DurationMinusBuffer(c.checkTimeout, checkMinBuffer, checkJitterFactor)
 }
 
-func (c *ConsulBackend) RunServiceDiscovery(waitGroup *sync.WaitGroup, shutdownCh ShutdownChannel, advertiseAddr string, activeFunc activeFunction, sealedFunc sealedFunction) (err error) {
+func (c *ConsulBackend) RunServiceDiscovery(shutdownTriggered *bool, waitGroup *sync.WaitGroup, shutdownCh ShutdownChannel, advertiseAddr string, activeFunc activeFunction, sealedFunc sealedFunction) (err error) {
 	if err := c.setAdvertiseAddr(advertiseAddr); err != nil {
 		return err
 	}
 
 	waitGroup.Add(1)
 
-	go c.runEventDemuxer(waitGroup, shutdownCh, advertiseAddr, activeFunc, sealedFunc)
+	go c.runEventDemuxer(shutdownTriggered, waitGroup, shutdownCh, advertiseAddr, activeFunc, sealedFunc)
 
 	return nil
 }
 
-func (c *ConsulBackend) runEventDemuxer(waitGroup *sync.WaitGroup, shutdownCh ShutdownChannel, advertiseAddr string, activeFunc activeFunction, sealedFunc sealedFunction) {
+func (c *ConsulBackend) runEventDemuxer(shutdownTriggered *bool, waitGroup *sync.WaitGroup, shutdownCh ShutdownChannel, advertiseAddr string, activeFunc activeFunction, sealedFunc sealedFunction) {
 	// Fire the reconcileTimer immediately upon starting the event demuxer
 	reconcileTimer := time.NewTimer(0)
 	defer reconcileTimer.Stop()
@@ -453,7 +453,7 @@ func (c *ConsulBackend) runEventDemuxer(waitGroup *sync.WaitGroup, shutdownCh Sh
 	var registeredServiceID string
 	var serviceRegLock int64
 shutdown:
-	for {
+	for !shutdown {
 		select {
 		case <-c.notifyActiveCh:
 			// Run reconcile immediately upon active state change notification
@@ -511,6 +511,7 @@ shutdown:
 			shutdown = true
 			break shutdown
 		}
+		shutdown = *shutdownTriggered
 	}
 
 	c.serviceLock.RLock()
diff --git a/physical/physical.go b/physical/physical.go
index 9e96beb6d8..e563e80e33 100644
--- a/physical/physical.go
+++ b/physical/physical.go
@@ -72,7 +72,7 @@ type ServiceDiscovery interface {
 
 	// Run executes any background service discovery tasks until the
 	// shutdown channel is closed.
-	RunServiceDiscovery(waitGroup *sync.WaitGroup, shutdownCh ShutdownChannel, advertiseAddr string, activeFunc activeFunction, sealedFunc sealedFunction) error
+	RunServiceDiscovery(shutdownTriggered *bool, waitGroup *sync.WaitGroup, shutdownCh ShutdownChannel, advertiseAddr string, activeFunc activeFunction, sealedFunc sealedFunction) error
 }
 
 type Lock interface {