use std::{collections::BTreeMap, ops::Bound};
use reifydb_codec::key::encoded::{EncodedKey, EncodedKeyRange};
use reifydb_core::interface::store::EntryKind;
use reifydb_store::row::page::{PageId, key_range_of, page_of};
use tracing::instrument;
use crate::{
MultiVersionScope,
tier::{
RangeBatch, RangeCursor, RawEntry,
read::{MultiReadBufferTier, PageEntry, ResidentPage, ServedChunk},
},
};
impl MultiReadBufferTier {
pub fn page_of_key(&self, key: &EncodedKey) -> PageId {
page_of(key, self.bucket_shift())
}
pub fn page_key_range(&self, page: PageId) -> Option<EncodedKeyRange> {
key_range_of(page, self.bucket_shift())
}
pub fn invalidate_page(&self, page: PageId) {
let mut shard = self.shard_for(&page).lock();
shard.pages.remove(&page);
}
pub fn populate_page(&self, page: PageId, entries: Vec<RawEntry>, complete: bool) {
let shift = self.bucket_shift();
let range_complete = complete && key_range_of(page, shift).is_some();
let mut shard = self.shard_for(&page).lock();
let next = shard.next_tick;
let resident = shard.pages.entry(page).or_insert_with(|| ResidentPage {
entries: BTreeMap::new(),
hot: false,
tick: next,
range_complete: false,
warm_blocked: false,
});
for entry in entries {
match resident.entries.get(&entry.key) {
Some(existing) if existing.version > entry.version => continue,
_ => {
resident.entries.insert(
entry.key,
PageEntry {
version: entry.version,
value: entry.value,
previous: None,
},
);
}
}
}
resident.range_complete = range_complete;
resident.tick = next;
shard.next_tick = next + 1;
shard.evict_to_capacity();
}
pub fn finish_warm(&self, page: PageId, entries: Vec<RawEntry>) -> bool {
let shift = self.bucket_shift();
let range_complete = key_range_of(page, shift).is_some();
let mut shard = self.shard_for(&page).lock();
let Some(dirty) = shard.warming.remove(&page) else {
return false;
};
if dirty || !range_complete {
return false;
}
let next = shard.next_tick;
let resident = shard.pages.entry(page).or_insert_with(|| ResidentPage {
entries: BTreeMap::new(),
hot: false,
tick: next,
range_complete: false,
warm_blocked: false,
});
for entry in entries {
match resident.entries.get(&entry.key) {
Some(existing) if existing.version > entry.version => continue,
_ => {
resident.entries.insert(
entry.key,
PageEntry {
version: entry.version,
value: entry.value,
previous: None,
},
);
}
}
}
resident.range_complete = true;
resident.tick = next;
shard.next_tick = next + 1;
shard.evict_to_capacity();
true
}
#[allow(clippy::too_many_arguments)]
#[instrument(name = "store::multi::read::serve", level = "trace", skip(self, cursor, start, end), fields(table = ?table, descending = descending))]
pub fn serve_persistent_chunk(
&self,
table: EntryKind,
cursor: &mut RangeCursor,
start: &[u8],
end: &[u8],
scope: MultiVersionScope,
batch_size: usize,
descending: bool,
) -> ServedChunk {
match table {
EntryKind::Source(_) => {}
EntryKind::Operator(_) | EntryKind::OperatorInternal(_) => {
return self.serve_operator_chunk(cursor, start, end, scope, batch_size, descending);
}
_ => return ServedChunk::Gap,
}
let shift = self.bucket_shift();
let range_lo = EncodedKey::new(start.to_vec());
let range_hi = EncodedKey::new(end.to_vec());
if range_lo > range_hi {
cursor.exhausted = true;
return ServedChunk::Served(RangeBatch::empty());
}
let mut out: Vec<RawEntry> = Vec::new();
let mut first = true;
let mut page = match &cursor.last_key {
Some(last) => page_of(last, shift),
None if descending => page_of(&range_hi, shift),
None => page_of(&range_lo, shift),
};
loop {
let Some(page_range) = key_range_of(page, shift) else {
if out.is_empty() {
return ServedChunk::Gap;
}
return served_chunk(out, cursor, false);
};
let (page_start, page_end) = match (page_range.start, page_range.end) {
(Bound::Included(s), Bound::Included(e)) => (s, e),
_ => {
if out.is_empty() {
return ServedChunk::Gap;
}
return served_chunk(out, cursor, false);
}
};
if descending {
if page_end < range_lo {
return served_chunk(out, cursor, true);
}
} else if page_start > range_hi {
return served_chunk(out, cursor, true);
}
let mut shard = self.shard_for(&page).lock();
let complete = shard.pages.get(&page).map(|p| p.range_complete).unwrap_or(false);
if !complete {
drop(shard);
if out.is_empty() {
return ServedChunk::Gap;
}
return served_chunk(out, cursor, false);
}
let tick = shard.next_tick;
let page_ref = shard.pages.get_mut(&page).expect("complete page present under lock");
let lo_bound: Bound<EncodedKey> = if first {
match &cursor.last_key {
Some(last) if !descending && *last >= range_lo => Bound::Excluded(last.clone()),
_ => Bound::Included(page_start.clone().max(range_lo.clone())),
}
} else {
Bound::Included(page_start.clone().max(range_lo.clone()))
};
let hi_bound: Bound<EncodedKey> = if first {
match &cursor.last_key {
Some(last) if descending && *last <= range_hi => Bound::Excluded(last.clone()),
_ => Bound::Included(page_end.clone().min(range_hi.clone())),
}
} else {
Bound::Included(page_end.clone().min(range_hi.clone()))
};
let mut full = false;
if descending {
for (key, entry) in page_ref.entries.range((lo_bound, hi_bound)).rev() {
if out.len() >= batch_size {
full = true;
break;
}
if scope.contains(entry.version) {
out.push(RawEntry {
key: key.clone(),
version: entry.version,
value: entry.value.clone(),
});
}
}
} else {
for (key, entry) in page_ref.entries.range((lo_bound, hi_bound)) {
if out.len() >= batch_size {
full = true;
break;
}
if scope.contains(entry.version) {
out.push(RawEntry {
key: key.clone(),
version: entry.version,
value: entry.value.clone(),
});
}
}
}
page_ref.hot = true;
page_ref.tick = tick;
shard.next_tick = tick + 1;
drop(shard);
if full {
return served_chunk(out, cursor, false);
}
if descending {
if page_start <= range_lo {
return served_chunk(out, cursor, true);
}
page = PageId {
kind: page.kind,
bucket: page.bucket + 1,
};
} else {
if page_end >= range_hi {
return served_chunk(out, cursor, true);
}
if page.bucket == 0 {
return served_chunk(out, cursor, true);
}
page = PageId {
kind: page.kind,
bucket: page.bucket - 1,
};
}
first = false;
}
}
fn serve_operator_chunk(
&self,
cursor: &mut RangeCursor,
start: &[u8],
end: &[u8],
scope: MultiVersionScope,
batch_size: usize,
descending: bool,
) -> ServedChunk {
let shift = self.bucket_shift();
let range_lo = EncodedKey::new(start.to_vec());
let range_hi = EncodedKey::new(end.to_vec());
if range_lo > range_hi {
cursor.exhausted = true;
return ServedChunk::Served(RangeBatch::empty());
}
let page = page_of(&range_lo, shift);
let Some(page_range) = key_range_of(page, shift) else {
return ServedChunk::Gap;
};
let (Bound::Included(page_start), Bound::Included(page_end)) = (page_range.start, page_range.end)
else {
return ServedChunk::Gap;
};
if range_lo < page_start || range_hi > page_end {
return ServedChunk::Gap;
}
let mut shard = self.shard_for(&page).lock();
let complete = shard.pages.get(&page).map(|p| p.range_complete).unwrap_or(false);
if !complete {
return ServedChunk::Gap;
}
let tick = shard.next_tick;
let page_ref = shard.pages.get_mut(&page).expect("complete page present under lock");
let lo_bound: Bound<EncodedKey> = match &cursor.last_key {
Some(last) if !descending && *last >= range_lo => Bound::Excluded(last.clone()),
_ => Bound::Included(range_lo.clone()),
};
let hi_bound: Bound<EncodedKey> = match &cursor.last_key {
Some(last) if descending && *last <= range_hi => Bound::Excluded(last.clone()),
_ => Bound::Included(range_hi.clone()),
};
let mut out: Vec<RawEntry> = Vec::new();
let mut full = false;
if descending {
for (key, entry) in page_ref.entries.range((lo_bound, hi_bound)).rev() {
if out.len() >= batch_size {
full = true;
break;
}
if entry.version > scope.read() {
return ServedChunk::Gap;
}
if scope.contains(entry.version) {
out.push(RawEntry {
key: key.clone(),
version: entry.version,
value: entry.value.clone(),
});
}
}
} else {
for (key, entry) in page_ref.entries.range((lo_bound, hi_bound)) {
if out.len() >= batch_size {
full = true;
break;
}
if entry.version > scope.read() {
return ServedChunk::Gap;
}
if scope.contains(entry.version) {
out.push(RawEntry {
key: key.clone(),
version: entry.version,
value: entry.value.clone(),
});
}
}
}
page_ref.hot = true;
page_ref.tick = tick;
shard.next_tick = tick + 1;
drop(shard);
served_chunk(out, cursor, !full)
}
pub fn page_is_complete(&self, page: PageId) -> bool {
let shard = self.shard_for(&page).lock();
shard.pages.get(&page).map(|p| p.range_complete).unwrap_or(false)
}
}
fn served_chunk(out: Vec<RawEntry>, cursor: &mut RangeCursor, exhausted: bool) -> ServedChunk {
if let Some(last) = out.last() {
cursor.last_key = Some(last.key.clone());
}
cursor.exhausted = exhausted;
ServedChunk::Served(RangeBatch {
entries: out,
has_more: !exhausted,
})
}