extras/hs-test/cpu.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "errors"
   6         "fmt"
   7         "os"
   8         "os/exec"
   9         "strings"
  10 )
  11
  12 var CgroupPath = "/sys/fs/cgroup/"
  13
  14 type CpuContext struct {
  15         cpuAllocator *CpuAllocatorT
  16         cpus         []int
  17 }
  18
  19 func (c *CpuContext) Release() {
  20         c.cpuAllocator.cpus = append(c.cpuAllocator.cpus, c.cpus...)
  21         c.cpus = c.cpus[:0] // empty the list
  22 }
  23
  24 type CpuAllocatorT struct {
  25         cpus []int
  26 }
  27
  28 var cpuAllocator *CpuAllocatorT = nil
  29
  30 func (c *CpuAllocatorT) Allocate(nCpus int) (*CpuContext, error) {
  31         var cpuCtx CpuContext
  32
  33         if len(c.cpus) < nCpus {
  34                 return nil, fmt.Errorf("could not allocate %d CPUs; available: %d", nCpus, len(c.cpus))
  35         }
  36         cpuCtx.cpus = c.cpus[0:nCpus]
  37         cpuCtx.cpuAllocator = c
  38         c.cpus = c.cpus[nCpus:]
  39         return &cpuCtx, nil
  40 }
  41
  42 func (c *CpuAllocatorT) readCpus() error {
  43         var first, last int
  44
  45         // Path depends on cgroup version. We need to check which version is in use.
  46         // For that following command can be used: 'stat -fc %T /sys/fs/cgroup/'
  47         // In case the output states 'cgroup2fs' then cgroups v2 is used, 'tmpfs' in case cgroups v1.
  48         cmd := exec.Command("stat", "-fc", "%T", "/sys/fs/cgroup/")
  49         byteOutput, err := cmd.CombinedOutput()
  50         if err != nil {
  51                 return err
  52         }
  53         CpuPath := CgroupPath
  54         if strings.Contains(string(byteOutput), "tmpfs") {
  55                 CpuPath += "cpuset/cpuset.effective_cpus"
  56         } else if strings.Contains(string(byteOutput), "cgroup2fs") {
  57                 CpuPath += "cpuset.cpus.effective"
  58         } else {
  59                 return errors.New("cgroup unknown fs: " + string(byteOutput))
  60         }
  61
  62         file, err := os.Open(CpuPath)
  63         if err != nil {
  64                 return err
  65         }
  66         defer file.Close()
  67
  68         sc := bufio.NewScanner(file)
  69         sc.Scan()
  70         line := sc.Text()
  71         _, err = fmt.Sscanf(line, "%d-%d", &first, &last)
  72         if err != nil {
  73                 return err
  74         }
  75         for i := first; i <= last; i++ {
  76                 c.cpus = append(c.cpus, i)
  77         }
  78         return nil
  79 }
  80
  81 func CpuAllocator() (*CpuAllocatorT, error) {
  82         if cpuAllocator == nil {
  83                 cpuAllocator = new(CpuAllocatorT)
  84                 err := cpuAllocator.readCpus()
  85                 if err != nil {
  86                         return nil, err
  87                 }
  88         }
  89         return cpuAllocator, nil
  90 }