refactor(collector): 🔧 Centralize command execution and add timeout

- Centralizes all external command executions into a new `collector.Execute` function. This function handles logging, errors, and context-based timeouts.
- Replaces all `exec.Command` calls in the collectors with the new centralized function, removing hardcoded binary paths and fixing latent bugs in argument passing.
- Makes the execution logic mockable for improved testing.
- Adds a `--command.timeout` flag (default 5s) to prevent the exporter from hanging on unresponsive Slurm commands.
- Adds a `--log.level` flag to control logging verbosity.
- Cleans up verbose debug logging for successful commands.

Author: SckyzO

collector/accounts.go

@@ -16,8 +16,6 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>. */
 package collector
 
 import (
-	"io"
-	"os/exec"
 	"regexp"
 	"strconv"
 	"strings"
@@ -28,22 +26,7 @@ import (
 )
 
 func AccountsData(logger log.Logger) ([]byte, error) {
-	cmd := exec.Command("squeue", "-a", "-r", "-h", "-o %A|%a|%T|%C")
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return nil, err
-	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return nil, err
-	}
-	out, _ := io.ReadAll(stdout)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err)
-		return nil, err
-	}
-	return out, nil
+	return Execute(logger, "squeue", []string{"-a", "-r", "-h", "-o", "%A|%a|%T|%C"})
 }
 
 type JobMetrics struct {

collector/cpus.go

@@ -16,8 +16,6 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>. */
 package collector
 
 import (
-	"io"
-	"os/exec"
 	"strconv"
 	"strings"
 
@@ -55,22 +53,7 @@ func ParseCPUsMetrics(input []byte) *CPUsMetrics {
 
 // Execute the sinfo command and return its output
 func CPUsData(logger log.Logger) ([]byte, error) {
-	cmd := exec.Command("sinfo", "-h", "-o %C")
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return nil, err
-	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return nil, err
-	}
-	out, _ := io.ReadAll(stdout)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err)
-		return nil, err
-	}
-	return out, nil
+	return Execute(logger, "sinfo", []string{"-h", "-o", "%C"})
 }
 
 /*

collector/execute.go

@@ -0,0 +1,43 @@
+package collector
+
+import (
+	"context"
+	"os/exec"
+	"strings"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+)
+
+var (
+	commandTimeout time.Duration
+)
+
+// SetCommandTimeout sets the timeout for external commands.
+func SetCommandTimeout(t time.Duration) {
+	commandTimeout = t
+}
+
+// Execute is a wrapper around exec.CommandContext to provide logging and a timeout.
+var Execute = func(logger log.Logger, command string, args []string) ([]byte, error) {
+	level.Debug(logger).Log("msg", "Executing command", "command", command, "args", strings.Join(args, " "))
+
+	ctx, cancel := context.WithTimeout(context.Background(), commandTimeout)
+	defer cancel()
+
+	cmd := exec.CommandContext(ctx, command, args...)
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		// Check if the error is due to the context deadline exceeding.
+		if ctx.Err() == context.DeadlineExceeded {
+			level.Error(logger).Log("msg", "Command timed out", "command", command, "args", strings.Join(args, " "), "timeout", commandTimeout)
+			return nil, ctx.Err()
+		}
+		level.Error(logger).Log("msg", "Failed to execute command", "command", command, "args", strings.Join(args, " "), "output", string(out), "err", err)
+		return nil, err
+	}
+
+	level.Debug(logger).Log("msg", "Command executed successfully", "command", command)
+	return out, nil
+}

collector/gpus.go

@@ -1,7 +1,6 @@
 package collector
 
 import (
-	"os/exec"
 	"regexp"
 	"strconv"
 	"strings"
@@ -211,33 +210,21 @@ func ParseGPUsMetrics(logger log.Logger) (*GPUsMetrics, error) {
 	return &gm, nil
 }
 
-var executeCommand = Execute
-
 func AllocatedGPUsData(logger log.Logger) ([]byte, error) {
 	args := []string{"-a", "-h", "--Format=Nodes: ,GresUsed:", "--state=allocated"}
-	return executeCommand(logger, "sinfo", args)
+	return Execute(logger, "sinfo", args)
 }
 
 func IdleGPUsData(logger log.Logger) ([]byte, error) {
 	args := []string{"-a", "-h", "--Format=Nodes: ,Gres: ,GresUsed:", "--state=idle,allocated"}
-	return executeCommand(logger, "sinfo", args)
+	return Execute(logger, "sinfo", args)
 }
 
 func TotalGPUsData(logger log.Logger) ([]byte, error) {
 	args := []string{"-a", "-h", "--Format=Nodes: ,Gres:"}
-	return executeCommand(logger, "sinfo", args)
+	return Execute(logger, "sinfo", args)
 }
 
-// Execute the sinfo command and return its output
-func Execute(logger log.Logger, command string, arguments []string) ([]byte, error) {
-	cmd := exec.Command(command, arguments...)
-	out, err := cmd.CombinedOutput()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to execute command", "command", command, "args", strings.Join(arguments, " "), "err", err)
-		return nil, err
-	}
-	return out, nil
-}
 
 /*
  * Implement the Prometheus Collector interface and feed the

collector/gpus_test.go

@@ -61,14 +61,14 @@ func TestGPUsMetrics(t *testing.T) {
 }
 
 func TestGPUsGetMetrics(t *testing.T) {
-	oldExecuteCommand := executeCommand
-	defer func() { executeCommand = oldExecuteCommand }()
+	oldExecute := Execute
+	defer func() { Execute = oldExecute }()
 
 	test_data_paths, _ := filepath.Glob("../test_data/slurm-*")
 	for _, test_data_path := range test_data_paths {
 		slurm_version := strings.TrimPrefix(test_data_path, "../test_data/slurm-")
 		t.Run(slurm_version, func(t *testing.T) {
-			executeCommand = func(logger log.Logger, command string, arguments []string) ([]byte, error) {
+			Execute = func(logger log.Logger, command string, arguments []string) ([]byte, error) {
 				var file string
 				if strings.Contains(arguments[2], "GresUsed:") && strings.Contains(arguments[2], "Gres:") {
 					file = filepath.Join(test_data_path, "sinfo_gpus_idle.txt")

collector/node.go

@@ -1,7 +1,6 @@
 package collector
 
 import (
-	"os/exec"
 	"sort"
 	"strconv"
 	"strings"
@@ -81,13 +80,8 @@ func ParseNodeMetrics(input []byte) map[string]*NodeMetrics {
 // NodeData executes the sinfo command to get data for each node
 // It returns the output of the sinfo command
 func NodeData(logger log.Logger) ([]byte, error) {
-	cmd := exec.Command("sinfo", "-h", "-N", "-O", "NodeList,AllocMem,Memory,CPUsState,StateLong,Partition")
-	out, err := cmd.Output()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to execute sinfo command", "err", err)
-		return nil, err
-	}
-	return out, nil
+	args := []string{"-h", "-N", "-O", "NodeList,AllocMem,Memory,CPUsState,StateLong,Partition"}
+	return Execute(logger, "sinfo", args)
 }
 
 type NodeCollector struct {

collector/nodes.go

@@ -1,8 +1,6 @@
 package collector
 
 import (
-	"io"
-	"os/exec"
 	"regexp"
 	"sort"
 	"strconv"
@@ -135,69 +133,41 @@ func ParseNodesMetrics(input []byte) *NodesMetrics {
 
 // Execute the sinfo command and return its output
 func NodesData(logger log.Logger, part string) ([]byte, error) {
-	cmd := exec.Command("sinfo", "-h", "-o %D|%T|%b", "-p", part, "| sort", "| uniq")
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return nil, err
-	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return nil, err
-	}
-	out, _ := io.ReadAll(stdout)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err)
-		return nil, err
-	}
-	return out, nil
+	return Execute(logger, "sinfo", []string{"-h", "-o", "%D|%T|%b", "-p", part})
 }
 
 func SlurmGetTotal(logger log.Logger) (float64, error) {
-	cmd := exec.Command("bash", "-c", "scontrol show nodes -o | grep -c NodeName=[a-z]*[0-9]*")
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return 0, err
-	}
-	stderr, err := cmd.StderrPipe()
+	out, err := Execute(logger, "scontrol", []string{"show", "nodes", "-o"})
 	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stderr pipe", "err", err)
 		return 0, err
 	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return 0, err
-	}
-	out, _ := io.ReadAll(stdout)
-	err_out, _ := io.ReadAll(stderr)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err, "stdout", string(out), "stderr", string(err_out))
-		return 0, err
+	// Filter out empty lines before counting
+	lines := strings.Split(string(out), "\n")
+	count := 0
+	for _, line := range lines {
+		if strings.TrimSpace(line) != "" {
+			count++
+		}
 	}
-	data := strings.Split(string(out), "\n")
-	total, _ := strconv.ParseFloat(data[0], 64)
-	return total, nil
+	return float64(count), nil
 }
 
 func SlurmGetPartitions(logger log.Logger) ([]string, error) {
-	cmd := exec.Command("sinfo", "-h", "-o %R", "| sort", "| uniq")
-	stdout, err := cmd.StdoutPipe()
+	out, err := Execute(logger, "sinfo", []string{"-h", "-o", "%R"})
 	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return nil, err
-	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return nil, err
-	}
-	out, _ := io.ReadAll(stdout)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err)
 		return nil, err
 	}
 	partitions := strings.Split(string(out), "\n")
-	return partitions, nil
+	// Trim whitespace and remove empty strings
+	var cleanedPartitions []string
+	for _, p := range partitions {
+		p = strings.TrimSpace(p)
+		if p != "" {
+			cleanedPartitions = append(cleanedPartitions, p)
+		}
+	}
+	sort.Strings(cleanedPartitions)
+	return RemoveDuplicates(cleanedPartitions), nil
 }
 
 /*

collector/partitions.go

@@ -1,8 +1,6 @@
 package collector
 
 import (
-	"io"
-	"os/exec"
 	"strconv"
 	"strings"
 
@@ -12,41 +10,11 @@ import (
 )
 
 func PartitionsData(logger log.Logger) ([]byte, error) {
-	cmd := exec.Command("sinfo", "-h", "-o%R,%C")
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return nil, err
-	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return nil, err
-	}
-	out, _ := io.ReadAll(stdout)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err)
-		return nil, err
-	}
-	return out, nil
+	return Execute(logger, "sinfo", []string{"-h", "-o", "%R,%C"})
 }
 
 func PartitionsPendingJobsData(logger log.Logger) ([]byte, error) {
-	cmd := exec.Command("squeue", "-a", "-r", "-h", "-o%P", "--states=PENDING")
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return nil, err
-	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return nil, err
-	}
-	out, _ := io.ReadAll(stdout)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err)
-		return nil, err
-	}
-	return out, nil
+	return Execute(logger, "squeue", []string{"-a", "-r", "-h", "-o", "%P", "--states=PENDING"})
 }
 
 type PartitionMetrics struct {

collector/queue.go

@@ -1,8 +1,6 @@
 package collector
 
 import (
-	"io"
-	"os/exec"
 	"strconv"
 	"strings"
 
@@ -150,22 +148,7 @@ func ParseQueueMetrics(input []byte) *QueueMetrics {
 
 // Execute the squeue command and return its output
 func QueueData(logger log.Logger) ([]byte, error) {
-	cmd := exec.Command("/usr/bin/squeue", "-h", "-o %P,%T,%C,%r,%u")
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		level.Error(logger).Log("msg", "Failed to create stdout pipe", "err", err)
-		return nil, err
-	}
-	if err := cmd.Start(); err != nil {
-		level.Error(logger).Log("msg", "Failed to start command", "err", err)
-		return nil, err
-	}
-	out, _ := io.ReadAll(stdout)
-	if err := cmd.Wait(); err != nil {
-		level.Error(logger).Log("msg", "Failed to wait for command", "err", err)
-		return nil, err
-	}
-	return out, nil
+	return Execute(logger, "squeue", []string{"-h", "-o", "%P,%T,%C,%r,%u"})
 }
 
 /*