reifydb_store_multi/tier/read/
pool.rs1use std::{
5 collections::{HashMap, hash_map::DefaultHasher},
6 hash::{Hash, Hasher},
7 sync::{
8 Arc,
9 atomic::{AtomicU8, Ordering},
10 },
11};
12
13use reifydb_runtime::sync::mutex::Mutex;
14use reifydb_store::row::page::PageId;
15
16use crate::tier::read::{MultiReadBufferTier, PoolInner, ReadBufferConfig, Shard};
17
18impl MultiReadBufferTier {
19 pub fn new(config: ReadBufferConfig) -> Self {
20 let shard_count = config.shards.max(1);
21 let page_cap = (config.resident_pages / shard_count).max(1);
22 let shards: Vec<Mutex<Shard>> = (0..shard_count)
23 .map(|_| {
24 Mutex::new(Shard {
25 pages: HashMap::new(),
26 warming: HashMap::new(),
27 next_tick: 0,
28 page_cap,
29 })
30 })
31 .collect();
32 Self {
33 inner: Arc::new(PoolInner {
34 shards: shards.into_boxed_slice(),
35 bucket_shift: AtomicU8::new(config.bucket_shift),
36 }),
37 }
38 }
39
40 pub(super) fn bucket_shift(&self) -> u8 {
41 self.inner.bucket_shift.load(Ordering::Relaxed)
42 }
43
44 pub(super) fn shard_for(&self, page: &PageId) -> &Mutex<Shard> {
45 let mut hasher = DefaultHasher::new();
46 page.hash(&mut hasher);
47 let index = (hasher.finish() % self.inner.shards.len() as u64) as usize;
48 &self.inner.shards[index]
49 }
50
51 #[cfg(test)]
52 pub fn len(&self) -> usize {
53 self.inner
54 .shards
55 .iter()
56 .map(|shard| shard.lock().pages.values().map(|page| page.entries.len()).sum::<usize>())
57 .sum()
58 }
59
60 #[cfg(test)]
61 pub fn resident_pages(&self) -> usize {
62 self.inner.shards.iter().map(|shard| shard.lock().pages.len()).sum()
63 }
64}
65
66impl Shard {
67 fn pick_victim(&self) -> Option<PageId> {
68 let mut probationary: Option<(u64, PageId)> = None;
69 let mut hot: Option<(u64, PageId)> = None;
70 for (id, page) in &self.pages {
71 let slot = if page.hot {
72 &mut hot
73 } else {
74 &mut probationary
75 };
76 if slot.map(|(tick, _)| page.tick < tick).unwrap_or(true) {
77 *slot = Some((page.tick, *id));
78 }
79 }
80 probationary.or(hot).map(|(_, id)| id)
81 }
82
83 pub(super) fn evict_to_capacity(&mut self) {
84 while self.pages.len() > self.page_cap {
85 let Some(victim) = self.pick_victim() else {
86 break;
87 };
88 self.pages.remove(&victim);
89 }
90 }
91}