Refactor memory

Author: nekohasekai

common/memory/memory.go

@@ -3,22 +3,23 @@ package memory
 import "runtime"
 
 func Total() uint64 {
-	if nativeAvailable {
-		return usageNative()
-	}
-	return Inuse()
+	return totalNative()
 }
 
-func Inuse() uint64 {
-	var memStats runtime.MemStats
-	runtime.ReadMemStats(&memStats)
-	return memStats.StackInuse + memStats.HeapInuse + memStats.HeapIdle - memStats.HeapReleased
+func TotalAvailable() bool {
+	return totalAvailable()
 }
 
 func Available() uint64 {
 	return availableNative()
 }
 
-func AvailableSupported() bool {
-	return availableNativeSupported()
+func AvailableAvailable() bool {
+	return availableAvailable()
+}
+
+func Inuse() uint64 {
+	var memStats runtime.MemStats
+	runtime.ReadMemStats(&memStats)
+	return memStats.StackInuse + memStats.HeapInuse + memStats.HeapIdle - memStats.HeapReleased
 }

common/memory/memory_available_darwin.go

@@ -1,18 +1,59 @@
 package memory
 
 // #include <mach/mach.h>
+// #include <stddef.h>
+// #include <dlfcn.h>
+//
+// typedef size_t (*proc_available_memory_func)(void);
+// static int resolved = 0;
+// static proc_available_memory_func fn = NULL;
+//
+// static void resolve_available_memory() {
+//     if (!resolved) {
+//         fn = (proc_available_memory_func)dlsym(RTLD_DEFAULT, "os_proc_available_memory");
+//         resolved = 1;
+//     }
+// }
+//
+// static size_t get_available_memory(int *supported) {
+//     resolve_available_memory();
+//     if (fn) {
+//         *supported = 1;
+//         return fn();
+//     }
+//     *supported = 0;
+//     return 0;
+// }
+//
+// static int is_available_memory_supported() {
+//     resolve_available_memory();
+//     return fn != NULL;
+// }
 import "C"
 import "unsafe"
 
-const nativeAvailable = true
-
-func usageNative() uint64 {
-	var memoryUsageInByte uint64
+func totalNative() uint64 {
 	var vmInfo C.task_vm_info_data_t
 	var count C.mach_msg_type_number_t = C.TASK_VM_INFO_COUNT
-	var kernelReturn C.kern_return_t = C.task_info(C.vm_map_t(C.mach_task_self_), C.TASK_VM_INFO, (*C.integer_t)(unsafe.Pointer(&vmInfo)), &count)
-	if kernelReturn == C.KERN_SUCCESS {
-		memoryUsageInByte = uint64(vmInfo.phys_footprint)
+	if C.task_info(C.vm_map_t(C.mach_task_self_), C.TASK_VM_INFO, (*C.integer_t)(unsafe.Pointer(&vmInfo)), &count) == C.KERN_SUCCESS {
+		return uint64(vmInfo.phys_footprint)
+	}
+	return 0
+}
+
+func totalAvailable() bool {
+	return true
+}
+
+func availableNative() uint64 {
+	var supported C.int
+	result := C.get_available_memory(&supported)
+	if supported == 0 {
+		return 0
 	}
-	return memoryUsageInByte
+	return uint64(result)
+}
+
+func availableAvailable() bool {
+	return C.is_available_memory_supported() != 0
 }

common/memory/memory_available_stub.go

@@ -8,8 +8,41 @@ import (
 	"strings"
 )
 
-func availableNativeSupported() bool {
-	return true
+var pageSize = uint64(os.Getpagesize())
+
+func totalNative() uint64 {
+	fd, err := os.Open("/proc/self/statm")
+	if err != nil {
+		return 0
+	}
+	defer fd.Close()
+	var buf [128]byte
+	n, _ := fd.Read(buf[:])
+	if n == 0 {
+		return 0
+	}
+	i := 0
+	for i < n && buf[i] != ' ' {
+		i++
+	}
+	i++
+	var rss uint64
+	for i < n && buf[i] >= '0' && buf[i] <= '9' {
+		rss = rss*10 + uint64(buf[i]-'0')
+		i++
+	}
+	return rss * pageSize
+}
+
+func totalAvailable() bool {
+	fd, err := os.Open("/proc/self/statm")
+	if err != nil {
+		return false
+	}
+	defer fd.Close()
+	var buf [1]byte
+	n, _ := fd.Read(buf[:])
+	return n > 0
 }
 
 func availableNative() uint64 {
@@ -20,6 +53,19 @@ func availableNative() uint64 {
 	return procMemAvailable()
 }
 
+func availableAvailable() bool {
+	_, ok := cgroupAvailable()
+	if ok {
+		return true
+	}
+	fd, err := os.Open("/proc/meminfo")
+	if err != nil {
+		return false
+	}
+	fd.Close()
+	return true
+}
+
 func cgroupAvailable() (uint64, bool) {
 	max, err := readCgroupUint("/sys/fs/cgroup/memory.max")
 	if err == nil && max != math.MaxUint64 {

common/memory/memory_darwin.go

@@ -1,9 +1,19 @@
-//go:build (darwin && !cgo) || !darwin
+//go:build (darwin && !cgo) || (!darwin && !linux && !windows)
 
 package memory
 
-const nativeAvailable = false
+func totalNative() uint64 {
+	return 0
+}
+
+func totalAvailable() bool {
+	return false
+}
 
-func usageNative() uint64 {
+func availableNative() uint64 {
 	return 0
 }
+
+func availableAvailable() bool {
+	return false
+}

common/memory/memory_available_linux.go common/memory/memory_linux.gocommon/memory/memory_available_linux.go renamed to common/memory/memory_linux.go

@@ -0,0 +1,35 @@
+package memory
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+func totalNative() uint64 {
+	var mem processMemoryCounters
+	mem.cb = uint32(unsafe.Sizeof(mem))
+	err := getProcessMemoryInfo(windows.CurrentProcess(), &mem, mem.cb)
+	if err != nil {
+		return 0
+	}
+	return uint64(mem.workingSetSize)
+}
+
+func totalAvailable() bool {
+	return true
+}
+
+func availableNative() uint64 {
+	var mem memoryStatusEx
+	mem.dwLength = uint32(unsafe.Sizeof(mem))
+	err := globalMemoryStatusEx(&mem)
+	if err != nil {
+		return 0
+	}
+	return mem.ullAvailPhys
+}
+
+func availableAvailable() bool {
+	return true
+}

common/memory/memory_stub.go

@@ -0,0 +1,34 @@
+package memory
+
+//go:generate go run golang.org/x/sys/windows/mkwinsyscall -output zsyscall_windows.go syscall_windows.go
+
+type processMemoryCounters struct {
+	cb                         uint32
+	pageFaultCount             uint32
+	peakWorkingSetSize         uintptr
+	workingSetSize             uintptr
+	quotaPeakPagedPoolUsage    uintptr
+	quotaPagedPoolUsage        uintptr
+	quotaPeakNonPagedPoolUsage uintptr
+	quotaNonPagedPoolUsage     uintptr
+	pagefileUsage              uintptr
+	peakPagefileUsage          uintptr
+}
+
+type memoryStatusEx struct {
+	dwLength                uint32
+	dwMemoryLoad            uint32
+	ullTotalPhys            uint64
+	ullAvailPhys            uint64
+	ullTotalPageFile        uint64
+	ullAvailPageFile        uint64
+	ullTotalVirtual         uint64
+	ullAvailVirtual         uint64
+	ullAvailExtendedVirtual uint64
+}
+
+// https://learn.microsoft.com/en-us/windows/win32/api/psapi/nf-psapi-getprocessmemoryinfo
+//sys getProcessMemoryInfo(process windows.Handle, ppsmemCounters *processMemoryCounters, cb uint32) (err error) = kernel32.K32GetProcessMemoryInfo
+
+// https://learn.microsoft.com/en-us/windows/win32/api/sysinfoapi/nf-sysinfoapi-globalmemorystatusex
+//sys globalMemoryStatusEx(lpBuffer *memoryStatusEx) (err error) = kernel32.GlobalMemoryStatusEx