From c7828c1f7cdef0af9977635bf96b40d86c7d8217 Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Fri, 17 Jul 2020 11:04:57 -0400
Subject: [PATCH] Set AgentShutdown

---
 command/agent/agent.go | 1 +
 nomad/config.go        | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index 1e381b4fb..4c1a700c9 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -430,6 +430,7 @@ func (a *Agent) finalizeServerConfig(c *nomad.Config) {
 	// Setup the plugin loaders
 	c.PluginLoader = a.pluginLoader
 	c.PluginSingletonLoader = a.pluginSingletonLoader
+	c.AgentShutdown = func() error { return a.Shutdown() }
 }
 
 // clientConfig is used to generate a new client configuration struct for
diff --git a/nomad/config.go b/nomad/config.go
index ea3d71ba2..2e25e0471 100644
--- a/nomad/config.go
+++ b/nomad/config.go
@@ -351,6 +351,10 @@ type Config struct {
 
 	// LicenseConfig is a tunable knob for enterprise license testing.
 	LicenseConfig *LicenseConfig
+
+	// AgentShutdown is used to call agent.Shutdown from the context of a Server
+	// It is used primarily for licensing
+	AgentShutdown func() error
 }
 
 // CheckVersion is used to check if the ProtocolVersion is valid