extras/hs-test/cpu.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "fmt"
   6         "os"
   7 )
   8
   9 var CPU_PATH = "/sys/fs/cgroup/cpuset.cpus.effective"
  10
  11 type CpuContext struct {
  12         cpuAllocator *CpuAllocatorT
  13         cpus         []int
  14 }
  15
  16 func (c *CpuContext) Release() {
  17         c.cpuAllocator.cpus = append(c.cpuAllocator.cpus, c.cpus...)
  18         c.cpus = c.cpus[:0] // empty the list
  19 }
  20
  21 type CpuAllocatorT struct {
  22         cpus []int
  23 }
  24
  25 var cpuAllocator *CpuAllocatorT = nil
  26
  27 func (c *CpuAllocatorT) Allocate(nCpus int) (*CpuContext, error) {
  28         var cpuCtx CpuContext
  29
  30         if len(c.cpus) < nCpus {
  31                 return nil, fmt.Errorf("could not allocate %d CPUs; available: %d", nCpus, len(c.cpus))
  32         }
  33         cpuCtx.cpus = c.cpus[0:nCpus]
  34         cpuCtx.cpuAllocator = c
  35         c.cpus = c.cpus[nCpus:]
  36         return &cpuCtx, nil
  37 }
  38
  39 func (c *CpuAllocatorT) readCpus(fname string) error {
  40         var first, last int
  41         file, err := os.Open(CPU_PATH)
  42         if err != nil {
  43                 return err
  44         }
  45         defer file.Close()
  46
  47         sc := bufio.NewScanner(file)
  48         sc.Scan()
  49         line := sc.Text()
  50         _, err = fmt.Sscanf(line, "%d-%d", &first, &last)
  51         if err != nil {
  52                 return err
  53         }
  54         for i := first; i <= last; i++ {
  55                 c.cpus = append(c.cpus, i)
  56         }
  57         return nil
  58 }
  59
  60 func CpuAllocator() (*CpuAllocatorT, error) {
  61         if cpuAllocator == nil {
  62                 cpuAllocator = new(CpuAllocatorT)
  63                 err := cpuAllocator.readCpus(CPU_PATH)
  64                 if err != nil {
  65                         return nil, err
  66                 }
  67         }
  68         return cpuAllocator, nil
  69 }