From 2c1ca98abaf6756d3246841cc3cf4ca3f617efca Mon Sep 17 00:00:00 2001
From: Jonas Schievink <jonasschievink@gmail.com>
Date: Wed, 9 Jun 2021 22:55:50 +0200
Subject: [PATCH] Make `MemoryUsage` work on Windows

---
 Cargo.lock                         |  1 +
 crates/profile/Cargo.toml          |  3 +++
 crates/profile/src/memory_usage.rs | 16 ++++++++++++++++
 3 files changed, 20 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 505263c645..2bd8897f05 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1170,6 +1170,7 @@ dependencies = [
  "once_cell",
  "perf-event",
  "tikv-jemalloc-ctl",
+ "winapi",
 ]
 
 [[package]]
diff --git a/crates/profile/Cargo.toml b/crates/profile/Cargo.toml
index 1a8c8f862c..653d3d9839 100644
--- a/crates/profile/Cargo.toml
+++ b/crates/profile/Cargo.toml
@@ -20,6 +20,9 @@ jemalloc-ctl = { version = "0.4.1", package = "tikv-jemalloc-ctl", optional = tr
 [target.'cfg(target_os = "linux")'.dependencies]
 perf-event = "0.4"
 
+[target.'cfg(windows)'.dependencies]
+winapi = { version = "0.3.8", features = ["psapi"] }
+
 [features]
 cpu_profiler = []
 jemalloc = ["jemalloc-ctl"]
diff --git a/crates/profile/src/memory_usage.rs b/crates/profile/src/memory_usage.rs
index 2917ded60f..6ef58c9c14 100644
--- a/crates/profile/src/memory_usage.rs
+++ b/crates/profile/src/memory_usage.rs
@@ -35,6 +35,22 @@ impl MemoryUsage {
                 // Note: This is incredibly slow.
                 let alloc = unsafe { libc::mallinfo() }.uordblks as isize;
                 MemoryUsage { allocated: Bytes(alloc) }
+            } else if #[cfg(windows)] {
+                // There doesn't seem to be an API for determining heap usage, so we try to
+                // approximate that by using the Commit Charge value.
+
+                use winapi::um::processthreadsapi::*;
+                use winapi::um::psapi::*;
+                use std::mem::{MaybeUninit, size_of};
+
+                let proc = unsafe { GetCurrentProcess() };
+                let mut mem_counters = MaybeUninit::uninit();
+                let cb = size_of::<PROCESS_MEMORY_COUNTERS>();
+                let ret = unsafe { GetProcessMemoryInfo(proc, mem_counters.as_mut_ptr(), cb as u32) };
+                assert!(ret != 0);
+
+                let usage = unsafe { mem_counters.assume_init().PagefileUsage };
+                MemoryUsage { allocated: Bytes(usage as isize) }
             } else {
                 MemoryUsage { allocated: Bytes(0) }
             }