cpuset.go

Documentation: k8s.io/utils/cpuset

     1  /*
     2  Copyright 2017 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package cpuset represents a collection of CPUs in a 'set' data structure.
    18  //
    19  // It can be used to represent core IDs, hyper thread siblings, CPU nodes, or processor IDs.
    20  //
    21  // The only special thing about this package is that
    22  // methods are provided to convert back and forth from Linux 'list' syntax.
    23  // See http://man7.org/linux/man-pages/man7/cpuset.7.html#FORMATS for details.
    24  //
    25  // Future work can migrate this to use a 'set' library, and relax the dubious 'immutable' property.
    26  //
    27  // This package was originally developed in the 'kubernetes' repository.
    28  package cpuset
    29  
    30  import (
    31  	"bytes"
    32  	"fmt"
    33  	"reflect"
    34  	"sort"
    35  	"strconv"
    36  	"strings"
    37  )
    38  
    39  // CPUSet is a thread-safe, immutable set-like data structure for CPU IDs.
    40  type CPUSet struct {
    41  	elems map[int]struct{}
    42  }
    43  
    44  // New returns a new CPUSet containing the supplied elements.
    45  func New(cpus ...int) CPUSet {
    46  	s := CPUSet{
    47  		elems: map[int]struct{}{},
    48  	}
    49  	for _, c := range cpus {
    50  		s.add(c)
    51  	}
    52  	return s
    53  }
    54  
    55  // add adds the supplied elements to the CPUSet.
    56  // It is intended for internal use only, since it mutates the CPUSet.
    57  func (s CPUSet) add(elems ...int) {
    58  	for _, elem := range elems {
    59  		s.elems[elem] = struct{}{}
    60  	}
    61  }
    62  
    63  // Size returns the number of elements in this set.
    64  func (s CPUSet) Size() int {
    65  	return len(s.elems)
    66  }
    67  
    68  // IsEmpty returns true if there are zero elements in this set.
    69  func (s CPUSet) IsEmpty() bool {
    70  	return s.Size() == 0
    71  }
    72  
    73  // Contains returns true if the supplied element is present in this set.
    74  func (s CPUSet) Contains(cpu int) bool {
    75  	_, found := s.elems[cpu]
    76  	return found
    77  }
    78  
    79  // Equals returns true if the supplied set contains exactly the same elements
    80  // as this set (s IsSubsetOf s2 and s2 IsSubsetOf s).
    81  func (s CPUSet) Equals(s2 CPUSet) bool {
    82  	return reflect.DeepEqual(s.elems, s2.elems)
    83  }
    84  
    85  // filter returns a new CPU set that contains all of the elements from this
    86  // set that match the supplied predicate, without mutating the source set.
    87  func (s CPUSet) filter(predicate func(int) bool) CPUSet {
    88  	r := New()
    89  	for cpu := range s.elems {
    90  		if predicate(cpu) {
    91  			r.add(cpu)
    92  		}
    93  	}
    94  	return r
    95  }
    96  
    97  // IsSubsetOf returns true if the supplied set contains all the elements
    98  func (s CPUSet) IsSubsetOf(s2 CPUSet) bool {
    99  	result := true
   100  	for cpu := range s.elems {
   101  		if !s2.Contains(cpu) {
   102  			result = false
   103  			break
   104  		}
   105  	}
   106  	return result
   107  }
   108  
   109  // Union returns a new CPU set that contains all of the elements from this
   110  // set and all of the elements from the supplied sets, without mutating
   111  // either source set.
   112  func (s CPUSet) Union(s2 ...CPUSet) CPUSet {
   113  	r := New()
   114  	for cpu := range s.elems {
   115  		r.add(cpu)
   116  	}
   117  	for _, cs := range s2 {
   118  		for cpu := range cs.elems {
   119  			r.add(cpu)
   120  		}
   121  	}
   122  	return r
   123  }
   124  
   125  // Intersection returns a new CPU set that contains all of the elements
   126  // that are present in both this set and the supplied set, without mutating
   127  // either source set.
   128  func (s CPUSet) Intersection(s2 CPUSet) CPUSet {
   129  	return s.filter(func(cpu int) bool { return s2.Contains(cpu) })
   130  }
   131  
   132  // Difference returns a new CPU set that contains all of the elements that
   133  // are present in this set and not the supplied set, without mutating either
   134  // source set.
   135  func (s CPUSet) Difference(s2 CPUSet) CPUSet {
   136  	return s.filter(func(cpu int) bool { return !s2.Contains(cpu) })
   137  }
   138  
   139  // List returns a slice of integers that contains all elements from
   140  // this set. The list is sorted.
   141  func (s CPUSet) List() []int {
   142  	result := s.UnsortedList()
   143  	sort.Ints(result)
   144  	return result
   145  }
   146  
   147  // UnsortedList returns a slice of integers that contains all elements from
   148  // this set.
   149  func (s CPUSet) UnsortedList() []int {
   150  	result := make([]int, 0, len(s.elems))
   151  	for cpu := range s.elems {
   152  		result = append(result, cpu)
   153  	}
   154  	return result
   155  }
   156  
   157  // String returns a new string representation of the elements in this CPU set
   158  // in canonical linux CPU list format.
   159  //
   160  // See: http://man7.org/linux/man-pages/man7/cpuset.7.html#FORMATS
   161  func (s CPUSet) String() string {
   162  	if s.IsEmpty() {
   163  		return ""
   164  	}
   165  
   166  	elems := s.List()
   167  
   168  	type rng struct {
   169  		start int
   170  		end   int
   171  	}
   172  
   173  	ranges := []rng{{elems[0], elems[0]}}
   174  
   175  	for i := 1; i < len(elems); i++ {
   176  		lastRange := &ranges[len(ranges)-1]
   177  		// if this element is adjacent to the high end of the last range
   178  		if elems[i] == lastRange.end+1 {
   179  			// then extend the last range to include this element
   180  			lastRange.end = elems[i]
   181  			continue
   182  		}
   183  		// otherwise, start a new range beginning with this element
   184  		ranges = append(ranges, rng{elems[i], elems[i]})
   185  	}
   186  
   187  	// construct string from ranges
   188  	var result bytes.Buffer
   189  	for _, r := range ranges {
   190  		if r.start == r.end {
   191  			result.WriteString(strconv.Itoa(r.start))
   192  		} else {
   193  			result.WriteString(fmt.Sprintf("%d-%d", r.start, r.end))
   194  		}
   195  		result.WriteString(",")
   196  	}
   197  	return strings.TrimRight(result.String(), ",")
   198  }
   199  
   200  // Parse CPUSet constructs a new CPU set from a Linux CPU list formatted string.
   201  //
   202  // See: http://man7.org/linux/man-pages/man7/cpuset.7.html#FORMATS
   203  func Parse(s string) (CPUSet, error) {
   204  	// Handle empty string.
   205  	if s == "" {
   206  		return New(), nil
   207  	}
   208  
   209  	result := New()
   210  
   211  	// Split CPU list string:
   212  	// "0-5,34,46-48" => ["0-5", "34", "46-48"]
   213  	ranges := strings.Split(s, ",")
   214  
   215  	for _, r := range ranges {
   216  		boundaries := strings.SplitN(r, "-", 2)
   217  		if len(boundaries) == 1 {
   218  			// Handle ranges that consist of only one element like "34".
   219  			elem, err := strconv.Atoi(boundaries[0])
   220  			if err != nil {
   221  				return New(), err
   222  			}
   223  			result.add(elem)
   224  		} else if len(boundaries) == 2 {
   225  			// Handle multi-element ranges like "0-5".
   226  			start, err := strconv.Atoi(boundaries[0])
   227  			if err != nil {
   228  				return New(), err
   229  			}
   230  			end, err := strconv.Atoi(boundaries[1])
   231  			if err != nil {
   232  				return New(), err
   233  			}
   234  			if start > end {
   235  				return New(), fmt.Errorf("invalid range %q (%d > %d)", r, start, end)
   236  			}
   237  			// start == end is acceptable (1-1 -> 1)
   238  
   239  			// Add all elements to the result.
   240  			// e.g. "0-5", "46-48" => [0, 1, 2, 3, 4, 5, 46, 47, 48].
   241  			for e := start; e <= end; e++ {
   242  				result.add(e)
   243  			}
   244  		}
   245  	}
   246  	return result, nil
   247  }
   248  
   249  // Clone returns a copy of this CPU set.
   250  func (s CPUSet) Clone() CPUSet {
   251  	r := New()
   252  	for elem := range s.elems {
   253  		r.add(elem)
   254  	}
   255  	return r
   256  }
   257
View as plain text