Skip to main content

diskann_platform/win/
perf.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5use std::mem::{self, zeroed};
6
7use windows_sys::Win32::{
8    Foundation::FILETIME,
9    System::{
10        ProcessStatus::{GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS},
11        SystemInformation::{GetSystemInfo, SYSTEM_INFO},
12        Threading::{GetCurrentProcess, GetProcessTimes, GetSystemTimes},
13        WindowsProgramming::QueryProcessCycleTime,
14    },
15};
16
17pub fn get_process_cycle_time() -> Option<u64> {
18    let mut cycle_time: u64 = 0;
19
20    // SAFETY: Call a Win32 API.
21    let handle = unsafe { GetCurrentProcess() };
22
23    let result = unsafe { QueryProcessCycleTime(handle, &mut cycle_time as *mut u64) } != 0;
24    if result {
25        return Some(cycle_time);
26    }
27    None
28}
29
30// Gets the process time in kernel and user modes.
31pub fn get_process_time() -> Option<u64> {
32    let handle = unsafe { GetCurrentProcess() };
33
34    let mut creation_time: FILETIME = FILETIME {
35        dwLowDateTime: 0,
36        dwHighDateTime: 0,
37    };
38    let mut exit_time: FILETIME = FILETIME {
39        dwLowDateTime: 0,
40        dwHighDateTime: 0,
41    };
42    let mut kernel_time: FILETIME = FILETIME {
43        dwLowDateTime: 0,
44        dwHighDateTime: 0,
45    };
46    let mut user_time: FILETIME = FILETIME {
47        dwLowDateTime: 0,
48        dwHighDateTime: 0,
49    };
50
51    let result = unsafe {
52        GetProcessTimes(
53            handle,
54            &mut creation_time,
55            &mut exit_time,
56            &mut kernel_time,
57            &mut user_time,
58        )
59    };
60
61    if result != 0 {
62        let kernel_time = filetime_to_u64(kernel_time);
63        let user_time = filetime_to_u64(user_time);
64        return Some(kernel_time + user_time);
65    }
66
67    None
68}
69
70// Gets the system time in kernel and user modes.
71pub fn get_system_time() -> Option<u64> {
72    let mut idle_time: FILETIME = FILETIME {
73        dwLowDateTime: 0,
74        dwHighDateTime: 0,
75    };
76    let mut kernel_time: FILETIME = FILETIME {
77        dwLowDateTime: 0,
78        dwHighDateTime: 0,
79    };
80    let mut user_time: FILETIME = FILETIME {
81        dwLowDateTime: 0,
82        dwHighDateTime: 0,
83    };
84
85    let result = unsafe { GetSystemTimes(&mut idle_time, &mut kernel_time, &mut user_time) };
86
87    if result != 0 {
88        let kernel_time = filetime_to_u64(kernel_time);
89        let user_time = filetime_to_u64(user_time);
90        return Some(kernel_time + user_time);
91    }
92
93    None
94}
95
96// Gets the number of processors.
97pub fn get_number_of_processors() -> Option<u64> {
98    let mut system_info: SYSTEM_INFO = unsafe { zeroed() };
99    unsafe { GetSystemInfo(&mut system_info) };
100    Some(system_info.dwNumberOfProcessors as u64)
101}
102
103#[inline(always)]
104const fn filetime_to_u64(ft: FILETIME) -> u64 {
105    ((ft.dwHighDateTime as u64) << 32) | (ft.dwLowDateTime as u64)
106}
107
108pub fn get_peak_workingset_size() -> Option<u64> {
109    // SAFETY: Call a Win32 API.
110    let handle = unsafe { GetCurrentProcess() };
111
112    let mut counters: PROCESS_MEMORY_COUNTERS = unsafe { mem::zeroed() };
113    counters.cb = mem::size_of::<PROCESS_MEMORY_COUNTERS>() as u32;
114
115    let result = unsafe {
116        GetProcessMemoryInfo(
117            handle,
118            &mut counters as *mut PROCESS_MEMORY_COUNTERS,
119            counters.cb,
120        )
121    };
122
123    if result != 0 {
124        let peak_working_set_size = counters.PeakWorkingSetSize;
125        return Some(peak_working_set_size as u64);
126    }
127    None
128}