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 libc::EINVAL;
6 use winapi::um::processthreadsapi::GetCurrentThread;
7 use winapi::um::winbase::SetThreadAffinityMask;
8
9 use super::errno_result;
10 use super::Error;
11 use super::Result;
12
13 /// Set the CPU affinity of the current thread to a given set of CPUs.
14 /// The cpus must be a subset of those in the process affinity mask.
15 /// If there are more than 64 processors, the affinity mask must specify
16 /// processors in the thread's current processor group.
17 ///
18 /// Returns the previous bits set for cpu_affinity.
19 /// # Examples
20 ///
21 /// Set the calling thread's CPU affinity so it will run on only CPUs
22 /// 0, 1, 5, and 6.
23 ///
24 /// ```
25 /// # use base::windows::set_cpu_affinity;
26 /// set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
27 /// ```
set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<usize>28 pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<usize> {
29 let mut affinity_mask: usize = 0;
30 for cpu in cpus {
31 if cpu >= 64 {
32 return Err(Error::new(EINVAL));
33 }
34 affinity_mask |= 1 << cpu;
35 }
36 set_cpu_affinity_mask(affinity_mask)
37 }
38
set_cpu_affinity_mask(affinity_mask: usize) -> Result<usize>39 pub fn set_cpu_affinity_mask(affinity_mask: usize) -> Result<usize> {
40 // SAFETY: trivially safe as return value is checked.
41 let res: usize = unsafe {
42 let thread_handle = GetCurrentThread();
43 SetThreadAffinityMask(thread_handle, affinity_mask)
44 };
45 if res == 0 {
46 return errno_result::<usize>();
47 }
48 Ok(res)
49 }
50
get_cpu_affinity() -> Result<Vec<usize>>51 pub fn get_cpu_affinity() -> Result<Vec<usize>> {
52 // Unimplemented for now, since this is unused on Windows. Thread affinity can be read by
53 // calling GetThreadAffinityMask twice, to get and restore the previous affinities.
54 Err(Error::new(EINVAL))
55 }
56
57 #[cfg(test)]
58 mod tests {
59 use winapi::um::processthreadsapi::GetCurrentProcess;
60 use winapi::um::winbase::GetProcessAffinityMask;
61
62 use super::super::sched::*;
63 #[test]
cpu_affinity()64 fn cpu_affinity() {
65 let mut process_affinity_mask: usize = 0;
66 let mut system_affinity_mask: usize = 0;
67 // SAFETY: trivially safe as return value is checked.
68 let res = unsafe {
69 GetProcessAffinityMask(
70 GetCurrentProcess(),
71 &mut process_affinity_mask as *mut usize,
72 &mut system_affinity_mask as *mut usize,
73 )
74 };
75 assert_ne!(res, 0);
76 let mut prev_thread_affinity = set_cpu_affinity(vec![0, 1]).unwrap();
77 assert_eq!(process_affinity_mask, prev_thread_affinity);
78 // reset to original process affinity and grab previous affinity.
79 prev_thread_affinity = set_cpu_affinity_mask(process_affinity_mask).unwrap();
80 // cpu 0 + cpu 1, means bit 0 and bit 1 are set (1 + 2 = 3). Verify that is the case.
81 assert_eq!(prev_thread_affinity, 3);
82 }
83 }
84