func: add control context to run agent to prepare for ha mode (#651)

hashicorp · Jul 5, 2023 · d715a75 · d715a75
1 parent d05ab8f
commit d715a75
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/agent/agent.go b/agent/agent.go
@@ -49,11 +49,11 @@ func NewAgent(c *config.Agent, configPaths []string, logger hclog.Logger) *Agent
 	}
 }
 
-func (a *Agent) Run() error {
+func (a *Agent) Run(ctx context.Context) error {
 	defer a.stop()
 
 	// Create context to handle propagation to downstream routines.
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	// Generate the Nomad client.

diff --git a/command/agent.go b/command/agent.go
@@ -4,6 +4,7 @@
 package command
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"sort"
@@ -335,7 +336,9 @@ func (c *AgentCommand) Run(args []string) int {
 	go c.httpServer.Start()
 	defer c.httpServer.Stop()
 
-	if err := c.agent.Run(); err != nil {
+	ctx := context.Background()
+
+	if err := c.agent.Run(ctx); err != nil {
 		logger.Error("failed to start agent", "error", err)
 		return 1
 	}