1extern crate lazy_static;
2extern crate num_cpus;
3use super::bithacks::round_up_to_power2;
4
5const DEFAULT_CPUS: u32 = 4;
6
7pub struct CpuShard(u32, u32);
11
12impl CpuShard {
13 #[inline]
14 pub fn new(max_shard_limit: Option<u32>) -> Self {
15 let mut cpus = num_cpus::get() as u32;
16 if cpus < DEFAULT_CPUS {
17 cpus = DEFAULT_CPUS;
18 }
19 if let Some(limit) = max_shard_limit {
20 if cpus > limit {
21 cpus = limit;
22 }
23 }
24 let (shard, shift) = round_up_to_power2(cpus);
25 Self(shard, shift)
26 }
27
28 #[inline]
29 pub fn shards(&self) -> u32 {
30 self.0
31 }
32
33 #[inline]
34 pub fn shift(&self) -> u32 {
35 self.1
36 }
37}
38
39#[cfg(test)]
40mod tests {
41 use super::*;
42
43 #[test]
44 fn test_static_shards() {
45 let cpu_shard = CpuShard::new(None);
46 let cpus = num_cpus::get();
47 println!("cpus={}", cpus);
48 if cpus <= 4 {
49 assert_eq!(cpu_shard.shards(), 4);
50 assert_eq!(cpu_shard.shift(), 2);
51 } else if cpus <= 8 {
52 assert_eq!(cpu_shard.shards(), 8);
53 assert_eq!(cpu_shard.shift(), 3);
54 }
55 }
56
57 #[test]
58 fn test_shards_limit() {
59 let cpu_shard = CpuShard::new(Some(2));
60 assert_eq!(cpu_shard.shards(), 2);
61 assert_eq!(cpu_shard.shift(), 1);
62 }
63}