use super::{CacheRef, tip::Buffer};
use crate::{Blob, Error, IoBufMut, IoBufs};
use commonware_utils::Widen;
use std::{num::NonZeroUsize, sync::Arc};
#[derive(Clone, Copy)]
pub(super) enum Tail<'a> {
Buffered(&'a Buffer),
Sealed(&'a [u8]),
}
impl<'a> Tail<'a> {
fn copy_into(self, mut offset: usize, dst: &mut [u8]) {
let tail = match self {
Self::Sealed(tail) => tail,
Self::Buffered(buffer) => {
let (_, tail) = buffer.parts();
let prefix_len = buffer.len() - tail.len();
if offset < prefix_len {
self.cursor().copy_into(offset, dst);
return;
}
offset -= prefix_len;
tail
}
};
dst.copy_from_slice(&tail[offset..offset + dst.len()]);
}
fn cursor(self) -> Cursor<'a, impl Iterator<Item = &'a [u8]>> {
let (prefix, tail) = match self {
Self::Buffered(buffer) => {
let (prefix, tail) = buffer.parts();
(Some(prefix), tail)
}
Self::Sealed(tail) => (None, tail),
};
Cursor {
chunks: prefix
.into_iter()
.flat_map(|bufs| bufs.iter().map(AsRef::as_ref))
.chain(std::iter::once(tail)),
chunk: &[],
position: 0,
}
}
}
struct Cursor<'a, I> {
chunks: I,
chunk: &'a [u8],
position: usize,
}
impl<'a, I: Iterator<Item = &'a [u8]>> Cursor<'a, I> {
fn copy_into(&mut self, offset: usize, mut dst: &mut [u8]) {
if dst.is_empty() {
return;
}
let mut skip = offset
.checked_sub(self.position)
.expect("tail reads must be sorted");
while skip >= self.chunk.len() {
skip -= self.chunk.len();
self.chunk = self.chunks.next().expect("tail read out of bounds");
}
self.chunk = &self.chunk[skip..];
self.position = offset + dst.len();
while !dst.is_empty() {
if self.chunk.is_empty() {
self.chunk = self.chunks.next().expect("tail read out of bounds");
continue;
}
let len = self.chunk.len().min(dst.len());
dst[..len].copy_from_slice(&self.chunk[..len]);
dst = &mut dst[len..];
self.chunk = &self.chunk[len..];
}
}
}
pub struct View<'a, B: Blob> {
pub(super) blob: &'a Arc<B>,
pub(super) cache_ref: &'a CacheRef,
pub(super) id: u64,
pub(super) size: u64,
pub(super) tail_offset: u64,
pub(super) tail: Tail<'a>,
}
impl<B: Blob> Clone for View<'_, B> {
fn clone(&self) -> Self {
*self
}
}
impl<B: Blob> Copy for View<'_, B> {}
impl<B: Blob> View<'_, B> {
fn copy_tail_overlap(&self, buf: &mut [u8], offset: u64) -> usize {
let tail_start = self.tail_offset.max(offset);
let prefix_len = (tail_start - offset) as usize;
let tail_offset = (tail_start - self.tail_offset) as usize;
let (_, tail_buf) = buf.split_at_mut(prefix_len);
self.tail.copy_into(tail_offset, tail_buf);
prefix_len
}
pub fn try_read_sync_into(&self, buf: &mut [u8], offset: u64) -> bool {
let Some(end_offset) = offset.checked_add(buf.len() as u64) else {
return false;
};
if end_offset > self.size {
return false;
}
if buf.is_empty() {
return true;
}
if end_offset <= self.tail_offset {
return self.cache_ref.read_cached(self.id, buf, offset) == buf.len();
}
let dst_start = self.copy_tail_overlap(buf, offset);
if dst_start == 0 {
return true;
}
self.cache_ref
.read_cached(self.id, &mut buf[..dst_start], offset)
== dst_start
}
pub async fn read_into(&self, buf: &mut [u8], offset: u64) -> Result<(), Error> {
let end_offset = offset
.checked_add(buf.len() as u64)
.ok_or(Error::OffsetOverflow)?;
if end_offset > self.size {
return Err(Error::BlobInsufficientLength);
}
let remaining = if end_offset <= self.tail_offset {
buf.len()
} else {
self.copy_tail_overlap(buf, offset)
};
if remaining == 0 {
return Ok(());
}
let cached = self
.cache_ref
.read_cached(self.id, &mut buf[..remaining], offset);
if cached == remaining {
return Ok(());
}
let uncached_offset = offset + cached as u64;
let uncached_len = remaining - cached;
self.cache_ref
.read_after_miss(
self.blob,
self.id,
&mut buf[cached..cached + uncached_len],
uncached_offset,
)
.await
}
pub async fn read_at(&self, offset: u64, len: usize) -> Result<IoBufs, Error> {
let mut buf = IoBufMut::with_capacity(len);
unsafe { buf.set_len(len) };
self.read_into(buf.as_mut(), offset).await?;
Ok(buf.into())
}
pub async fn read_up_to(
&self,
offset: u64,
len: usize,
bufs: impl Into<IoBufMut> + Send,
) -> Result<(IoBufMut, usize), Error> {
let mut bufs = bufs.into();
if len == 0 {
bufs.truncate(0);
return Ok((bufs, 0));
}
let available = self.size.saturating_sub(offset).min(Widen::widen(len)) as usize;
if available == 0 {
return Err(Error::BlobInsufficientLength);
}
unsafe { bufs.set_len(available) };
self.read_into(bufs.as_mut(), offset).await?;
Ok((bufs, available))
}
pub async fn read_many_into(
&self,
buf: &mut [u8],
offsets: &[u64],
item_size: NonZeroUsize,
) -> Result<usize, Error> {
self.read_into_slots(buf, offsets.iter().map(|&o| (o, item_size.get())))
.await
}
pub async fn read_ranges_into(
&self,
buf: &mut [u8],
ranges: &[(u64, usize)],
) -> Result<usize, Error> {
self.read_into_slots(buf, ranges.iter().copied()).await
}
async fn read_into_slots(
&self,
buf: &mut [u8],
ranges: impl ExactSizeIterator<Item = (u64, usize)> + Clone,
) -> Result<usize, Error> {
super::validate_read_ranges(buf.len(), ranges.clone(), self.size)?;
let count = ranges.len();
if count == 0 {
return Ok(0);
}
let mut cache_ranges = split_read_ranges(buf, ranges, self.tail_offset, self.tail);
self.cache_ref.read_cached_many(self.id, &mut cache_ranges);
let blob_reads = cache_ranges.len();
if cache_ranges.is_empty() {
return Ok(count);
}
self.cache_ref
.read_after_misses(self.blob, self.id, cache_ranges)
.await?;
Ok(count - blob_reads)
}
pub fn try_read_many_sync_into(
&self,
buf: &mut [u8],
offsets: &[u64],
item_size: NonZeroUsize,
) -> Vec<usize> {
let ranges = || offsets.iter().map(|&o| (o, item_size.get()));
if super::validate_read_ranges(buf.len(), ranges(), self.size).is_err() {
return (0..offsets.len()).collect();
}
if offsets.is_empty() {
return Vec::new();
}
let mut cache_ranges = split_read_ranges(buf, ranges(), self.tail_offset, self.tail);
if cache_ranges.is_empty() {
return Vec::new();
}
self.cache_ref.read_cached_many(self.id, &mut cache_ranges);
map_misses(cache_ranges, |idx| (offsets[idx], item_size.get()))
}
pub fn try_read_ranges_sync_into(&self, buf: &mut [u8], ranges: &[(u64, usize)]) -> Vec<usize> {
if super::validate_read_ranges(buf.len(), ranges.iter().copied(), self.size).is_err() {
return (0..ranges.len()).collect();
}
if ranges.is_empty() {
return Vec::new();
}
let mut cache_ranges =
split_read_ranges(buf, ranges.iter().copied(), self.tail_offset, self.tail);
if cache_ranges.is_empty() {
return Vec::new();
}
self.cache_ref.read_cached_many(self.id, &mut cache_ranges);
map_misses(cache_ranges, |idx| ranges[idx])
}
}
fn split_read_ranges<'a>(
mut buf: &'a mut [u8],
ranges: impl ExactSizeIterator<Item = (u64, usize)>,
tail_offset: u64,
tail: Tail<'_>,
) -> Vec<(&'a mut [u8], u64)> {
let mut cache_ranges = Vec::with_capacity(ranges.len());
let mut cursor = tail.cursor();
for (offset, len) in ranges {
let (slot, rest) = buf.split_at_mut(len);
buf = rest;
if len == 0 {
continue;
}
let end = offset + len as u64;
if end <= tail_offset {
cache_ranges.push((slot, offset));
} else if offset >= tail_offset {
let src = (offset - tail_offset) as usize;
cursor.copy_into(src, slot);
} else {
let prefix_len = (tail_offset - offset) as usize;
let (prefix, suffix) = slot.split_at_mut(prefix_len);
cursor.copy_into(0, suffix);
cache_ranges.push((prefix, offset));
}
}
cache_ranges
}
fn map_misses(
missed: Vec<(&mut [u8], u64)>,
mut slot: impl FnMut(usize) -> (u64, usize),
) -> Vec<usize> {
let mut misses = Vec::with_capacity(missed.len());
let mut idx = 0;
for (_, offset) in missed {
loop {
let (slot_offset, slot_len) = slot(idx);
if offset >= slot_offset && offset - slot_offset < slot_len as u64 {
break;
}
idx += 1;
}
misses.push(idx);
idx += 1;
}
misses
}
#[cfg(test)]
mod tests {
use super::map_misses;
use crate::{Runner as _, Storage as _, buffer::paged::Writer, deterministic};
use commonware_utils::{NZU16, NZUsize};
use std::num::NonZeroU16;
const PAGE_SIZE: NonZeroU16 = NZU16!(103);
const BUFFER_SIZE: usize = PAGE_SIZE.get() as usize * 2;
#[test]
fn test_reads_cross_buffer_chunks_and_cache() {
deterministic::Runner::default().start(|context| async move {
let cache = super::CacheRef::new(PAGE_SIZE, NZUsize!(4));
let (blob, size) = context
.open("test_partition", b"chunk_reads")
.await
.unwrap();
let page = PAGE_SIZE.get() as usize;
let mut writer = Writer::new(blob, size, page * 32, cache.clone())
.await
.unwrap();
let data: Vec<_> = (0..page * 20 + 13).map(|i| (i % 251) as u8).collect();
let persisted = page * 3 + 17;
(writer, _) = writer.append(&data[..persisted]).await.unwrap();
writer = writer.sync().await.unwrap();
for chunk in data[persisted..].chunks(97) {
(writer, _) = writer.append(chunk).await.unwrap();
}
let mut all = vec![0; data.len()];
assert!(writer.try_read_sync_into(&mut all, 0));
assert_eq!(all, data);
let offsets = [
page - 8,
page * 3 - 8,
page * 4 - 8,
page * 8 - 8,
page * 16 - 8,
]
.map(|offset| offset as u64);
let expected: Vec<_> = offsets
.iter()
.flat_map(|&offset| data[offset as usize..offset as usize + 17].iter().copied())
.collect();
let mut batch = vec![0; expected.len()];
assert!(
writer
.try_read_many_sync_into(&mut batch, &offsets, NZUsize!(17))
.is_empty()
);
assert_eq!(batch, expected);
let ranges = [
(0, 0),
((page - 4) as u64, 16),
((page * 3 - 4) as u64, 15),
((page * 8 - 6) as u64, 19),
((data.len() - 13) as u64, 13),
(data.len() as u64, 0),
];
let expected_ranges: Vec<_> = ranges
.iter()
.flat_map(|&(offset, len)| {
data[offset as usize..offset as usize + len].iter().copied()
})
.collect();
let mut ranged = vec![0; expected_ranges.len()];
assert!(
writer
.try_read_ranges_sync_into(&mut ranged, &ranges)
.is_empty()
);
assert_eq!(ranged, expected_ranges);
cache.clear();
writer
.read_many_into(&mut batch, &offsets, NZUsize!(17))
.await
.unwrap();
assert_eq!(batch, expected);
assert_eq!(
writer
.read_at(0, data.len())
.await
.unwrap()
.coalesce()
.as_ref(),
data
);
writer = writer.sync().await.unwrap();
drop(writer);
cache.clear();
let (blob, size) = context
.open("test_partition", b"chunk_reads")
.await
.unwrap();
let writer = Writer::new(blob, size, page * 32, cache).await.unwrap();
assert_eq!(
writer
.read_at(0, data.len())
.await
.unwrap()
.coalesce()
.as_ref(),
data
);
});
}
#[test]
fn test_read_ranges_into_mixed_lengths_and_tail() {
deterministic::Runner::default().start(|context| async move {
let cache = super::CacheRef::new(PAGE_SIZE, NZUsize!(1));
let (blob, size) = context.open("test_partition", b"ranges").await.unwrap();
let page = PAGE_SIZE.get() as usize;
let mut writer = Writer::new(blob, size, BUFFER_SIZE, cache.clone())
.await
.unwrap();
let data: Vec<_> = (0..3 * page + 37).map(|i| (i % 251) as u8).collect();
(writer, _) = writer.append(&data[..3 * page]).await.unwrap();
writer = writer.sync().await.unwrap();
(writer, _) = writer.append(&data[3 * page..]).await.unwrap();
let ranges = [
(0, 0),
(1, 3),
((page + 2) as u64, 5),
((2 * page - 2) as u64, 4),
((3 * page - 3) as u64, 8),
((3 * page + 9) as u64, 9),
(data.len() as u64, 0),
];
let expected: Vec<_> = ranges
.iter()
.flat_map(|&(offset, len)| data[offset as usize..offset as usize + len].iter())
.copied()
.collect();
let mut out = vec![0xEE; expected.len()];
cache.clear();
assert_eq!(writer.read_ranges_into(&mut out, &ranges).await.unwrap(), 3);
assert_eq!(out, expected);
cache.clear();
writer.read_at(page as u64 + 1, 1).await.unwrap();
out.fill(0xEE);
assert_eq!(writer.read_ranges_into(&mut out, &ranges).await.unwrap(), 4);
assert_eq!(out, expected);
assert_eq!(writer.read_ranges_into(&mut [], &[]).await.unwrap(), 0);
});
}
#[test]
fn test_map_misses_with_cached_prefixes() {
let slots = [
(0, 4),
(10, 0),
(10, 8),
(18, 0),
(18, 5),
(u64::MAX - 4, 4),
];
let mut suffix_a = [0; 2];
let mut suffix_b = [0; 5];
let mut suffix_c = [0; 1];
let missed = vec![
(suffix_a.as_mut_slice(), 16),
(suffix_b.as_mut_slice(), 18),
(suffix_c.as_mut_slice(), u64::MAX - 1),
];
assert_eq!(map_misses(missed, |idx| slots[idx]), vec![2, 4, 5]);
}
#[test]
fn test_view_try_read_sync_straddles_cache_and_tail() {
let executor = deterministic::Runner::default();
executor.start(|context: deterministic::Context| async move {
let cache_ref = super::CacheRef::new(PAGE_SIZE, NZUsize!(BUFFER_SIZE));
let (blob, blob_size) = context
.open("test_partition", b"view_straddle")
.await
.unwrap();
let mut writer = Writer::new(blob, blob_size, BUFFER_SIZE, cache_ref)
.await
.unwrap();
let page_size = PAGE_SIZE.get() as usize;
(writer, _) = writer.append(&vec![0xAA; page_size]).await.unwrap();
(writer, _) = writer.append(b"TAIL").await.unwrap();
writer = writer.sync().await.unwrap();
writer.read_at(0, page_size).await.unwrap();
let mut buf = [0u8; 4];
assert!(writer.try_read_sync_into(&mut buf, page_size as u64 - 2));
assert_eq!(&buf, &[0xAA, 0xAA, b'T', b'A']);
});
}
}