diff --git a/numcpus_linux.go b/numcpus_linux.go index 7b991da..d05ee98 100644 --- a/numcpus_linux.go +++ b/numcpus_linux.go @@ -47,10 +47,12 @@ func readCPURangeWith[T any](file string, f func(cpus string) (T, error)) (T, er if err != nil { return zero, err } - return f(strings.Trim(string(buf), "\n ")) + return f(string(buf)) } func countCPURange(cpus string) (int, error) { + cpus = strings.Trim(cpus, "\n ") + // Treat empty file as valid. This might be the case if there are no offline CPUs in which // case /sys/devices/system/cpu/offline is empty. if cpus == "" { @@ -58,7 +60,7 @@ func countCPURange(cpus string) (int, error) { } n := int(0) - for _, cpuRange := range strings.Split(cpus, ",") { + for cpuRange := range strings.SplitSeq(cpus, ",") { if cpuRange == "" { return 0, fmt.Errorf("empty CPU range in CPU string %q", cpus) } @@ -84,13 +86,15 @@ func countCPURange(cpus string) (int, error) { } func listCPURange(cpus string) ([]int, error) { + cpus = strings.Trim(cpus, "\n ") + // See comment in countCPURange. if cpus == "" { return []int{}, nil } list := []int{} - for _, cpuRange := range strings.Split(cpus, ",") { + for cpuRange := range strings.SplitSeq(cpus, ",") { if cpuRange == "" { return nil, fmt.Errorf("empty CPU range in CPU string %q", cpus) } diff --git a/numcpus_linux_test.go b/numcpus_linux_test.go index f7cc9cd..f6f757c 100644 --- a/numcpus_linux_test.go +++ b/numcpus_linux_test.go @@ -31,6 +31,11 @@ func TestCPURange(t *testing.T) { wantCount: 0, wantList: []int{}, }, + { + str: "\n", + wantCount: 0, + wantList: []int{}, + }, { str: "0", wantCount: 1, @@ -86,6 +91,11 @@ func TestCPURange(t *testing.T) { wantCount: 9, wantList: []int{0, 2, 3, 4, 6, 8, 9, 10, 31}, }, + { + str: " 0,2-7\n", + wantCount: 7, + wantList: []int{0, 2, 3, 4, 5, 6, 7}, + }, { str: "invalid", wantErr: true, @@ -130,6 +140,10 @@ func TestCPURange(t *testing.T) { str: "0,5-3", wantErr: true, }, + { + str: "0, 5 - 3", + wantErr: true, + }, } for _, tc := range testCases {