[cpu][linux] Add support for logical arg in Counts #640 #628

cpu/cpu.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"runtime"
 	"strconv"
 	"strings"
 	"sync"
@@ -63,14 +62,11 @@ func init() {
 	lastCPUPercent.Unlock()
 }
 
+// Counts returns the number of physical or logical cores in the system
 func Counts(logical bool) (int, error) {
 	return CountsWithContext(context.Background(), logical)
 }
 
-func CountsWithContext(ctx context.Context, logical bool) (int, error) {
-	return runtime.NumCPU(), nil
-}
-
 func (c TimesStat) String() string {
 	v := []string{
 		`"cpu":"` + c.CPU + `"`,

cpu/cpu_darwin.go

@@ -4,6 +4,7 @@ package cpu
 
 import (
 	"context"
+	"runtime"
 	"strconv"
 	"strings"
 
@@ -85,3 +86,7 @@ func InfoWithContext(ctx context.Context) ([]InfoStat, error) {
 
 	return append(ret, c), nil
 }
+
+func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+	return runtime.NumCPU(), nil
+}

cpu/cpu_fallback.go

@@ -4,6 +4,7 @@ package cpu
 
 import (
 	"context"
+	"runtime"
 
 	"github.com/shirou/gopsutil/internal/common"
 )
@@ -23,3 +24,7 @@ func Info() ([]InfoStat, error) {
 func InfoWithContext(ctx context.Context) ([]InfoStat, error) {
 	return []InfoStat{}, common.ErrNotImplementedError
 }
+
+func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+	return runtime.NumCPU(), nil
+}

cpu/cpu_freebsd.go

@@ -6,6 +6,7 @@ import (
 	"os/exec"
 	"reflect"
 	"regexp"
+	"runtime"
 	"strconv"
 	"strings"
 	"unsafe"
@@ -166,3 +167,7 @@ func parseDmesgBoot(fileName string) (InfoStat, int, error) {
 
 	return c, cpuNum, nil
 }
+
+func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+	return runtime.NumCPU(), nil
+}

cpu/cpu_linux.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"os/exec"
+	"regexp"
 	"strconv"
 	"strings"
 
@@ -282,3 +283,74 @@ func parseStatLine(line string) (*TimesStat, error) {
 
 	return ct, nil
 }
+
+func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+	if logical {
+		ret := 0
+		// https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_pslinux.py#L599
+		procCpuinfo := common.HostProc("cpuinfo")
+		lines, err := common.ReadLines(procCpuinfo)
+		if err == nil {
+			for _, line := range lines {
+				line = strings.ToLower(line)
+				if strings.HasPrefix(line, "processor") {
+					ret++
+				}
+			}
+		}
+		if ret == 0 {
+			procStat := common.HostProc("stat")
+			lines, err = common.ReadLines(procStat)
+			if err != nil {
+				return 0, err
+			}
+			re := regexp.MustCompile(`cpu\d`)
+			for _, line := range lines {
+				line = strings.Split(line, " ")[0]
+				if re.MatchString(line) {
+					ret++
+				}
+			}
+		}
+		return ret, nil
+	}
+	// physical cores https://github.com/giampaolo/psutil/blob/d01a9eaa35a8aadf6c519839e987a49d8be2d891/psutil/_pslinux.py#L628
+	filename := common.HostProc("cpuinfo")
+	lines, err := common.ReadLines(filename)
+	if err != nil {
+		return 0, err
+	}
+	mapping := make(map[int]int)
+	currentInfo := make(map[string]int)
+	for _, line := range lines {
+		line = strings.ToLower(strings.TrimSpace(line))
+		if line == "" {
+			// new section
+			id, okID := currentInfo["physical id"]
+			cores, okCores := currentInfo["cpu cores"]
+			if okID && okCores {
+				mapping[id] = cores
+			}
+			currentInfo = make(map[string]int)
+			continue
+		}
+		fields := strings.Split(line, ":")
+		if len(fields) < 2 {
+			continue
+		}
+		fields[0] = strings.TrimSpace(fields[0])
+		if fields[0] == "physical id" || fields[0] == "cpu cores" {
+			val, err := strconv.Atoi(strings.TrimSpace(fields[1]))
+			if err != nil {
+				fmt.Println(err)
+				continue
+			}
+			currentInfo[fields[0]] = val
+		}
+	}
+	ret := 0
+	for _, v := range mapping {
+		ret += v
+	}
+	return ret, nil
+}

cpu/cpu_openbsd.go

@@ -8,6 +8,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"os/exec"
+	"runtime"
 	"strconv"
 	"strings"
 
@@ -145,3 +146,7 @@ func InfoWithContext(ctx context.Context) ([]InfoStat, error) {
 
 	return append(ret, c), nil
 }
+
+func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+	return runtime.NumCPU(), nil
+}

cpu/cpu_solaris.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os/exec"
 	"regexp"
+	"runtime"
 	"sort"
 	"strconv"
 	"strings"
@@ -196,3 +197,7 @@ func parseProcessorInfo(cmdOutput string) ([]InfoStat, error) {
 	}
 	return result, nil
 }
+
+func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+	return runtime.NumCPU(), nil
+}

cpu/cpu_windows.go

@@ -5,6 +5,7 @@ package cpu
 import (
 	"context"
 	"fmt"
+	"runtime"
 	"unsafe"
 
 	"github.com/StackExchange/wmi"
@@ -199,3 +200,7 @@ func perfInfo() ([]win32_SystemProcessorPerformanceInformation, error) {
 
 	return resultBuffer, nil
 }
+
+func CountsWithContext(ctx context.Context, logical bool) (int, error) {
+	return runtime.NumCPU(), nil
+}