From d19e6b0fbb6830c531e216482426d1adaa2ef7ae Mon Sep 17 00:00:00 2001
From: Kasefuchs <kasefuchs@protonmail.com>
Date: Thu, 7 Nov 2024 11:47:02 +0300
Subject: [PATCH] :bug: Fix Nomad provider

---
 pkg/plugin/options.go            |  4 +--
 pkg/provider/nomad/allocation.go | 47 ++++++++++++++------------------
 2 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/pkg/plugin/options.go b/pkg/plugin/options.go
index 4befae4..e620313 100644
--- a/pkg/plugin/options.go
+++ b/pkg/plugin/options.go
@@ -20,10 +20,10 @@ type providerSelector func() (provider.Provider, error)
 // queuesSelector represents function used to select queues to use.
 type queuesSelector func() ([]queue.Queue, error)
 
-// configLoader represents function used to load pconfig configuration.
+// configLoader represents function used to load plugin configuration.
 type configLoader func() (*pconfig.Config, error)
 
-// Options contains customizable pconfig options.
+// Options contains customizable plugin options.
 type Options struct {
 	ProviderSelector providerSelector // Selector of provider.
 	QueuesSelector   queuesSelector   // Selector of available queues.
diff --git a/pkg/provider/nomad/allocation.go b/pkg/provider/nomad/allocation.go
index c1de583..f3c6e44 100644
--- a/pkg/provider/nomad/allocation.go
+++ b/pkg/provider/nomad/allocation.go
@@ -28,15 +28,24 @@ func NewAllocation(client *api.Client, item *item) *Allocation {
 	}
 }
 
-func (a *Allocation) scale(count *int64) error {
+func (a *Allocation) taskGroup() (*api.TaskGroup, error) {
 	job, _, err := a.client.Jobs().Info(*a.item.job.ID, nil)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	group := job.LookupTaskGroup(*a.item.group.Name)
 	if group == nil {
-		return fmt.Errorf("task group %s not found", *a.item.group.Name)
+		return nil, fmt.Errorf("task group %s not found", *a.item.group.Name)
+	}
+
+	return group, nil
+}
+
+func (a *Allocation) scale(count *int64) error {
+	group, err := a.taskGroup()
+	if err != nil {
+		return err
 	}
 
 	if int64(*group.Count) != *count {
@@ -66,45 +75,31 @@ func (a *Allocation) Start() error {
 }
 
 func (a *Allocation) State() provider.AllocationState {
-	info, _, err := a.client.Jobs().Info(*a.item.job.ID, nil)
+	group, err := a.taskGroup()
 	if err != nil {
 		return provider.AllocationStateUnknown
 	}
 
-	for _, group := range info.TaskGroups {
-		if group.Name != a.item.group.Name {
-			continue
-		}
-
-		if *group.Count == 0 {
-			return provider.AllocationStateStopped
-		}
-
-		return provider.AllocationStateStarted
+	if *group.Count == 0 {
+		return provider.AllocationStateStopped
 	}
 
-	return provider.AllocationStateUnknown
+	return provider.AllocationStateStarted
 }
 
 func (a *Allocation) Config() (*allocation.Config, error) {
-	info, _, err := a.client.Jobs().Info(*a.item.job.ID, nil)
+	group, err := a.taskGroup()
 	if err != nil {
 		return nil, err
 	}
 
-	for _, group := range info.TaskGroups {
-		if group.Name != a.item.group.Name {
+	for _, service := range group.Services {
+		cfg, err := allocation.ParseTags(service.Tags)
+		if err != nil {
 			continue
 		}
 
-		for _, service := range group.Services {
-			cfg, err := allocation.ParseTags(service.Tags)
-			if err != nil {
-				continue
-			}
-
-			return cfg, nil
-		}
+		return cfg, nil
 	}
 
 	return nil, fmt.Errorf("no allocation configuration found")