1
2
3
4
5 package main
6
7 import (
8 "bytes"
9 "fmt"
10 "os"
11 "os/exec"
12 "regexp"
13 "runtime"
14 "strconv"
15 "strings"
16 "syscall"
17 )
18
19 var (
20 cpuSetRE = regexp.MustCompile(`(\d,?)+`)
21 )
22
23 func init() {
24 register("FreeBSDNumCPU", FreeBSDNumCPU)
25 register("FreeBSDNumCPUHelper", FreeBSDNumCPUHelper)
26 }
27
28 func FreeBSDNumCPUHelper() {
29 fmt.Printf("%d\n", runtime.NumCPU())
30 }
31
32 func FreeBSDNumCPU() {
33 _, err := exec.LookPath("cpuset")
34 if err != nil {
35
36 fmt.Println("OK")
37 return
38 }
39 _, err = exec.LookPath("sysctl")
40 if err != nil {
41
42 fmt.Println("OK")
43 return
44 }
45 cmd := exec.Command("sysctl", "-n", "kern.smp.active")
46 output, err := cmd.CombinedOutput()
47 if err != nil {
48 fmt.Printf("fail to launch '%s', error: %s, output: %s\n", strings.Join(cmd.Args, " "), err, output)
49 return
50 }
51 if bytes.Equal(output, []byte("1\n")) == false {
52
53 fmt.Println("OK")
54 return
55 }
56
57 list, err := getList()
58 if err != nil {
59 fmt.Printf("%s\n", err)
60 return
61 }
62 err = checkNCPU(list)
63 if err != nil {
64 fmt.Printf("%s\n", err)
65 return
66 }
67 if len(list) >= 2 {
68 err = checkNCPU(list[:len(list)-1])
69 if err != nil {
70 fmt.Printf("%s\n", err)
71 return
72 }
73 }
74 fmt.Println("OK")
75 return
76 }
77
78 func getList() ([]string, error) {
79 pid := syscall.Getpid()
80
81
82 cmd := exec.Command("cpuset", "-g", "-p", strconv.Itoa(pid))
83 cmdline := strings.Join(cmd.Args, " ")
84 output, err := cmd.CombinedOutput()
85 if err != nil {
86 return nil, fmt.Errorf("fail to execute '%s': %s", cmdline, err)
87 }
88 pos := bytes.IndexRune(output, '\n')
89 if pos == -1 {
90 return nil, fmt.Errorf("invalid output from '%s', '\\n' not found: %s", cmdline, output)
91 }
92 output = output[0:pos]
93
94 pos = bytes.IndexRune(output, ':')
95 if pos == -1 {
96 return nil, fmt.Errorf("invalid output from '%s', ':' not found: %s", cmdline, output)
97 }
98
99 var list []string
100 for _, val := range bytes.Split(output[pos+1:], []byte(",")) {
101 index := string(bytes.TrimSpace(val))
102 if len(index) == 0 {
103 continue
104 }
105 list = append(list, index)
106 }
107 if len(list) == 0 {
108 return nil, fmt.Errorf("empty CPU list from '%s': %s", cmdline, output)
109 }
110 return list, nil
111 }
112
113 func checkNCPU(list []string) error {
114 listString := strings.Join(list, ",")
115 if len(listString) == 0 {
116 return fmt.Errorf("could not check against an empty CPU list")
117 }
118
119 cListString := cpuSetRE.FindString(listString)
120 if len(cListString) == 0 {
121 return fmt.Errorf("invalid cpuset output '%s'", listString)
122 }
123
124 cmd := exec.Command("cpuset", "-l", cListString, os.Args[0], "FreeBSDNumCPUHelper")
125 cmdline := strings.Join(cmd.Args, " ")
126 output, err := cmd.CombinedOutput()
127 if err != nil {
128 return fmt.Errorf("fail to launch child '%s', error: %s, output: %s", cmdline, err, output)
129 }
130
131
132 output = bytes.TrimSpace(output)
133 n, err := strconv.Atoi(string(output))
134 if err != nil {
135 return fmt.Errorf("fail to parse output from child '%s', error: %s, output: %s", cmdline, err, output)
136 }
137 if n != len(list) {
138 return fmt.Errorf("runtime.NumCPU() expected to %d, got %d when run with CPU list %s", len(list), n, cListString)
139 }
140 return nil
141 }
142
View as plain text