From 03e954560e9efa3eabfa4c7fef4d7d2a0213d454 Mon Sep 17 00:00:00 2001
From: Michael Zhang <mail@mzhang.io>
Date: Mon, 21 Oct 2024 10:33:44 -0500
Subject: [PATCH] determine windows free disk size

---
 .gitignore         |  3 ++-
 Cargo.lock         |  1 +
 src/rpc/Cargo.toml |  9 ++++---
 src/rpc/system.rs  | 60 ++++++++++++++++++++++++++++++++++++----------
 4 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/.gitignore b/.gitignore
index ef7a56eb..c8bd9441 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,5 @@
 /pki
 **/*.rs.bk
 *.swp
-/.direnv
\ No newline at end of file
+/.direnv
+Packet.lib
\ No newline at end of file
diff --git a/Cargo.lock b/Cargo.lock
index fa313874..fa9b4375 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1544,6 +1544,7 @@ dependencies = [
  "tokio",
  "tokio-stream",
  "tracing",
+ "winapi",
 ]
 
 [[package]]
diff --git a/src/rpc/Cargo.toml b/src/rpc/Cargo.toml
index acde0911..aceb9505 100644
--- a/src/rpc/Cargo.toml
+++ b/src/rpc/Cargo.toml
@@ -51,7 +51,10 @@ tokio.workspace = true
 tokio-stream.workspace = true
 opentelemetry.workspace = true
 
+[target.'cfg(windows)'.dependencies]
+winapi = { version = "0.3.9", features = ["fileapi", "impl-default"] }
+
 [features]
-kubernetes-discovery = [ "kube", "k8s-openapi", "schemars" ]
-consul-discovery = [ "reqwest", "err-derive" ]
-system-libs = [ "sodiumoxide/use-pkg-config" ]
+kubernetes-discovery = ["kube", "k8s-openapi", "schemars"]
+consul-discovery = ["reqwest", "err-derive"]
+system-libs = ["sodiumoxide/use-pkg-config"]
diff --git a/src/rpc/system.rs b/src/rpc/system.rs
index d49bec8f..c0958736 100644
--- a/src/rpc/system.rs
+++ b/src/rpc/system.rs
@@ -817,22 +817,58 @@ impl NodeStatus {
 		}
 	}
 
+	#[cfg(windows)]
 	fn update_disk_usage(&mut self, meta_dir: &Path, data_dir: &DataDirEnum) {
-		#[cfg(unix)]
-		let mount_avail = {
-			use nix::sys::statvfs::statvfs;
-			|path: &Path| match statvfs(path) {
-				Ok(x) => {
-					let avail = x.blocks_available() as u64 * x.fragment_size() as u64;
-					let total = x.blocks() as u64 * x.fragment_size() as u64;
-					Some((x.filesystem_id(), avail, total))
-				}
-				Err(_) => None,
+		use winapi::um::fileapi::GetDiskFreeSpaceExA;
+		use winapi::um::winnt::ULARGE_INTEGER;
+
+		let mount_avail = |path: &Path| -> Option<(u64, u64)> {
+			let mut path = path.to_path_buf();
+			path.push(""); // Ensure trailing slash
+
+			let mut a: ULARGE_INTEGER = Default::default();
+			let mut total: ULARGE_INTEGER = Default::default();
+			let mut free: ULARGE_INTEGER = Default::default();
+
+			let path_ptr = path.as_os_str().as_encoded_bytes().as_ptr();
+			let result = unsafe {
+				GetDiskFreeSpaceExA(path_ptr as *const i8, &mut a, &mut total, &mut free)
+			};
+
+			if result == 0 {
+				return None;
 			}
+
+			let free = unsafe { *free.QuadPart() };
+			let total = unsafe { *total.QuadPart() };
+
+			Some((free, total))
 		};
 
-		#[cfg(windows)]
-		let mount_avail = |_path: &Path| None::<(u64, _, _)>;
+		self.meta_disk_avail = mount_avail(meta_dir);
+		self.data_disk_avail = match data_dir {
+			DataDirEnum::Single(path_buf) => mount_avail(path_buf),
+
+			// TODO: THIS IS WRONG!! Does not take into account multiple dirs on the same partition
+			// Will have to deduplicate by the filesystem mount path
+			DataDirEnum::Multiple(dirs) => dirs
+				.into_iter()
+				.filter_map(|dir| mount_avail(&dir.path))
+				.reduce(|(a1, b1), (a2, b2)| (a1 + a2, b1 + b2)),
+		};
+	}
+
+	#[cfg(unix)]
+	fn update_disk_usage(&mut self, meta_dir: &Path, data_dir: &DataDirEnum) {
+		use nix::sys::statvfs::statvfs;
+		let mount_avail = |path: &Path| match statvfs(path) {
+			Ok(x) => {
+				let avail = x.blocks_available() as u64 * x.fragment_size() as u64;
+				let total = x.blocks() as u64 * x.fragment_size() as u64;
+				Some((x.filesystem_id(), avail, total))
+			}
+			Err(_) => None,
+		};
 
 		self.meta_disk_avail = mount_avail(meta_dir).map(|(_, a, t)| (a, t));
 		self.data_disk_avail = match data_dir {