Merge pull request #1768 from uubulb/darwin

refactor(darwin): optimize string allocation
shirou · Dec 31, 2024 · 252dcbf · 252dcbf
2 parents bcd0c0a + 1f39525
commit 252dcbf
Show file tree

Hide file tree

Showing 7 changed files with 98 additions and 33 deletions.
diff --git a/cpu/cpu_darwin_arm64.go b/cpu/cpu_darwin_arm64.go
@@ -55,10 +55,10 @@ func getFrequency() (float64, error) {
 			break
 		}
 
-		buf := make([]byte, 512)
-		ioRegistryEntryGetName(service, &buf[0])
+		buf := common.NewCStr(512)
+		ioRegistryEntryGetName(service, buf)
 
-		if common.GoString(&buf[0]) == "pmgr" {
+		if buf.GoString() == "pmgr" {
 			pCoreRef := ioRegistryEntryCreateCFProperty(service, uintptr(pCorekey), common.KCFAllocatorDefault, common.KNilOptions)
 			length := cfDataGetLength(uintptr(pCoreRef))
 			data := cfDataGetBytePtr(uintptr(pCoreRef))

diff --git a/disk/disk_darwin.go b/disk/disk_darwin.go
@@ -5,6 +5,7 @@ package disk
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"unsafe"
 
@@ -234,17 +235,17 @@ func (i *ioCounters) getDriveStat(d uint32) (*IOCountersStat, error) {
 	key := i.cfStr(kIOBSDNameKey)
 	defer i.cfRelease(uintptr(key))
 	name := i.cfDictionaryGetValue(uintptr(props), uintptr(key))
-	length := cfStringGetLength(uintptr(name)) + 1
-	buf := make([]byte, length-1)
-	cfStringGetCString(uintptr(name), &buf[0], length, common.KCFStringEncodingUTF8)
+
+	buf := common.NewCStr(cfStringGetLength(uintptr(name)))
+	cfStringGetCString(uintptr(name), buf, buf.Length(), common.KCFStringEncodingUTF8)
 
 	stat, err := i.fillStat(parent)
 	if err != nil {
 		return nil, err
 	}
 
 	if stat != nil {
-		stat.Name = string(buf)
+		stat.Name = buf.GoString()
 		return stat, nil
 	}
 	return nil, nil
@@ -263,9 +264,10 @@ func (i *ioCounters) fillStat(d uint32) (*IOCountersStat, error) {
 
 	key := i.cfStr(kIOBlockStorageDriverStatisticsKey)
 	defer i.cfRelease(uintptr(key))
+
 	v := i.cfDictionaryGetValue(uintptr(props), uintptr(key))
 	if v == nil {
-		return nil, fmt.Errorf("CFDictionaryGetValue failed")
+		return nil, errors.New("CFDictionaryGetValue failed")
 	}
 
 	var stat IOCountersStat
@@ -280,10 +282,10 @@ func (i *ioCounters) fillStat(d uint32) (*IOCountersStat, error) {
 
 	for key, off := range statstab {
 		s := i.cfStr(key)
-		defer i.cfRelease(uintptr(s))
 		if num := i.cfDictionaryGetValue(uintptr(v), uintptr(s)); num != nil {
-			i.cfNumberGetValue(uintptr(num), common.KCFNumberSInt64Type, uintptr(unsafe.Pointer(uintptr(unsafe.Pointer(&stat))+off)))
+			i.cfNumberGetValue(uintptr(num), common.KCFNumberSInt64Type, uintptr(unsafe.Add(unsafe.Pointer(&stat), off)))
 		}
+		i.cfRelease(uintptr(s))
 	}
 
 	return &stat, nil

diff --git a/internal/common/common_darwin.go b/internal/common/common_darwin.go
@@ -125,7 +125,7 @@ type (
 	IOServiceOpenFunc                     func(service, owningTask, connType uint32, connect *uint32) int
 	IOServiceCloseFunc                    func(connect uint32) int
 	IOIteratorNextFunc                    func(iterator uint32) uint32
-	IORegistryEntryGetNameFunc            func(entry uint32, name *byte) int
+	IORegistryEntryGetNameFunc            func(entry uint32, name CStr) int
 	IORegistryEntryGetParentEntryFunc     func(entry uint32, plane string, parent *uint32) int
 	IORegistryEntryCreateCFPropertyFunc   func(entry uint32, key, allocator uintptr, options uint32) unsafe.Pointer
 	IORegistryEntryCreateCFPropertiesFunc func(entry uint32, properties unsafe.Pointer, allocator uintptr, options uint32) int
@@ -191,7 +191,7 @@ type (
 	CFArrayGetValueAtIndexFunc    func(theArray uintptr, index int32) unsafe.Pointer
 	CFStringCreateMutableFunc     func(alloc uintptr, maxLength int32) unsafe.Pointer
 	CFStringGetLengthFunc         func(theString uintptr) int32
-	CFStringGetCStringFunc        func(theString uintptr, buffer *byte, bufferSize int32, encoding uint32)
+	CFStringGetCStringFunc        func(theString uintptr, buffer CStr, bufferSize int32, encoding uint32)
 	CFStringCreateWithCStringFunc func(alloc uintptr, cStr string, encoding uint32) unsafe.Pointer
 	CFDataGetLengthFunc           func(theData uintptr) int32
 	CFDataGetBytePtrFunc          func(theData uintptr) unsafe.Pointer
@@ -348,6 +348,44 @@ func (s *SMC) Close() error {
 	return nil
 }
 
+type CStr []byte
+
+func NewCStr(length int32) CStr {
+	return make(CStr, length)
+}
+
+func (s CStr) Length() int32 {
+	// Include null terminator to make CFStringGetCString properly functions
+	return int32(len(s)) + 1
+}
+
+func (s CStr) Ptr() *byte {
+	if len(s) < 1 {
+		return nil
+	}
+
+	return &s[0]
+}
+
+func (c CStr) Addr() uintptr {
+	return uintptr(unsafe.Pointer(c.Ptr()))
+}
+
+func (s CStr) GoString() string {
+	if s == nil {
+		return ""
+	}
+
+	var length int
+	for _, char := range s {
+		if char == '\x00' {
+			break
+		}
+		length++
+	}
+	return string(s[:length])
+}
+
 // https://github.com/ebitengine/purego/blob/main/internal/strings/strings.go#L26
 func GoString(cStr *byte) string {
 	if cStr == nil {

diff --git a/process/process_darwin.go b/process/process_darwin.go
@@ -310,14 +310,14 @@ func (p *Process) ExeWithContext(ctx context.Context) (string, error) {
 	}
 	defer lib.Close()
 
-	buf := make([]byte, common.PROC_PIDPATHINFO_MAXSIZE)
-	ret := procPidPath(p.Pid, uintptr(unsafe.Pointer(&buf[0])), common.PROC_PIDPATHINFO_MAXSIZE)
+	buf := common.NewCStr(common.PROC_PIDPATHINFO_MAXSIZE)
+	ret := procPidPath(p.Pid, buf.Addr(), common.PROC_PIDPATHINFO_MAXSIZE)
 
 	if ret <= 0 {
 		return "", fmt.Errorf("unknown error: proc_pidpath returned %d", ret)
 	}
 
-	return common.GoString(&buf[0]), nil
+	return buf.GoString(), nil
 }
 
 // sys/proc_info.h
@@ -339,6 +339,7 @@ func (p *Process) CwdWithContext(ctx context.Context) (string, error) {
 	}
 	defer lib.Close()
 
+	// Lock OS thread to ensure the errno does not change
 	runtime.LockOSThread()
 	defer runtime.UnlockOSThread()
 
@@ -366,6 +367,8 @@ func procArgs(pid int32) ([]byte, int, error) {
 	if err != nil {
 		return nil, 0, err
 	}
+
+	// The first 4 bytes indicate the number of arguments.
 	nargs := procargs[:4]
 	return procargs, int(binary.LittleEndian.Uint32(nargs)), nil
 }
@@ -434,8 +437,7 @@ func (p *Process) NumThreadsWithContext(ctx context.Context) (int32, error) {
 	defer lib.Close()
 
 	var ti ProcTaskInfo
-	const tiSize = int32(unsafe.Sizeof(ti))
-	procPidInfo(p.Pid, common.PROC_PIDTASKINFO, 0, uintptr(unsafe.Pointer(&ti)), tiSize)
+	procPidInfo(p.Pid, common.PROC_PIDTASKINFO, 0, uintptr(unsafe.Pointer(&ti)), int32(unsafe.Sizeof(ti)))
 
 	return int32(ti.Threadnum), nil
 }
@@ -448,8 +450,7 @@ func (p *Process) TimesWithContext(ctx context.Context) (*cpu.TimesStat, error)
 	defer lib.Close()
 
 	var ti ProcTaskInfo
-	const tiSize = int32(unsafe.Sizeof(ti))
-	procPidInfo(p.Pid, common.PROC_PIDTASKINFO, 0, uintptr(unsafe.Pointer(&ti)), tiSize)
+	procPidInfo(p.Pid, common.PROC_PIDTASKINFO, 0, uintptr(unsafe.Pointer(&ti)), int32(unsafe.Sizeof(ti)))
 
 	timescaleToNanoSeconds := getTimeScaleToNanoSeconds()
 	ret := &cpu.TimesStat{
@@ -468,8 +469,7 @@ func (p *Process) MemoryInfoWithContext(ctx context.Context) (*MemoryInfoStat, e
 	defer lib.Close()
 
 	var ti ProcTaskInfo
-	const tiSize = int32(unsafe.Sizeof(ti))
-	procPidInfo(p.Pid, common.PROC_PIDTASKINFO, 0, uintptr(unsafe.Pointer(&ti)), tiSize)
+	procPidInfo(p.Pid, common.PROC_PIDTASKINFO, 0, uintptr(unsafe.Pointer(&ti)), int32(unsafe.Sizeof(ti)))
 
 	ret := &MemoryInfoStat{
 		RSS:  uint64(ti.Resident_size),

diff --git a/sensors/sensors_darwin.go b/sensors/sensors_darwin.go
@@ -24,7 +24,7 @@ func TemperaturesWithContext(ctx context.Context) ([]TemperatureStat, error) {
 	}
 	defer smc.Close()
 
-	var temperatures []TemperatureStat
+	temperatures := make([]TemperatureStat, 0, len(temperatureKeys))
 	for _, key := range temperatureKeys {
 		temperatures = append(temperatures, TemperatureStat{
 			SensorKey:   key,

diff --git a/sensors/sensors_darwin_arm64.go b/sensors/sensors_darwin_arm64.go
@@ -94,8 +94,6 @@ func (ta *temperatureArm) getProductNames(system unsafe.Pointer) []string {
 	cfStringGetLength := common.GetFunc[common.CFStringGetLengthFunc](ta.cf, common.CFStringGetLengthSym)
 	cfStringGetCString := common.GetFunc[common.CFStringGetCStringFunc](ta.cf, common.CFStringGetCStringSym)
 
-	var names []string
-
 	ta.ioHIDEventSystemClientSetMatching(uintptr(system), uintptr(ta.sensors))
 	matchingsrvs := ta.ioHIDEventSystemClientCopyServices(uintptr(system))
 
@@ -110,18 +108,19 @@ func (ta *temperatureArm) getProductNames(system unsafe.Pointer) []string {
 	str := ta.cfStr("Product")
 	defer ta.cfRelease(uintptr(str))
 
+	names := make([]string, 0, count)
 	for i = 0; i < count; i++ {
 		sc := ta.cfArrayGetValueAtIndex(uintptr(matchingsrvs), i)
 		name := ioHIDServiceClientCopyProperty(uintptr(sc), uintptr(str))
 
 		if name != nil {
-			length := cfStringGetLength(uintptr(name)) + 1 // include null terminator
-			buf := make([]byte, length)                    // allocate buffer with full length
-			cfStringGetCString(uintptr(name), &buf[0], length, common.KCFStringEncodingUTF8)
+			buf := common.NewCStr(cfStringGetLength(uintptr(name)))
+			cfStringGetCString(uintptr(name), buf, buf.Length(), common.KCFStringEncodingUTF8)
 
-			names = append(names, string(buf[:length-1])) // remove null terminator
+			names = append(names, buf.GoString())
 			ta.cfRelease(uintptr(name))
 		} else {
+			// make sure the number of names and values are consistent
 			names = append(names, "noname")
 		}
 	}
@@ -166,11 +165,17 @@ func (ta *temperatureArm) matching(page, usage int) {
 	cfDictionaryCreate := common.GetFunc[common.CFDictionaryCreateFunc](ta.cf, common.CFDictionaryCreateSym)
 
 	pageNum := cfNumberCreate(common.KCFAllocatorDefault, common.KCFNumberIntType, uintptr(unsafe.Pointer(&page)))
+	defer ta.cfRelease(uintptr(pageNum))
+
 	usageNum := cfNumberCreate(common.KCFAllocatorDefault, common.KCFNumberIntType, uintptr(unsafe.Pointer(&usage)))
+	defer ta.cfRelease(uintptr(usageNum))
 
 	k1 := ta.cfStr("PrimaryUsagePage")
 	k2 := ta.cfStr("PrimaryUsage")
 
+	defer ta.cfRelease(uintptr(k1))
+	defer ta.cfRelease(uintptr(k2))
+
 	keys := []unsafe.Pointer{k1, k2}
 	values := []unsafe.Pointer{pageNum, usageNum}
 
@@ -180,11 +185,6 @@ func (ta *temperatureArm) matching(page, usage int) {
 	ta.sensors = cfDictionaryCreate(common.KCFAllocatorDefault, &keys[0], &values[0], 2,
 		kCFTypeDictionaryKeyCallBacks,
 		kCFTypeDictionaryValueCallBacks)
-
-	ta.cfRelease(uintptr(pageNum))
-	ta.cfRelease(uintptr(usageNum))
-	ta.cfRelease(uintptr(k1))
-	ta.cfRelease(uintptr(k2))
 }
 
 func (ta *temperatureArm) cfStr(str string) unsafe.Pointer {

diff --git a/sensors/sensors_test.go b/sensors/sensors_test.go
@@ -3,8 +3,12 @@
 package sensors
 
 import (
+	"errors"
 	"fmt"
+	"os"
 	"testing"
+
+	"github.com/shirou/gopsutil/v4/internal/common"
 )
 
 func TestTemperatureStat_String(t *testing.T) {
@@ -19,3 +23,24 @@ func TestTemperatureStat_String(t *testing.T) {
 		t.Errorf("TemperatureStat string is invalid, %v", fmt.Sprintf("%v", v))
 	}
 }
+
+func skipIfNotImplementedErr(t *testing.T, err error) {
+	if errors.Is(err, common.ErrNotImplementedError) {
+		t.Skip("not implemented")
+	}
+}
+
+func TestTemperatures(t *testing.T) {
+	if os.Getenv("CI") != "" {
+		t.Skip("Skip CI")
+	}
+	v, err := SensorsTemperatures()
+	skipIfNotImplementedErr(t, err)
+	if err != nil {
+		t.Errorf("error %v", err)
+	}
+	if len(v) == 0 {
+		t.Errorf("Could not get temperature %v", v)
+	}
+	t.Log(v)
+}