sync_utils/
cpu.rs

1extern crate lazy_static;
2extern crate num_cpus;
3use super::bithacks::round_up_to_power2;
4
5const DEFAULT_CPUS: u32 = 4;
6
7/// Determine the right sharding by cpu.
8/// Minimise the cost of hashing function of remainder calculation.
9/// Refer to `bithacks` module.
10pub 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}