txmap/
custodian.rs

1use crate::{
2    lock_guards::LockGuards,
3    lock_policies::lock_policy::LockPolicy,
4    new_types::{BitMask, ShardCount, ShardIndex},
5    shard::Shard,
6};
7use hashbrown::HashTable;
8use intmap::IntMap;
9
10pub(crate) struct Custodian<K, V, L>
11where
12    L: LockPolicy,
13{
14    pub(crate) shard_count: ShardCount,
15    pub(crate) shards: Vec<L::Lock<Shard<K, V>>>,
16}
17
18impl<K, V, L> Custodian<K, V, L>
19where
20    L: LockPolicy,
21{
22    pub fn new(shard_count: ShardCount, capacity: usize) -> Self {
23        let mut shards = Vec::with_capacity(shard_count.0 as usize);
24        let capacity_per_shard = capacity.div_ceil(shard_count.0 as usize);
25        for _ in 0..shard_count.0 {
26            shards.push(L::new(HashTable::with_capacity(capacity_per_shard)));
27        }
28        Self {
29            shard_count,
30            shards,
31        }
32    }
33    pub fn all_read_guards(&self) -> IntMap<u8, L::ReadGuard<'_, Shard<K, V>>> {
34        self.lock_guards(self.all_bitmask(), BitMask::ZERO).read
35    }
36    pub fn all_write_guards(&self) -> IntMap<u8, L::WriteGuard<'_, Shard<K, V>>> {
37        self.lock_guards(BitMask::ZERO, self.all_bitmask()).write
38    }
39    fn all_bitmask(&self) -> BitMask {
40        let bitmask = if self.shard_count.0 == 128 {
41            !0u128
42        } else {
43            (1 << self.shard_count.0) - 1
44        };
45        BitMask(bitmask)
46    }
47    pub fn lock_guards(&self, read: BitMask, write: BitMask) -> LockGuards<'_, K, V, L> {
48        let mut read_guards = IntMap::new();
49        let mut write_guards = IntMap::new();
50        for i in 0..self.shard_count.0 {
51            let bitmask = ShardIndex(i).bitmask();
52            let shard_lock = &self.shards[i as usize];
53            if (write & bitmask) != BitMask::ZERO {
54                let write_guard = L::write(shard_lock);
55                write_guards.insert(i, write_guard);
56            } else if (read & bitmask) != BitMask::ZERO {
57                let read_guard = L::read(shard_lock);
58                read_guards.insert(i, read_guard);
59            };
60        }
61        LockGuards {
62            read: read_guards,
63            write: write_guards,
64            write_bitmask: write,
65        }
66    }
67    pub fn write_guards(&self, write: BitMask) -> IntMap<u8, L::WriteGuard<'_, Shard<K, V>>> {
68        let mut write_guards = IntMap::new();
69        for i in 0..self.shard_count.0 {
70            let bitmask = ShardIndex(i).bitmask();
71            if (write & bitmask) != BitMask::ZERO {
72                let shard_lock = &self.shards[i as usize];
73                let write_guard = L::write(shard_lock);
74                write_guards.insert(i, write_guard);
75            };
76        }
77        write_guards
78    }
79    pub fn read_guard_at(&self, shard_index: ShardIndex) -> L::ReadGuard<'_, Shard<K, V>> {
80        let shard_lock = &self.shards[shard_index.0 as usize];
81        L::read(shard_lock)
82    }
83    pub fn write_guard_at(&self, shard_index: ShardIndex) -> L::WriteGuard<'_, Shard<K, V>> {
84        let shard_lock = &self.shards[shard_index.0 as usize];
85        L::write(shard_lock)
86    }
87}
txmap/custodian.rs

txmap/
custodian.rs
