xref: /aosp_15_r20/external/crosvm/base/src/sys/windows/priority.rs (revision bb4ee6a4ae7042d18b07a98463b9c8b875e44b39)
1 // Copyright 2022 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 use std::os::windows::raw::HANDLE;
6 
7 use log::warn;
8 use winapi::shared::minwindef::FALSE;
9 use winapi::um::avrt::AvRevertMmThreadCharacteristics;
10 use winapi::um::avrt::AvSetMmThreadCharacteristicsA;
11 use winapi::um::errhandlingapi::GetLastError;
12 use winapi::um::processthreadsapi::GetCurrentThread;
13 use winapi::um::processthreadsapi::SetThreadPriority;
14 
15 use super::errno_result;
16 use super::Result;
17 
set_audio_thread_priority() -> Result<SafeMultimediaHandle>18 pub fn set_audio_thread_priority() -> Result<SafeMultimediaHandle> {
19     // SAFETY:
20     // Safe because we know Pro Audio is part of windows and we down task_index.
21     let multimedia_handle = unsafe {
22         let mut task_index: u32 = 0;
23         // "Pro Audio" is defined in:
24         // HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows
25         // NT\CurrentVersion\Multimedia\SystemProfile\Tasks\Pro Audio
26         let pro_audio = std::ffi::CString::new("Pro Audio").unwrap();
27         AvSetMmThreadCharacteristicsA(pro_audio.as_ptr(), &mut task_index)
28     };
29 
30     if multimedia_handle.is_null() {
31         warn!(
32             "Failed to set audio thread to Pro Audio. Error: {}",
33             // SAFETY: trivially safe
34             unsafe { GetLastError() }
35         );
36         errno_result()
37     } else {
38         Ok(SafeMultimediaHandle { multimedia_handle })
39     }
40 }
41 
set_thread_priority(thread_priority: i32) -> Result<()>42 pub fn set_thread_priority(thread_priority: i32) -> Result<()> {
43     let res =
44     // SAFETY:
45         // Safe because priority level value is valid and a valid thread handle will be passed in
46         unsafe { SetThreadPriority(GetCurrentThread(), thread_priority) };
47     if res == 0 {
48         errno_result()
49     } else {
50         Ok(())
51     }
52 }
53 
54 pub struct SafeMultimediaHandle {
55     multimedia_handle: HANDLE,
56 }
57 
58 impl Drop for SafeMultimediaHandle {
drop(&mut self)59     fn drop(&mut self) {
60         // SAFETY:
61         // Safe because we `multimedia_handle` is defined in the same thread and is created in the
62         // function above. `multimedia_handle` needs be created from
63         // `AvSetMmThreadCharacteristicsA`. This will also drop the `mulitmedia_handle`.
64         if unsafe { AvRevertMmThreadCharacteristics(self.multimedia_handle) } == FALSE {
65             warn!(
66                 "Failed to revert audio thread. Error: {}",
67                 // SAFETY: trivially safe
68                 unsafe { GetLastError() }
69             );
70         }
71     }
72 }
73 
74 #[cfg(test)]
75 mod test {
76     use winapi::um::processthreadsapi::GetCurrentThread;
77     use winapi::um::processthreadsapi::GetThreadPriority;
78     use winapi::um::winbase::THREAD_PRIORITY_NORMAL;
79     use winapi::um::winbase::THREAD_PRIORITY_TIME_CRITICAL;
80 
81     use super::*;
82 
83     // TODO(b/223733375): Enable ignored flaky tests.
84     #[test]
85     #[ignore]
test_mm_handle_is_dropped()86     fn test_mm_handle_is_dropped() {
87         // SAFETY:
88         // Safe because the only the only unsafe functions called are to get the thread
89         // priority.
90         unsafe {
91             let thread_priority = GetThreadPriority(GetCurrentThread());
92             assert_eq!(thread_priority, THREAD_PRIORITY_NORMAL as i32);
93             {
94                 let _handle = set_audio_thread_priority();
95                 let thread_priority = GetThreadPriority(GetCurrentThread());
96                 assert_eq!(thread_priority, THREAD_PRIORITY_TIME_CRITICAL as i32);
97             }
98             let thread_priority = GetThreadPriority(GetCurrentThread());
99             assert_eq!(thread_priority, THREAD_PRIORITY_NORMAL as i32);
100         }
101     }
102 }
103