provide vault server flag to exit on core shutdown (#7561)

* provide vault server flag to exit on core shutdown

* Update command/server.go

Co-Authored-By: Jeff Mitchell <jeffrey.mitchell@gmail.com>

Co-authored-by: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Co-authored-by: Brian Kassouf <briankassouf@users.noreply.github.com>

command/server.go

@@ -126,6 +126,7 @@ type ServerCommand struct {
 	flagCombineLogs      bool
 	flagTestServerConfig bool
 	flagDevConsul        bool
+  flagExitOnCoreShutdown bool
 }
 
 type ServerListener struct {
@@ -204,6 +205,13 @@ func (c *ServerCommand) Flags() *FlagSets {
 		Usage:      `Log format. Supported values are "standard" and "json".`,
 	})
 
+	f.BoolVar(&BoolVar{
+		Name:    "exit-on-core-shutdown",
+		Target:  &c.flagExitOnCoreShutdown,
+		Default: false,
+		Usage:   "Exit the vault server if the vault core is shutdown.",
+	})
+
 	f.BoolVar(&BoolVar{
 		Name:   "recovery",
 		Target: &c.flagRecovery,
@@ -1787,26 +1795,24 @@ CLUSTER_SYNTHESIS_COMPLETE:
 		}
 	}()
 
+	var coreShutdownDoneCh <-chan struct{}
+	if c.flagExitOnCoreShutdown {
+		coreShutdownDoneCh = core.ShutdownDone()
+	}
+
 	// Wait for shutdown
 	shutdownTriggered := false
+	retCode := 0
 
 	for !shutdownTriggered {
 		select {
+		case <-coreShutdownDoneCh:
+			c.UI.Output("==> Vault core was shut down")
+			retCode = 1
+			shutdownTriggered = true
 		case <-c.ShutdownCh:
 			c.UI.Output("==> Vault shutdown triggered")
-
-			// Stop the listeners so that we don't process further client requests.
-			c.cleanupGuard.Do(listenerCloseFunc)
-
-			// Shutdown will wait until after Vault is sealed, which means the
-			// request forwarding listeners will also be closed (and also
-			// waited for).
-			if err := core.Shutdown(); err != nil {
-				c.UI.Error(fmt.Sprintf("Error with core shutdown: %s", err))
-			}
-
 			shutdownTriggered = true
-
 		case <-c.SighupCh:
 			c.UI.Output("==> Vault reload triggered")
 
@@ -1867,9 +1873,19 @@ CLUSTER_SYNTHESIS_COMPLETE:
 		}
 	}
 
+	// Stop the listeners so that we don't process further client requests.
+	c.cleanupGuard.Do(listenerCloseFunc)
+
+	// Shutdown will wait until after Vault is sealed, which means the
+	// request forwarding listeners will also be closed (and also
+	// waited for).
+	if err := core.Shutdown(); err != nil {
+		c.UI.Error(fmt.Sprintf("Error with core shutdown: %s", err))
+	}
+
 	// Wait for dependent goroutines to complete
 	c.WaitGroup.Wait()
-	return 0
+	return retCode
 }
 
 func (c *ServerCommand) enableDev(core *vault.Core, coreConfig *vault.CoreConfig) (*vault.InitResult, error) {

vault/core.go

@@ -267,6 +267,9 @@ type Core struct {
 	keepHALockOnStepDown *uint32
 	heldHALock           physical.Lock
 
+	// shutdownDoneCh is used to notify when Shutdown() completes
+	shutdownDoneCh chan struct{}
+
 	// unlockInfo has the keys provided to Unseal until the threshold number of parts is available, as well as the operation nonce
 	unlockInfo *unlockInformation
 
@@ -737,6 +740,7 @@ func NewCore(conf *CoreConfig) (*Core, error) {
 		clusterPeerClusterAddrsCache: cache.New(3*cluster.HeartbeatInterval, time.Second),
 		enableMlock:                  !conf.DisableMlock,
 		rawEnabled:                   conf.EnableRaw,
+		shutdownDoneCh:               make(chan struct{}),
 		replicationState:             new(uint32),
 		atomicPrimaryClusterAddrs:    new(atomic.Value),
 		atomicPrimaryFailoverAddrs:   new(atomic.Value),
@@ -933,7 +937,21 @@ func NewCore(conf *CoreConfig) (*Core, error) {
 // happens as quickly as possible.
 func (c *Core) Shutdown() error {
 	c.logger.Debug("shutdown called")
-	return c.sealInternal()
+	err := c.sealInternal()
+
+	c.stateLock.Lock()
+	defer c.stateLock.Unlock()
+	if c.shutdownDoneCh != nil {
+		close(c.shutdownDoneCh)
+		c.shutdownDoneCh = nil
+	}
+
+	return err
+}
+
+// ShutdownDone returns a channel that will be closed after Shutdown completes
+func (c *Core) ShutdownDone() <-chan struct{} {
+	return c.shutdownDoneCh
 }
 
 // CORSConfig returns the current CORS configuration

vault/core_test.go

@@ -247,6 +247,25 @@ func TestCore_Shutdown(t *testing.T) {
 	}
 }
 
+// verify the channel returned by ShutdownDone is closed after Shutdown
+func TestCore_ShutdownDone(t *testing.T) {
+	c, _, _ := TestCoreUnsealed(t)
+	doneCh := c.ShutdownDone()
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		c.Shutdown()
+	}()
+
+	select {
+	case <-doneCh:
+		if !c.Sealed() {
+			t.Fatalf("shutdown done called prematurely!")
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatalf("shutdown notification not received")
+	}
+}
+
 // Attempt to seal bad token
 func TestCore_Seal_BadToken(t *testing.T) {
 	c, _, _ := TestCoreUnsealed(t)