light_cache/map/
builder.rs

1use super::{max_parrellism, LightMap, Shard};
2
3use hashbrown::raw::RawTable;
4use hashbrown::HashMap;
5
6use std::hash::BuildHasher;
7
8use parking_lot::{Mutex, RwLock};
9
10pub struct MapBuilder {
11    pub(crate) shards: Option<usize>,
12    pub(crate) estimated_size: Option<usize>,
13}
14
15impl MapBuilder {
16    pub fn new() -> Self {
17        MapBuilder {
18            shards: None,
19            estimated_size: None,
20        }
21    }
22
23    /// Set the number of shards
24    /// 
25    /// This will be rounded up to the next power of two, but by default it will be 4 times the number of cores
26    /// since this is the number of possible parrallel operations
27    pub fn shards(mut self, shards: usize) -> Self {
28        self.shards = Some(shards);
29        self
30    }
31
32    /// Set the estimated size of the map
33    /// 
34    /// Used to preallocate the map
35    pub fn estimated_size(mut self, estimated_size: usize) -> Self {
36        self.estimated_size = Some(estimated_size);
37        self
38    }
39
40    pub fn build<K, V, S: BuildHasher>(self, build_hasher: S) -> LightMap<K, V, S> {
41        let shards = self
42            .shards
43            .unwrap_or_else(|| max_parrellism() * 4)
44            .next_power_of_two();
45
46        if let Some(estimated_size) = self.estimated_size {
47            if estimated_size > shards {
48                let per_shard = (estimated_size / shards) * 2;
49
50                let shards = (0..shards)
51                    .map(|_| Shard {
52                        waiters: Mutex::new(HashMap::new()),
53                        table: RwLock::new(RawTable::with_capacity(per_shard)),
54                    })
55                    .collect();
56
57                return LightMap {
58                    shards,
59                    build_hasher,
60                };
61            }
62        }
63
64        // we have no estimated size or theres more shards
65        let shards = (0..shards)
66            .map(|_| Shard {
67                waiters: Mutex::new(HashMap::new()),
68                table: RwLock::new(RawTable::new()),
69            })
70            .collect();
71
72        LightMap {
73            shards,
74            build_hasher,
75        }
76    }
77}