From 82c0c3b0e5561008255b7f023e7eee43db593946 Mon Sep 17 00:00:00 2001
From: Clement Tsang <34804052+ClementTsang@users.noreply.github.com>
Date: Wed, 8 Mar 2023 00:08:02 -0500
Subject: [PATCH] refactor: use internal committed memory calculation for
 Windows (#1051)

* refactor: use internal committed memory calculation for Windows

* comments
---
 Cargo.toml                               |  3 ++-
 src/app/data_harvester.rs                |  5 +++-
 src/app/data_harvester/memory.rs         | 12 ++++++++-
 src/app/data_harvester/memory/sysinfo.rs |  1 +
 src/app/data_harvester/memory/windows.rs | 32 ++++++++++++++++++++++++
 5 files changed, 50 insertions(+), 3 deletions(-)
 create mode 100644 src/app/data_harvester/memory/windows.rs

diff --git a/Cargo.toml b/Cargo.toml
index 6a721237..0b9c7bf1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -119,8 +119,9 @@ mach2 = "0.4.1"
 [target.'cfg(target_os = "windows")'.dependencies]
 heim = { version = "0.1.0-rc.1", features = ["disk"] }
 windows = { version = "0.44.0", features = [
-    "Win32_System_Threading",
     "Win32_Foundation",
+    "Win32_System_ProcessStatus",
+    "Win32_System_Threading",
 ] }
 
 [target.'cfg(target_os = "freebsd")'.dependencies]
diff --git a/src/app/data_harvester.rs b/src/app/data_harvester.rs
index 9d8a02a1..0f1eee0a 100644
--- a/src/app/data_harvester.rs
+++ b/src/app/data_harvester.rs
@@ -397,7 +397,10 @@ impl DataCollector {
     fn update_memory_usage(&mut self) {
         if self.widgets_to_harvest.use_mem {
             self.data.memory = memory::get_ram_usage(&self.sys);
-            self.data.swap = memory::get_swap_usage(&self.sys);
+            self.data.swap = memory::get_swap_usage(
+                #[cfg(not(target_os = "windows"))]
+                &self.sys,
+            );
 
             #[cfg(feature = "zfs")]
             {
diff --git a/src/app/data_harvester/memory.rs b/src/app/data_harvester/memory.rs
index 18ac94cd..02c0fa9c 100644
--- a/src/app/data_harvester/memory.rs
+++ b/src/app/data_harvester/memory.rs
@@ -1,7 +1,17 @@
 //! Memory data collection.
 
 pub mod sysinfo;
-pub(crate) use self::sysinfo::{get_ram_usage, get_swap_usage};
+pub(crate) use self::sysinfo::get_ram_usage;
+
+cfg_if::cfg_if! {
+    if #[cfg(target_os = "windows")] {
+        pub mod windows;
+        pub(crate) use self::windows::get_swap_usage;
+    } else {
+        pub(crate) use self::sysinfo::get_swap_usage;
+
+    }
+}
 
 #[cfg(feature = "gpu")]
 pub mod gpu;
diff --git a/src/app/data_harvester/memory/sysinfo.rs b/src/app/data_harvester/memory/sysinfo.rs
index d0c5c8ca..4cb8bdb0 100644
--- a/src/app/data_harvester/memory/sysinfo.rs
+++ b/src/app/data_harvester/memory/sysinfo.rs
@@ -21,6 +21,7 @@ pub(crate) fn get_ram_usage(sys: &System) -> Option<MemHarvest> {
 }
 
 /// Returns SWAP usage.
+#[cfg(not(target_os = "windows"))]
 pub(crate) fn get_swap_usage(sys: &System) -> Option<MemHarvest> {
     let mem_used_in_kib = sys.used_swap() / 1024;
     let mem_total_in_kib = sys.total_swap() / 1024;
diff --git a/src/app/data_harvester/memory/windows.rs b/src/app/data_harvester/memory/windows.rs
new file mode 100644
index 00000000..5ead4b52
--- /dev/null
+++ b/src/app/data_harvester/memory/windows.rs
@@ -0,0 +1,32 @@
+use std::mem::{size_of, zeroed};
+use windows::Win32::Foundation::TRUE;
+use windows::Win32::System::ProcessStatus::{K32GetPerformanceInfo, PERFORMANCE_INFORMATION};
+
+use crate::data_harvester::memory::MemHarvest;
+
+// TODO: Note this actually calculates the total *committed* usage. Rename and change label for accuracy!
+/// Get the committed memory usage.
+///
+/// Code based on [sysinfo's](https://github.com/GuillaumeGomez/sysinfo/blob/6f8178495adcf3ca4696a9ec548586cf6a621bc8/src/windows/system.rs#L169).
+pub(crate) fn get_swap_usage() -> Option<MemHarvest> {
+    // SAFETY: The safety invariant is that we only touch what's in `perf_info` if it succeeds, and that
+    // the bindings are "safe" to use with how we call them.
+    unsafe {
+        let mut perf_info: PERFORMANCE_INFORMATION = zeroed();
+        if K32GetPerformanceInfo(&mut perf_info, size_of::<PERFORMANCE_INFORMATION>() as u32)
+            == TRUE
+        {
+            // Saturating sub by perf_info.PhysicalTotal for what sysinfo does.
+            let swap_total = perf_info.PageSize.saturating_mul(perf_info.CommitLimit) as u64;
+            let swap_used = perf_info.PageSize.saturating_mul(perf_info.CommitTotal) as u64;
+
+            Some(MemHarvest {
+                total_kib: swap_total / 1024,
+                used_kib: swap_used / 1024,
+                use_percent: Some(swap_used as f64 / swap_total as f64 * 100.0),
+            })
+        } else {
+            None
+        }
+    }
+}