numalign: add command line parameters

We are starting to have too many env vars; is not practical to control the behaviour ONLY using them. We want to keep them because they are good if used sparingly, and for backward compatibility. Signed-off-by: Francesco Romani <[email protected]>
ffromani · Sep 24, 2020 · 801c86b · 801c86b
1 parent 7ce3477
commit 801c86b
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 19 deletions.
diff --git a/cmd/numalign/main.go b/cmd/numalign/main.go
@@ -23,37 +23,52 @@ import (
 	"strconv"
 	"time"
 
+	flag "github.com/spf13/pflag"
+
 	"github.com/fromanirh/numalign/internal/pkg/numalign"
 )
 
 func main() {
-	var val string
 	var sleepTime time.Duration
 
+	var sleepHoursParam = flag.StringP("sleep-hours", "S", "", "sleep hours once done.")
+	var scriptPathParam = flag.StringP("script-path", "P", "", "save test script to this path.")
+	flag.Parse()
+
 	if _, ok := os.LookupEnv("NUMALIGN_DEBUG"); !ok {
 		log.SetOutput(ioutil.Discard)
 	}
 
-	val = os.Getenv("NUMALIGN_SLEEP_HOURS")
-	if val != "" {
-		hours, err := strconv.Atoi(val)
+	sleepHours := *sleepHoursParam
+	if sleepHours == "" {
+		os.Getenv("NUMALIGN_SLEEP_HOURS")
+	}
+	if sleepHours != "" {
+		hours, err := strconv.Atoi(sleepHours)
 		if err != nil {
 			log.Fatalf("%v", err)
 		}
-		sleepTime = time.Duration(hours) * time.Hour
+		if hours > 0 {
+			sleepTime = time.Duration(hours) * time.Hour
+		}
 	}
 
 	log.Printf("SYS: sleep for %v after the check", sleepTime)
 
-	R, err := numalign.NewResources()
+	R, err := numalign.NewResources(flag.Args())
 	if err != nil {
 		log.Fatalf("%v", err)
 	}
 	ret := numalign.Validate(R)
 
-	if val = os.Getenv("NUMALIGN_VALIDATION_SCRIPT"); val != "" {
+	scriptPath := *scriptPathParam
+	if scriptPath == "" {
+		scriptPath = os.Getenv("NUMALIGN_VALIDATION_SCRIPT")
+	}
+
+	if scriptPath != "" {
 		code := []byte(R.MakeValidationScript())
-		err := ioutil.WriteFile(val, code, 0755)
+		err := ioutil.WriteFile(scriptPath, code, 0755)
 		if err != nil {
 			log.Printf("SYS: validation script creation failed: %v", err)
 		}

diff --git a/internal/pkg/numalign/numalign.go b/internal/pkg/numalign/numalign.go
@@ -227,7 +227,7 @@ func GetCPUNUMANodes(cpusPerNUMA map[int][]int) map[int]int {
 	return CPUToNUMANode
 }
 
-func NewResources() (*Resources, error) {
+func NewResources(pids []string) (*Resources, error) {
 	var err error
 
 	cpuRes, err := cpus.NewCPUs("/sys")
@@ -250,8 +250,8 @@ func NewResources() (*Resources, error) {
 	}
 
 	var pidStrings []string
-	if len(os.Args) > 1 {
-		pidStrings = append(pidStrings, os.Args[1:]...)
+	if len(pids) > 1 {
+		pidStrings = append(pidStrings, pids...)
 	} else {
 		pidStrings = append(pidStrings, "self")
 	}
@@ -303,11 +303,3 @@ func Validate(R *Resources) int {
 	fmt.Printf("NUMA NODE=%v\n", nodeNum)
 	return 0
 }
-
-func Execute() int {
-	R, err := NewResources()
-	if err != nil {
-		log.Fatalf("%v", err)
-	}
-	return Validate(R)
-}