trace logic

exporter/trace.go

@@ -8,7 +8,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/fs"
 	"log"
 	"net/http"
 	"os"
@@ -109,22 +108,11 @@ type TraceCollector struct {
 
 func NewTraceCollector(config *Config) *TraceCollector {
 	traceConfig := config.TraceConf
-	templateRootDir := "."
-	// path to look for the /templates directory. Defaults to cwd
-	if path, ok := os.LookupEnv("TRACE_ROOT_PATH"); ok {
-		templateRootDir = path
-	}
-	traceDir := ""
-	err := filepath.WalkDir(templateRootDir, func(path string, d fs.DirEntry, err error) error {
-		if err == nil && d.IsDir() && d.Name() == templateDirName {
-			traceDir = path
-			return nil
-		}
-		return nil
-	})
-	if err != nil || traceDir == "" {
+	traceDir := detectTraceTemplatePath()
+	if traceDir == "" {
 		log.Fatal("no template found")
 	}
+	slog.Debug("using trace template path: " + traceDir)
 	return &TraceCollector{
 		ProcessFetcher: NewAtomicProFetcher(traceConfig.rate),
 		squeueFetcher:  traceConfig.sharedFetcher,
@@ -203,21 +191,22 @@ func (c *TraceCollector) uploadTrace(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-// detectTracePath returns the trace_root path based on the following criteria:
+// detectTraceTemplatePath returns the trace_root path based on the following criteria:
 // 1. If TRACE_ROOT_PATH is specified, search that directory. If we don't find a templates dir, let's panic and crash the program.
 // 2. If TRACE_ROOT_PATH isn't specified, we can search cwd and /usr/share/prometheus-slurm-exporter.
-func detectTracePath() string {
-	templateRootDir := ""
-	if path, ok := os.LookupEnv("TRACE_ROOT_PATH"); ok {
-		templateRootDir = path
-		if _, err := os.Stat(filepath.Join(templateRootDir, templateDirName)); err != nil {
-			panic("TRACE_ROOT_PATH must include a directory called: " + templateDirName)
+// If no templates path is found, returns an empty string
+func detectTraceTemplatePath() string {
+	if rpath, ok := os.LookupEnv("TRACE_ROOT_PATH"); ok {
+		templateP := filepath.Join(rpath, templateDirName)
+		if _, err := os.Stat(templateP); err != nil {
+			panic("TRACE_ROOT_PATH must include a directory called: templates")
 		}
-		return templateRootDir
+		return templateP
 	}
-	for _, p := range []string{".", "/usr/share/prometheus-slurm-exporter"} {
-		if _, err := os.Stat(filepath.Join(p, templateDirName)); err == nil {
-			return p
+	for _, rpath := range []string{".", "/usr/share/prometheus-slurm-exporter"} {
+		templateP := filepath.Join(rpath, templateDirName)
+		if _, err := os.Stat(templateP); err == nil {
+			return templateP
 		}
 	}
 	return ""

exporter/trace_test.go

@@ -211,11 +211,11 @@ func TestDetectTraceRootPath_Env(t *testing.T) {
 	testDir := t.TempDir()
 	t.Setenv("TRACE_ROOT_PATH", testDir)
 	// Ensure that the function panics if given a TRACE_ROOT_PATh with no 'templates' subdirectory
-	assert.PanicsWithValue(t, "TRACE_ROOT_PATH must include a directory called: templates", func() { detectTracePath() })
+	assert.PanicsWithValue(t, "TRACE_ROOT_PATH must include a directory called: templates", func() { detectTraceTemplatePath() })
 	require.NoError(t, os.Mkdir(filepath.Join(testDir, templateDirName), 0o700))
 
 	// Now that we have a 'templates' subdir, it should no longer panic
-	assert.Equal(t, testDir, detectTracePath())
+	assert.Equal(t, filepath.Join(testDir, templateDirName), detectTraceTemplatePath())
 }
 
 func TestDetectTraceRootPath_Default(t *testing.T) {
@@ -224,9 +224,9 @@ func TestDetectTraceRootPath_Default(t *testing.T) {
 	os.Chdir(testDir)
 
 	// Should come back empty if since we don't yet have a 'templates' subdir
-	assert.Equal(t, detectTracePath(), "")
+	assert.Equal(t, detectTraceTemplatePath(), "")
 	require.NoError(t, os.Mkdir(filepath.Join(testDir, templateDirName), 0o700))
 
 	// Now that we have 'templates' subdir, cwd is a valid path
-	assert.Equal(t, detectTracePath(), ".")
+	assert.Equal(t, templateDirName, detectTraceTemplatePath())
 }