1// Copyright 2017 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package main
6
7import (
8	"bytes"
9	"fmt"
10	"os"
11	"os/exec"
12	"regexp"
13	"runtime"
14	"strconv"
15	"strings"
16	"syscall"
17)
18
19var (
20	cpuSetRE = regexp.MustCompile(`(\d,?)+`)
21)
22
23func init() {
24	register("FreeBSDNumCPU", FreeBSDNumCPU)
25	register("FreeBSDNumCPUHelper", FreeBSDNumCPUHelper)
26}
27
28func FreeBSDNumCPUHelper() {
29	fmt.Printf("%d\n", runtime.NumCPU())
30}
31
32func FreeBSDNumCPU() {
33	_, err := exec.LookPath("cpuset")
34	if err != nil {
35		// Can not test without cpuset command.
36		fmt.Println("OK")
37		return
38	}
39	_, err = exec.LookPath("sysctl")
40	if err != nil {
41		// Can not test without sysctl command.
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")) {
52		// SMP mode deactivated in kernel.
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
78func getList() ([]string, error) {
79	pid := syscall.Getpid()
80
81	// Launch cpuset to print a list of available CPUs: pid <PID> mask: 0, 1, 2, 3.
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	output, _, ok := bytes.Cut(output, []byte("\n"))
89	if !ok {
90		return nil, fmt.Errorf("invalid output from '%s', '\\n' not found: %s", cmdline, output)
91	}
92
93	_, cpus, ok := bytes.Cut(output, []byte(":"))
94	if !ok {
95		return nil, fmt.Errorf("invalid output from '%s', ':' not found: %s", cmdline, output)
96	}
97
98	var list []string
99	for _, val := range bytes.Split(cpus, []byte(",")) {
100		index := string(bytes.TrimSpace(val))
101		if len(index) == 0 {
102			continue
103		}
104		list = append(list, index)
105	}
106	if len(list) == 0 {
107		return nil, fmt.Errorf("empty CPU list from '%s': %s", cmdline, output)
108	}
109	return list, nil
110}
111
112func checkNCPU(list []string) error {
113	listString := strings.Join(list, ",")
114	if len(listString) == 0 {
115		return fmt.Errorf("could not check against an empty CPU list")
116	}
117
118	cListString := cpuSetRE.FindString(listString)
119	if len(cListString) == 0 {
120		return fmt.Errorf("invalid cpuset output '%s'", listString)
121	}
122	// Launch FreeBSDNumCPUHelper() with specified CPUs list.
123	cmd := exec.Command("cpuset", "-l", cListString, os.Args[0], "FreeBSDNumCPUHelper")
124	cmdline := strings.Join(cmd.Args, " ")
125	output, err := cmd.CombinedOutput()
126	if err != nil {
127		return fmt.Errorf("fail to launch child '%s', error: %s, output: %s", cmdline, err, output)
128	}
129
130	// NumCPU from FreeBSDNumCPUHelper come with '\n'.
131	output = bytes.TrimSpace(output)
132	n, err := strconv.Atoi(string(output))
133	if err != nil {
134		return fmt.Errorf("fail to parse output from child '%s', error: %s, output: %s", cmdline, err, output)
135	}
136	if n != len(list) {
137		return fmt.Errorf("runtime.NumCPU() expected to %d, got %d when run with CPU list %s", len(list), n, cListString)
138	}
139	return nil
140}
141