vendor/num_cpus-1.7.0/src/lib.rs

   1 //! # num_cpus
   2 //!
   3 //! A crate with utilities to determine the number of CPUs available on the
   4 //! current system.
   5 //!
   6 //! ## Example
   7 //!
   8 //! ```
   9 //! let cpus = num_cpus::get();
  10 //! ```
  11 #![cfg_attr(test, deny(warnings))]
  12 #![deny(missing_docs)]
  13 #![doc(html_root_url = "https://docs.rs/num_cpus/1.7.0")]
  14 #![allow(non_snake_case)]
  15
  16 #[cfg(not(windows))]
  17 extern crate libc;
  18
  19
  20 /// Returns the number of available CPUs of the current system.
  21 ///
  22 /// # Note
  23 ///
  24 /// This will check sched affinity on Linux.
  25 #[inline]
  26 pub fn get() -> usize {
  27     get_num_cpus()
  28 }
  29
  30 /// Returns the number of physical cores of the current system.
  31 ///
  32 /// If not possible on the particular architecture, returns same as `get()`
  33 /// which is the logical CPUs.
  34 #[inline]
  35 pub fn get_physical() -> usize {
  36     get_num_physical_cpus()
  37 }
  38
  39
  40 #[cfg(not(any(target_os = "linux", target_os = "windows", target_os="macos")))]
  41 #[inline]
  42 fn get_num_physical_cpus() -> usize {
  43     // Not implemented, fallback
  44     get_num_cpus()
  45 }
  46
  47 #[cfg(target_os = "windows")]
  48 fn get_num_physical_cpus() -> usize {
  49     match get_num_physical_cpus_windows() {
  50         Some(num) => num,
  51         None => get_num_cpus()
  52     }
  53 }
  54
  55 #[cfg(target_os = "windows")]
  56 fn get_num_physical_cpus_windows() -> Option<usize> {
  57     // Inspired by https://msdn.microsoft.com/en-us/library/ms683194
  58
  59     use std::ptr;
  60     use std::mem;
  61
  62     #[allow(non_upper_case_globals)]
  63     const RelationProcessorCore: u32 = 0;
  64
  65     #[repr(C)]
  66     #[allow(non_camel_case_types)]
  67     struct SYSTEM_LOGICAL_PROCESSOR_INFORMATION {
  68         mask: usize,
  69         relationship: u32,
  70         _unused: [u64; 2]
  71     }
  72
  73     extern "system" {
  74         fn GetLogicalProcessorInformation(
  75             info: *mut SYSTEM_LOGICAL_PROCESSOR_INFORMATION,
  76             length: &mut u32
  77         ) -> u32;
  78     }
  79
  80     // First we need to determine how much space to reserve.
  81
  82     // The required size of the buffer, in bytes.
  83     let mut needed_size = 0;
  84
  85     unsafe {
  86         GetLogicalProcessorInformation(ptr::null_mut(), &mut needed_size);
  87     }
  88
  89     let struct_size = mem::size_of::<SYSTEM_LOGICAL_PROCESSOR_INFORMATION>() as u32;
  90
  91     // Could be 0, or some other bogus size.
  92     if needed_size == 0 || needed_size < struct_size || needed_size % struct_size != 0 {
  93         return None;
  94     }
  95
  96     let count = needed_size / struct_size;
  97
  98     // Allocate some memory where we will store the processor info.
  99     let mut buf = Vec::with_capacity(count as usize);
 100
 101     let result;
 102
 103     unsafe {
 104         result = GetLogicalProcessorInformation(buf.as_mut_ptr(), &mut needed_size);
 105     }
 106
 107     // Failed for any reason.
 108     if result == 0 {
 109         return None;
 110     }
 111
 112     let count = needed_size / struct_size;
 113
 114     unsafe {
 115         buf.set_len(count as usize);
 116     }
 117
 118     let phys_proc_count = buf.iter()
 119         // Only interested in processor packages (physical processors.)
 120         .filter(|proc_info| proc_info.relationship == RelationProcessorCore)
 121         .count();
 122
 123     if phys_proc_count == 0 {
 124         None
 125     } else {
 126         Some(phys_proc_count)
 127     }
 128 }
 129
 130 #[cfg(target_os = "linux")]
 131 fn get_num_physical_cpus() -> usize {
 132     use std::io::BufReader;
 133     use std::io::BufRead;
 134     use std::fs::File;
 135     use std::collections::HashSet;
 136
 137     let file = match File::open("/proc/cpuinfo") {
 138         Ok(val) => val,
 139         Err(_) => {return get_num_cpus()},
 140     };
 141     let reader = BufReader::new(file);
 142     let mut set = HashSet::new();
 143     let mut coreid: u32 = 0;
 144     let mut physid: u32 = 0;
 145     let mut chgcount = 0;
 146     for line in reader.lines().filter_map(|result| result.ok()) {
 147         let parts: Vec<&str> = line.split(':').map(|s| s.trim()).collect();
 148         if parts.len() != 2 {
 149             continue
 150         }
 151         if parts[0] == "core id" || parts[0] == "physical id" {
 152             let value = match parts[1].trim().parse() {
 153               Ok(val) => val,
 154               Err(_) => break,
 155             };
 156             match parts[0] {
 157                 "core id"     => coreid = value,
 158                 "physical id" => physid = value,
 159                 _ => {},
 160             }
 161             chgcount += 1;
 162         }
 163         if chgcount == 2 {
 164             set.insert((physid, coreid));
 165             chgcount = 0;
 166         }
 167     }
 168     let count = set.len();
 169     if count == 0 { get_num_cpus() } else { count }
 170 }
 171
 172 #[cfg(windows)]
 173 fn get_num_cpus() -> usize {
 174     #[repr(C)]
 175     struct SYSTEM_INFO {
 176         wProcessorArchitecture: u16,
 177         wReserved: u16,
 178         dwPageSize: u32,
 179         lpMinimumApplicationAddress: *mut u8,
 180         lpMaximumApplicationAddress: *mut u8,
 181         dwActiveProcessorMask: *mut u8,
 182         dwNumberOfProcessors: u32,
 183         dwProcessorType: u32,
 184         dwAllocationGranularity: u32,
 185         wProcessorLevel: u16,
 186         wProcessorRevision: u16,
 187     }
 188
 189     extern "system" {
 190         fn GetSystemInfo(lpSystemInfo: *mut SYSTEM_INFO);
 191     }
 192
 193     unsafe {
 194         let mut sysinfo: SYSTEM_INFO = std::mem::uninitialized();
 195         GetSystemInfo(&mut sysinfo);
 196         sysinfo.dwNumberOfProcessors as usize
 197     }
 198 }
 199
 200 #[cfg(any(target_os = "freebsd",
 201           target_os = "dragonfly",
 202           target_os = "bitrig",
 203           target_os = "netbsd"))]
 204 fn get_num_cpus() -> usize {
 205     let mut cpus: libc::c_uint = 0;
 206     let mut cpus_size = std::mem::size_of_val(&cpus);
 207
 208     unsafe {
 209         cpus = libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as libc::c_uint;
 210     }
 211     if cpus < 1 {
 212         let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0];
 213         unsafe {
 214             libc::sysctl(mib.as_mut_ptr(),
 215                          2,
 216                          &mut cpus as *mut _ as *mut _,
 217                          &mut cpus_size as *mut _ as *mut _,
 218                          0 as *mut _,
 219                          0);
 220         }
 221         if cpus < 1 {
 222             cpus = 1;
 223         }
 224     }
 225     cpus as usize
 226 }
 227
 228 #[cfg(target_os = "openbsd")]
 229 fn get_num_cpus() -> usize {
 230     let mut cpus: libc::c_uint = 0;
 231     let mut cpus_size = std::mem::size_of_val(&cpus);
 232     let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0];
 233
 234     unsafe {
 235         libc::sysctl(mib.as_mut_ptr(),
 236                      2,
 237                      &mut cpus as *mut _ as *mut _,
 238                      &mut cpus_size as *mut _ as *mut _,
 239                      0 as *mut _,
 240                      0);
 241     }
 242     if cpus < 1 {
 243         cpus = 1;
 244     }
 245     cpus as usize
 246 }
 247
 248
 249 #[cfg(target_os = "macos")]
 250 fn get_num_physical_cpus() -> usize {
 251     use std::ffi::CStr;
 252     use std::ptr;
 253
 254     let mut cpus: i32 = 0;
 255     let mut cpus_size = std::mem::size_of_val(&cpus);
 256
 257     let sysctl_name = CStr::from_bytes_with_nul(b"hw.physicalcpu\0")
 258         .expect("byte literal is missing NUL");
 259
 260     unsafe {
 261         if 0 != libc::sysctlbyname(sysctl_name.as_ptr(),
 262                                    &mut cpus as *mut _ as *mut _,
 263                                    &mut cpus_size as *mut _ as *mut _,
 264                                    ptr::null_mut(),
 265                                    0) {
 266             return get_num_cpus();
 267         }
 268     }
 269     cpus as usize
 270 }
 271
 272 #[cfg(target_os = "linux")]
 273 fn get_num_cpus() -> usize {
 274     let mut set:  libc::cpu_set_t = unsafe { std::mem::zeroed() };
 275     if unsafe { libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
 276         let mut count: u32 = 0;
 277         for i in 0..libc::CPU_SETSIZE as usize {
 278             if unsafe { libc::CPU_ISSET(i, &set) } {
 279                 count += 1
 280             }
 281         }
 282         count as usize
 283     } else {
 284         let cpus = unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) };
 285         if cpus < 1 {
 286             1
 287         } else {
 288             cpus as usize
 289         }
 290     }
 291 }
 292
 293 #[cfg(any(
 294     target_os = "nacl",
 295     target_os = "macos",
 296     target_os = "ios",
 297     target_os = "android",
 298     target_os = "solaris",
 299     target_os = "fuchsia")
 300 )]
 301 fn get_num_cpus() -> usize {
 302     // On ARM targets, processors could be turned off to save power.
 303     // Use `_SC_NPROCESSORS_CONF` to get the real number.
 304     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
 305     const CONF_NAME: libc::c_int = libc::_SC_NPROCESSORS_CONF;
 306     #[cfg(not(any(target_arch = "arm", target_arch = "aarch64")))]
 307     const CONF_NAME: libc::c_int = libc::_SC_NPROCESSORS_ONLN;
 308
 309     let cpus = unsafe { libc::sysconf(CONF_NAME) };
 310     if cpus < 1 {
 311         1
 312     } else {
 313         cpus as usize
 314     }
 315 }
 316
 317 #[cfg(any(target_os = "emscripten", target_os = "redox", target_os = "haiku"))]
 318 fn get_num_cpus() -> usize {
 319     1
 320 }
 321
 322 #[cfg(test)]
 323 mod tests {
 324     fn env_var(name: &'static str) -> Option<usize> {
 325         ::std::env::var(name).ok().map(|val| val.parse().unwrap())
 326     }
 327
 328     #[test]
 329     fn test_get() {
 330         let num = super::get();
 331         if let Some(n) = env_var("NUM_CPUS_TEST_GET") {
 332             assert_eq!(num, n);
 333         } else {
 334             assert!(num > 0);
 335             assert!(num < 236_451);
 336         }
 337     }
 338
 339     #[test]
 340     fn test_get_physical() {
 341         let num = super::get_physical();
 342         if let Some(n) = env_var("NUM_CPUS_TEST_GET_PHYSICAL") {
 343             assert_eq!(num, n);
 344         } else {
 345             assert!(num > 0);
 346             assert!(num < 236_451);
 347         }
 348     }
 349 }