Skip to main content

reifydb_store_multi/tier/read/
pool.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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}