From d78bf1122d2f21738bb69c6cc06a29c4462d6c62 Mon Sep 17 00:00:00 2001
From: kbkpbot <kbkpbot@sina.com>
Date: Mon, 3 Feb 2025 01:32:21 +0800
Subject: [PATCH] os: add disk_usage/1 (#23634)

---
 vlib/builtin/cfns.c.v  |  4 ++++
 vlib/os/os.c.v         |  7 +++++++
 vlib/os/os_nix.c.v     | 29 +++++++++++++++++++++++++++++
 vlib/os/os_test.c.v    |  7 +++++++
 vlib/os/os_windows.c.v | 23 +++++++++++++++++++++++
 5 files changed, 70 insertions(+)

diff --git a/vlib/builtin/cfns.c.v b/vlib/builtin/cfns.c.v
index bffec65d28e793..f8332c25cc0bfe 100644
--- a/vlib/builtin/cfns.c.v
+++ b/vlib/builtin/cfns.c.v
@@ -142,6 +142,8 @@ fn C.stat(&char, voidptr) int
 
 fn C.lstat(path &char, buf &C.stat) int
 
+fn C.statvfs(const_path &char, buf &C.statvfs) int
+
 fn C.rename(old_filename &char, new_filename &char) int
 
 fn C.fgets(str &char, n int, stream &C.FILE) int
@@ -513,3 +515,5 @@ fn C.WrappedNSLog(str &u8)
 // absolute value
 @[trusted]
 fn C.abs(number int) int
+
+fn C.GetDiskFreeSpaceExA(const_path &char, free_bytes_available_to_caller &u64, total_number_of_bytes &u64, total_number_of_free_bytes &u64) bool
diff --git a/vlib/os/os.c.v b/vlib/os/os.c.v
index a5ebd2ee062373..ec74beb3bcea02 100644
--- a/vlib/os/os.c.v
+++ b/vlib/os/os.c.v
@@ -1073,3 +1073,10 @@ pub fn error_win32(e SystemError) IError {
 		panic('Win32 API not available on this platform.')
 	}
 }
+
+pub struct DiskUsage {
+pub:
+	total     u64
+	available u64
+	used      u64
+}
diff --git a/vlib/os/os_nix.c.v b/vlib/os/os_nix.c.v
index 681b88e98cbdbf..194e0233a0211c 100644
--- a/vlib/os/os_nix.c.v
+++ b/vlib/os/os_nix.c.v
@@ -7,6 +7,7 @@ import strings
 #include <fcntl.h>
 #include <sys/utsname.h>
 #include <sys/types.h>
+#include <sys/statvfs.h>
 #include <utime.h>
 
 // path_separator is the platform specific separator string, used between the folders
@@ -530,3 +531,31 @@ fn C.sysconf(name int) i64
 pub fn page_size() int {
 	return int(C.sysconf(C._SC_PAGESIZE))
 }
+
+struct C.statvfs {
+	f_bsize  usize
+	f_blocks usize
+	f_bfree  usize
+	f_bavail usize
+}
+
+// disk_usage returns disk usage of `path`
+@[manualfree]
+pub fn disk_usage(path string) !DiskUsage {
+	mpath := if path == '' { '.' } else { path }
+	defer { unsafe { mpath.free() } }
+	mut vfs := C.statvfs{}
+	ret := unsafe { C.statvfs(&char(mpath.str), &vfs) }
+	if ret == -1 {
+		return error('cannot get disk usage of path')
+	}
+	f_bsize := u64(vfs.f_bsize)
+	f_blocks := u64(vfs.f_blocks)
+	f_bavail := u64(vfs.f_bavail)
+	f_bfree := u64(vfs.f_bfree)
+	return DiskUsage{
+		total:     f_bsize * f_blocks
+		available: f_bsize * f_bavail
+		used:      f_bsize * (f_blocks - f_bfree)
+	}
+}
diff --git a/vlib/os/os_test.c.v b/vlib/os/os_test.c.v
index 7d8e155306fe73..3ea7ee401dae4c 100644
--- a/vlib/os/os_test.c.v
+++ b/vlib/os/os_test.c.v
@@ -1110,3 +1110,10 @@ fn test_mkdir_at_file_dst() {
 	}
 	assert false
 }
+
+fn test_disk_usage() {
+	usage := os.disk_usage('.')!
+	assert usage.total > 0
+	assert usage.available > 0
+	assert usage.used > 0
+}
diff --git a/vlib/os/os_windows.c.v b/vlib/os/os_windows.c.v
index e36a30d4e52486..06f35107e13ad4 100644
--- a/vlib/os/os_windows.c.v
+++ b/vlib/os/os_windows.c.v
@@ -601,3 +601,26 @@ pub fn page_size() int {
 	C.GetSystemInfo(&sinfo)
 	return int(sinfo.dwPageSize)
 }
+
+// disk_usage returns disk usage of `path`
+pub fn disk_usage(path string) !DiskUsage {
+	mut free_bytes_available_to_caller := u64(0)
+	mut total := u64(0)
+	mut available := u64(0)
+	mut ret := false
+	if path == '.' || path == '' {
+		ret = C.GetDiskFreeSpaceExA(&char(0), &free_bytes_available_to_caller, &total,
+			&available)
+	} else {
+		ret = C.GetDiskFreeSpaceExA(&char(path.str), &free_bytes_available_to_caller,
+			&total, &available)
+	}
+	if ret == false {
+		return error('cannot get disk usage of path')
+	}
+	return DiskUsage{
+		total:     total
+		available: available
+		used:      total - available
+	}
+}