gpu(darwin): format the string & fix behaviors on intel (#60)

* gpu(darwin): format the string

* gpu(darwin): fix behaviors on intel

* return string directly

pkg/gpu/gpu_darwin.go

@@ -7,6 +7,8 @@ import (
 	"unsafe"
 
 	"github.com/ebitengine/purego"
+
+	"github.com/nezhahq/agent/pkg/util"
 )
 
 type (
@@ -30,9 +32,13 @@ type (
 	CFStringCreateWithCStringFunc = func(alloc uintptr, cStr string, encoding CFStringEncoding) CFStringRef
 	CFGetTypeIDFunc               = func(cf uintptr) CFTypeID
 	CFStringGetTypeIDFunc         = func() CFTypeID
+	CFStringGetLengthFunc         = func(theString uintptr) int32
 	CFStringGetCStringFunc        = func(cfStr uintptr, buffer *byte, size CFIndex, encoding CFStringEncoding) bool
 	CFDictionaryGetTypeIDFunc     = func() CFTypeID
 	CFDictionaryGetValueFunc      = func(dict, key uintptr) unsafe.Pointer
+	CFDataGetTypeIDFunc           = func() CFTypeID
+	CFDataGetBytePtrFunc          = func(theData uintptr) unsafe.Pointer
+	CFDataGetLengthFunc           = func(theData uintptr) CFIndex
 	CFNumberGetValueFunc          = func(number uintptr, theType CFNumberType, valuePtr uintptr) bool
 	CFReleaseFunc                 = func(cf uintptr)
 
@@ -63,30 +69,19 @@ var (
 var (
 	coreFoundation, _ = purego.Dlopen("/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation", purego.RTLD_LAZY|purego.RTLD_GLOBAL)
 	ioKit, _          = purego.Dlopen("/System/Library/Frameworks/IOKit.framework/IOKit", purego.RTLD_LAZY|purego.RTLD_GLOBAL)
-
-	cfStringCreateWithCString, _ = purego.Dlsym(coreFoundation, "CFStringCreateWithCString")
-	cfGetTypeID, _               = purego.Dlsym(coreFoundation, "CFGetTypeID")
-	cfStringGetTypeID, _         = purego.Dlsym(coreFoundation, "CFStringGetTypeID")
-	cfStringGetCString, _        = purego.Dlsym(coreFoundation, "CFStringGetCString")
-	cfDictionaryGetTypeID, _     = purego.Dlsym(coreFoundation, "CFDictionaryGetTypeID")
-	cfDictionaryGetValue, _      = purego.Dlsym(coreFoundation, "CFDictionaryGetValue")
-	cfNumberGetValue, _          = purego.Dlsym(coreFoundation, "CFNumberGetValue")
-	cfRelease, _                 = purego.Dlsym(coreFoundation, "CFRelease")
-
-	ioServiceGetMatchingServices, _    = purego.Dlsym(ioKit, "IOServiceGetMatchingServices")
-	ioIteratorNext, _                  = purego.Dlsym(ioKit, "IOIteratorNext")
-	ioServiceMatching, _               = purego.Dlsym(ioKit, "IOServiceMatching")
-	ioRegistryEntrySearchCFProperty, _ = purego.Dlsym(ioKit, "IORegistryEntrySearchCFProperty")
-	ioObjectRelease, _                 = purego.Dlsym(ioKit, "IOObjectRelease")
 )
 
 var (
 	CFStringCreateWithCString CFStringCreateWithCStringFunc
 	CFGetTypeID               CFGetTypeIDFunc
 	CFStringGetTypeID         CFStringGetTypeIDFunc
+	CFStringGetLength         CFStringGetLengthFunc
 	CFStringGetCString        CFStringGetCStringFunc
 	CFDictionaryGetTypeID     CFDictionaryGetTypeIDFunc
 	CFDictionaryGetValue      CFDictionaryGetValueFunc
+	CFDataGetTypeID           CFDataGetTypeIDFunc
+	CFDataGetBytePtr          CFDataGetBytePtrFunc
+	CFDataGetLength           CFDataGetLengthFunc
 	CFNumberGetValue          CFNumberGetValueFunc
 	CFRelease                 CFReleaseFunc
 
@@ -97,25 +92,37 @@ var (
 	IOObjectRelease                 IOObjectReleaseFunc
 )
 
-func init() {
+var validVendors = []string{
-	purego.RegisterFunc(&CFStringCreateWithCString, cfStringCreateWithCString)
+	"AMD", "Intel", "NVIDIA", "Apple",
-	purego.RegisterFunc(&CFGetTypeID, cfGetTypeID)
+}
-	purego.RegisterFunc(&CFStringGetTypeID, cfStringGetTypeID)
-	purego.RegisterFunc(&CFStringGetCString, cfStringGetCString)
-	purego.RegisterFunc(&CFDictionaryGetTypeID, cfDictionaryGetTypeID)
-	purego.RegisterFunc(&CFDictionaryGetValue, cfDictionaryGetValue)
-	purego.RegisterFunc(&CFNumberGetValue, cfNumberGetValue)
-	purego.RegisterFunc(&CFRelease, cfRelease)
 
-	purego.RegisterFunc(&IOServiceGetMatchingServices, ioServiceGetMatchingServices)
+func init() {
-	purego.RegisterFunc(&IOIteratorNext, ioIteratorNext)
+	purego.RegisterLibFunc(&CFStringCreateWithCString, coreFoundation, "CFStringCreateWithCString")
-	purego.RegisterFunc(&IOServiceMatching, ioServiceMatching)
+	purego.RegisterLibFunc(&CFGetTypeID, coreFoundation, "CFGetTypeID")
-	purego.RegisterFunc(&IORegistryEntrySearchCFProperty, ioRegistryEntrySearchCFProperty)
+	purego.RegisterLibFunc(&CFStringGetTypeID, coreFoundation, "CFStringGetTypeID")
-	purego.RegisterFunc(&IOObjectRelease, ioObjectRelease)
+	purego.RegisterLibFunc(&CFStringGetLength, coreFoundation, "CFStringGetLength")
+	purego.RegisterLibFunc(&CFStringGetCString, coreFoundation, "CFStringGetCString")
+	purego.RegisterLibFunc(&CFDictionaryGetTypeID, coreFoundation, "CFDictionaryGetTypeID")
+	purego.RegisterLibFunc(&CFDictionaryGetValue, coreFoundation, "CFDictionaryGetValue")
+	purego.RegisterLibFunc(&CFDataGetTypeID, coreFoundation, "CFDataGetTypeID")
+	purego.RegisterLibFunc(&CFDataGetBytePtr, coreFoundation, "CFDataGetBytePtr")
+	purego.RegisterLibFunc(&CFDataGetLength, coreFoundation, "CFDataGetLength")
+	purego.RegisterLibFunc(&CFNumberGetValue, coreFoundation, "CFNumberGetValue")
+	purego.RegisterLibFunc(&CFRelease, coreFoundation, "CFRelease")
+
+	purego.RegisterLibFunc(&IOServiceGetMatchingServices, ioKit, "IOServiceGetMatchingServices")
+	purego.RegisterLibFunc(&IOIteratorNext, ioKit, "IOIteratorNext")
+	purego.RegisterLibFunc(&IOServiceMatching, ioKit, "IOServiceMatching")
+	purego.RegisterLibFunc(&IORegistryEntrySearchCFProperty, ioKit, "IORegistryEntrySearchCFProperty")
+	purego.RegisterLibFunc(&IOObjectRelease, ioKit, "IOObjectRelease")
 }
 
 func GetGPUModel() ([]string, error) {
-	return findDevices("model")
+	models, err := findDevices("model")
+	if err != nil {
+		return nil, err
+	}
+	return util.RemoveDuplicate(models), nil
 }
 
 func FindUtilization(key, dictKey string) (int, error) {
@@ -125,6 +132,7 @@ func FindUtilization(key, dictKey string) (int, error) {
 func findDevices(key string) ([]string, error) {
 	var iterator ioIterator
 	var results []string
+	done := false
 
 	iv := IOServiceGetMatchingServices(kIOMainPortDefault, uintptr(IOServiceMatching(IOSERVICE_GPU)), &iterator)
 	if iv != KERN_SUCCESS {
@@ -144,15 +152,16 @@ func findDevices(key string) ([]string, error) {
 		result, _, _ := findProperties(service, uintptr(cfStr), 0)
 		IOObjectRelease(service)
 
-		if result != nil {
+		if util.ContainsStr(validVendors, result) {
-			results = append(results, string(result))
+			results = append(results, result)
 			index++
-		} else if key == "model" {
+		} else if key == "model" && !done {
 			IOObjectRelease(iterator)
 			iv = IOServiceGetMatchingServices(kIOMainPortDefault, uintptr(IOServiceMatching(IOSERVICE_PCI)), &iterator)
 			if iv != KERN_SUCCESS {
 				return nil, fmt.Errorf("error retrieving GPU entry")
 			}
+			done = true
 		}
 	}
 
@@ -162,17 +171,22 @@ func findDevices(key string) ([]string, error) {
 
 func findUtilization(key, dictKey string) (int, error) {
 	var iterator ioIterator
-	var result int
 	var err error
+	result := 0
 
 	iv := IOServiceGetMatchingServices(kIOMainPortDefault, uintptr(IOServiceMatching(IOSERVICE_GPU)), &iterator)
 	if iv != KERN_SUCCESS {
 		return 0, fmt.Errorf("error retrieving GPU entry")
 	}
 
-	// Only retrieving the utilization of first GPU here
+	// Only retrieving the utilization of a single GPU here
-	service := IOIteratorNext(iterator)
+	var service ioObject
-	if service != MACH_PORT_NULL {
+	for {
+		service = IOIteratorNext(iterator)
+		if service == MACH_PORT_NULL {
+			break
+		}
+
 		cfStr := CFStringCreateWithCString(kCFAllocatorDefault, key, CFStringEncoding(kCFStringEncodingUTF8))
 		cfDictStr := CFStringCreateWithCString(kCFAllocatorDefault, dictKey, CFStringEncoding(kCFStringEncodingUTF8))
 
@@ -182,44 +196,50 @@ func findUtilization(key, dictKey string) (int, error) {
 		CFRelease(uintptr(cfDictStr))
 
 		if err != nil {
-			return 0, fmt.Errorf("failed retrieving GPU utilization: %v", err)
-		}
-	} else {
 			IOObjectRelease(service)
-		IOObjectRelease(iterator)
+			continue
-		return 0, fmt.Errorf("no GPU utilization entry found")
+		} else if result != 0 {
+			break
+		}
 	}
 
 	IOObjectRelease(service)
 	IOObjectRelease(iterator)
-	return result, nil
+
+	return result, err
 }
 
-func findProperties(service ioRegistryEntry, key, dictKey uintptr) ([]byte, int, error) {
+func findProperties(service ioRegistryEntry, key, dictKey uintptr) (string, int, error) {
 	properties := IORegistryEntrySearchCFProperty(service, kIOServicePlane, key, kCFAllocatorDefault, kIORegistryIterateRecursively)
 	ptrValue := uintptr(properties)
 	if properties != nil {
 		switch CFGetTypeID(ptrValue) {
 		// model
 		case CFStringGetTypeID():
-			buf := make([]byte, 1024)
+			length := CFStringGetLength(ptrValue) + 1 // null terminator
-			CFStringGetCString(ptrValue, &buf[0], int32(unsafe.Sizeof(buf)), uint32(kCFStringEncodingUTF8))
+			buf := make([]byte, length-1)
+			CFStringGetCString(ptrValue, &buf[0], length, uint32(kCFStringEncodingUTF8))
 			CFRelease(ptrValue)
-			return buf, 0, nil
+			return string(buf), 0, nil
+		case CFDataGetTypeID():
+			length := CFDataGetLength(ptrValue)
+			bin := unsafe.String((*byte)(CFDataGetBytePtr(ptrValue)), length)
+			CFRelease(ptrValue)
+			return bin, 0, nil
 		// PerformanceStatistics
 		case CFDictionaryGetTypeID():
 			cfValue := CFDictionaryGetValue(ptrValue, dictKey)
 			if cfValue != nil {
 				var value int
 				if CFNumberGetValue(uintptr(cfValue), kCFNumberIntType, uintptr(unsafe.Pointer(&value))) {
-					return nil, value, nil
+					return "", value, nil
 				} else {
-					return nil, 0, fmt.Errorf("failed to exec CFNumberGetValue")
+					return "", 0, fmt.Errorf("failed to exec CFNumberGetValue")
 				}
 			} else {
-				return nil, 0, fmt.Errorf("failed to exec CFDictionaryGetValue")
+				return "", 0, fmt.Errorf("failed to exec CFDictionaryGetValue")
 			}
 		}
 	}
-	return nil, 0, fmt.Errorf("failed to exec IORegistryEntrySearchCFProperty")
+	return "", 0, fmt.Errorf("failed to exec IORegistryEntrySearchCFProperty")
 }

pkg/monitor/monitor.go

@@ -350,7 +350,7 @@ func updateGPUStat() float64 {
 			gs, err := gpustat.GetGPUStat()
 			if err != nil {
 				statDataFetchAttempts["GPU"]++
-				println("gpustat.GetGPUStat error: ", err, ", attempt: ", statDataFetchAttempts["GPU"])
+				printf("gpustat.GetGPUStat error: %v, attempt: %d", err, statDataFetchAttempts["GPU"])
 				return 0
 			} else {
 				statDataFetchAttempts["GPU"] = 0

pkg/util/util.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"strings"
 	"time"
 
 	jsoniter "github.com/json-iterator/go"
@@ -40,3 +41,26 @@ func BrowserHeaders() *http.Header {
 		"User-Agent":      {MacOSChromeUA},
 	}
 }
+
+func ContainsStr(slice []string, str string) bool {
+	if str != "" {
+		for _, item := range slice {
+			if strings.Contains(str, item) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func RemoveDuplicate[T comparable](sliceList []T) []T {
+	allKeys := make(map[T]bool)
+	list := []T{}
+	for _, item := range sliceList {
+		if _, value := allKeys[item]; !value {
+			allKeys[item] = true
+			list = append(list, item)
+		}
+	}
+	return list
+}