diff --git a/lib/cpuset/cpuset.go b/lib/cpuset/cpuset.go index 4f03d3f74..7058e46df 100644 --- a/lib/cpuset/cpuset.go +++ b/lib/cpuset/cpuset.go @@ -27,6 +27,38 @@ func New(cpus ...uint16) CPUSet { return cpuset } +// String returns the cpuset as a comma delimited set of core values and ranged +func (c CPUSet) String() string { + if c.Size() == 0 { + return "" + } + cores := c.ToSlice() + cpusetStrs := []string{} + cur := [2]uint16{cores[0], cores[0]} + for i := 1; i < len(cores); i++ { + if cores[i] == cur[1]+1 { + cur[1] = cores[i] + continue + } + + if cur[0] == cur[1] { + cpusetStrs = append(cpusetStrs, fmt.Sprintf("%d", cur[0])) + } else { + cpusetStrs = append(cpusetStrs, fmt.Sprintf("%d-%d", cur[0], cur[1])) + } + + // new range + cur = [2]uint16{cores[i], cores[i]} + } + if cur[0] == cur[1] { + cpusetStrs = append(cpusetStrs, fmt.Sprintf("%d", cur[0])) + } else { + cpusetStrs = append(cpusetStrs, fmt.Sprintf("%d-%d", cur[0], cur[1])) + } + + return strings.Join(cpusetStrs, ",") +} + // Size returns to the number of cpus contained in the CPUSet func (c CPUSet) Size() int { return len(c.cpus) @@ -88,6 +120,16 @@ func (s CPUSet) IsSupersetOf(other CPUSet) bool { return true } +// ContainsAny returns true if any cpus in other CPUSet are present +func (s CPUSet) ContainsAny(other CPUSet) bool { + for cpu := range other.cpus { + if _, ok := s.cpus[cpu]; ok { + return true + } + } + return false +} + // Equals tests the equality of the elements in the CPUSet func (s CPUSet) Equals(other CPUSet) bool { return reflect.DeepEqual(s.cpus, other.cpus) @@ -98,6 +140,7 @@ func (s CPUSet) Equals(other CPUSet) bool { // Ref: http://man7.org/linux/man-pages/man7/cpuset.7.html#FORMATS func Parse(s string) (CPUSet, error) { cpuset := New() + s = strings.TrimSpace(s) if s == "" { return cpuset, nil } diff --git a/lib/cpuset/cpuset_test.go b/lib/cpuset/cpuset_test.go index bd1bc1f7d..ea7e72108 100644 --- a/lib/cpuset/cpuset_test.go +++ b/lib/cpuset/cpuset_test.go @@ -146,13 +146,34 @@ func TestCPUSet_IsSupersetOf(t *testing.T) { } } +func TestCPUSet_ContainsAny(t *testing.T) { + cases := []struct { + a CPUSet + b CPUSet + containsAny bool + }{ + {New(0), New(0), true}, + {New(0), New(), false}, + {New(), New(0), false}, + {New(0, 1, 2, 3), New(0), true}, + {New(0, 1, 2, 3), New(2, 3), true}, + {New(0, 1, 2, 3), New(2, 3, 4), true}, + } + + for _, c := range cases { + require.Equal(t, c.containsAny, c.a.ContainsAny(c.b)) + } +} + func TestParse(t *testing.T) { cases := []struct { cpuset string expected CPUSet }{ {"", New()}, + {"\n", New()}, {"1", New(1)}, + {"1\n", New(1)}, {"0,1,2,3", New(0, 1, 2, 3)}, {"0-3", New(0, 1, 2, 3)}, {"0,2-3,5", New(0, 2, 3, 5)}, @@ -164,3 +185,19 @@ func TestParse(t *testing.T) { require.True(t, result.Equals(c.expected)) } } + +func TestCPUSet_String(t *testing.T) { + cases := []struct { + cpuset CPUSet + expected string + }{ + {New(), ""}, + {New(0, 1, 2, 3), "0-3"}, + {New(1, 3), "1,3"}, + {New(0, 2, 3, 5), "0,2-3,5"}, + } + + for _, c := range cases { + require.Equal(t, c.expected, c.cpuset.String()) + } +}