func: add lock leaser on nomad agent run

agent/agent.go

@@ -24,10 +24,11 @@ import (
 )
 
 type Agent struct {
+	NomadClient *api.Client
+
 	logger        hclog.Logger
 	config        *config.Agent
 	configPaths   []string
-	nomadClient   *api.Client
 	pluginManager *manager.PluginManager
 	policySources map[policy.SourceName]policy.Source
 	policyManager *policy.Manager
@@ -56,11 +57,6 @@ func (a *Agent) Run(ctx context.Context) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	// Generate the Nomad client.
-	if err := a.generateNomadClient(); err != nil {
-		return err
-	}
-
 	// launch plugins
 	if err := a.setupPlugins(); err != nil {
 		return fmt.Errorf("failed to setup plugins: %v", err)
@@ -150,7 +146,7 @@ func (a *Agent) setupPolicyManager() (chan *sdk.ScalingEvaluation, error) {
 
 		switch policy.SourceName(s.Name) {
 		case policy.SourceNameNomad:
-			sources[policy.SourceNameNomad] = nomadPolicy.NewNomadSource(a.logger, a.nomadClient, policyProcessor)
+			sources[policy.SourceNameNomad] = nomadPolicy.NewNomadSource(a.logger, a.NomadClient, policyProcessor)
 		case policy.SourceNameFile:
 			// Only setup the file source if operators have configured a
 			// scaling policy directory to read from.
@@ -180,14 +176,14 @@ func (a *Agent) stop() {
 }
 
 // generateNomadClient creates a Nomad client for use within the agent.
-func (a *Agent) generateNomadClient() error {
+func (a *Agent) GenerateNomadClient() error {
 
 	// Generate the Nomad client.
 	client, err := api.NewClient(a.nomadCfg)
 	if err != nil {
 		return fmt.Errorf("failed to instantiate Nomad client: %v", err)
 	}
-	a.nomadClient = client
+	a.NomadClient = client
 
 	return nil
 }
@@ -209,7 +205,7 @@ func (a *Agent) reload() {
 	a.config = newCfg
 	a.nomadCfg = nomadHelper.MergeDefaultWithAgentConfig(newCfg.Nomad)
 
-	if err := a.generateNomadClient(); err != nil {
+	if err := a.GenerateNomadClient(); err != nil {
 		a.logger.Error("failed to reload Autoscaler configuration", "error", err)
 		os.Exit(1)
 	}
@@ -218,7 +214,7 @@ func (a *Agent) reload() {
 	// Set new Nomad client in the Nomad policy source.
 	ps, ok := a.policySources[policy.SourceNameNomad]
 	if ok {
-		ps.(*nomadPolicy.Source).SetNomadClient(a.nomadClient)
+		ps.(*nomadPolicy.Source).SetNomadClient(a.NomadClient)
 	}
 	a.policyManager.ReloadSources()
 

agent/agent_test.go

@@ -36,10 +36,10 @@ func TestAgent_generateNomadClient(t *testing.T) {
 	}
 
 	for _, tc := range testCases {
-		actualOutputErr := tc.inputAgent.generateNomadClient()
+		actualOutputErr := tc.inputAgent.GenerateNomadClient()
 		assert.Equal(t, tc.expectedOutputEr, actualOutputErr, tc.name)
 		if actualOutputErr == nil {
-			assert.Equal(t, tc.inputAgent.nomadCfg.Address, tc.inputAgent.nomadClient.Address(), tc.name)
+			assert.Equal(t, tc.inputAgent.nomadCfg.Address, tc.inputAgent.NomadClient.Address(), tc.name)
 		}
 	}
 }

command/agent.go

@@ -19,6 +19,7 @@ import (
 	flaghelper "github.com/hashicorp/nomad-autoscaler/sdk/helper/flag"
 	"github.com/hashicorp/nomad-autoscaler/sdk/helper/ptr"
 	"github.com/hashicorp/nomad-autoscaler/version"
+	"github.com/hashicorp/nomad/api"
 	"golang.org/x/text/cases"
 	"golang.org/x/text/language"
 )
@@ -338,7 +339,22 @@ func (c *AgentCommand) Run(args []string) int {
 
 	ctx := context.Background()
 
-	if err := c.agent.Run(ctx); err != nil {
+	asLock := api.Variable{
+		Path: "autoscaler/lock/",
+	}
+
+	// Generate the Nomad client.
+	if err := c.agent.GenerateNomadClient(); err != nil {
+		logger.Error("failed to start the nomad client", "error", err)
+		return 1
+	}
+
+	ll, err := c.agent.NomadClient.NewLockLeaser(api.WriteOptions{
+		Region:    parsedConfig.Nomad.Region,
+		Namespace: parsedConfig.Nomad.Region,
+	}, &asLock, 5*time.Minute, "")
+
+	if err := ll.Start(ctx, c.agent.Run); err != nil {
 		logger.Error("failed to start agent", "error", err)
 		return 1
 	}