use super::allocation::{AllocationError, AllocationLedger, LogicalVecCapacity, VisitedObjectSet};
use super::codec::{
CanonicalDecodeError, DecodeLimits, DecodedPageKind, DecodedPageShape, Decoder, Encoder,
};
use super::frontier;
use super::frontier::Machine as _;
use crate::async_storage::{
AsyncObjectStore, DecodedCacheAdmission, DecodedCacheKey, DecodedCacheValue,
};
use crate::cancellation::CancellationToken;
use crate::performance::{OperationFailure, WorkBudget, WorkCounters, WorkError};
use crate::speculation::{ResidencyHint, ResidencyReason};
use crate::storage::{
ByteRange, ObjectId, ObjectKind, ObjectRead, ObjectReadRequest, ObjectReadRetention,
ObjectReceipt, ObjectStore, ObjectStoreError, object_digest,
};
use bytes::Bytes;
use std::future::Future;
use std::io::Read;
use std::mem::size_of;
use std::sync::Arc;
use thiserror::Error;
const DOMAIN: &[u8; 8] = b"ACYFSBLB";
const VERSION: u16 = 1;
const BLOB_PAGE_HEADER_BYTES: usize = 8 + 2 + 8 + 8 + 1 + 4;
const BLOB_PAGE_ITEM_BYTES: usize = 8 + 8 + 32;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct BlobChunkRef {
pub first_offset: u64,
pub end_offset: u64,
pub chunk: ObjectId,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct BlobChild {
pub first_offset: u64,
pub end_offset: u64,
pub page: ObjectId,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum BlobNode {
Leaf(Vec<BlobChunkRef>),
Internal(Vec<BlobChild>),
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BlobPage {
pub first_offset: u64,
pub end_offset: u64,
pub node: BlobNode,
}
impl BlobPage {
fn validate(&self, maximum_items: u32) -> Result<(), CanonicalDecodeError> {
let (count, ranges_valid, kinds_valid) = match &self.node {
BlobNode::Leaf(items) => (
items.len(),
contiguous(
self.first_offset,
self.end_offset,
items
.iter()
.map(|item| (item.first_offset, item.end_offset)),
),
items
.iter()
.all(|item| item.chunk.kind == ObjectKind::BlobChunk),
),
BlobNode::Internal(items) => (
items.len(),
contiguous(
self.first_offset,
self.end_offset,
items
.iter()
.map(|item| (item.first_offset, item.end_offset)),
),
items.iter().all(|item| item.page.kind == ObjectKind::Blob),
),
};
if u32::try_from(count).unwrap_or(u32::MAX) > maximum_items {
return Err(CanonicalDecodeError::FieldTooLarge {
observed: u32::try_from(count).unwrap_or(u32::MAX),
maximum: maximum_items,
});
}
if count == 0 && (self.first_offset != 0 || self.end_offset != 0) {
return Err(invariant("non-empty blob bounds require entries"));
}
if count != 0 && (!ranges_valid || self.first_offset >= self.end_offset) {
return Err(invariant("blob ranges are not positive and contiguous"));
}
if !kinds_valid {
return Err(invariant("blob reference has the wrong object kind"));
}
Ok(())
}
}
fn contiguous(first: u64, end: u64, mut ranges: impl Iterator<Item = (u64, u64)>) -> bool {
let Some((item_first, item_end)) = ranges.next() else {
return first == 0 && end == 0;
};
if item_first != first || item_first >= item_end {
return false;
}
let mut prior_end = item_end;
for (next_first, next_end) in ranges {
if next_first != prior_end || next_first >= next_end {
return false;
}
prior_end = next_end;
}
prior_end == end
}
pub fn encode_blob_page(
page: &BlobPage,
maximum_items: u32,
) -> Result<Vec<u8>, CanonicalDecodeError> {
let encoded_length = blob_page_encoded_length(page, maximum_items)?;
let mut encoder = Encoder::with_exact_capacity(DOMAIN, VERSION, encoded_length)?;
encoder.u64(page.first_offset);
encoder.u64(page.end_offset);
match &page.node {
BlobNode::Leaf(items) => {
encoder.u8(1);
encoder.u32(count(items.len())?);
for item in items {
encoder.u64(item.first_offset);
encoder.u64(item.end_offset);
encoder.fixed(item.chunk.digest.as_bytes());
}
}
BlobNode::Internal(items) => {
encoder.u8(2);
encoder.u32(count(items.len())?);
for item in items {
encoder.u64(item.first_offset);
encoder.u64(item.end_offset);
encoder.fixed(item.page.digest.as_bytes());
}
}
}
Ok(encoder.finish())
}
fn blob_page_encoded_length(
page: &BlobPage,
maximum_items: u32,
) -> Result<usize, CanonicalDecodeError> {
page.validate(maximum_items)?;
let items = match &page.node {
BlobNode::Leaf(items) => items.len(),
BlobNode::Internal(items) => items.len(),
};
items
.checked_mul(BLOB_PAGE_ITEM_BYTES)
.and_then(|bytes| BLOB_PAGE_HEADER_BYTES.checked_add(bytes))
.ok_or(CanonicalDecodeError::LengthOverflow)
}
pub fn decode_blob_page(
bytes: &[u8],
limits: DecodeLimits,
) -> Result<BlobPage, CanonicalDecodeError> {
let mut decoder = Decoder::new(bytes, DOMAIN, VERSION, limits.maximum_page_object_bytes())?;
let first_offset = decoder.u64()?;
let end_offset = decoder.u64()?;
let tag = decoder.u8()?;
let item_count = decoder.u32()?;
if item_count > limits.maximum_page_items {
return Err(CanonicalDecodeError::FieldTooLarge {
observed: item_count,
maximum: limits.maximum_page_items,
});
}
let mut leaf = Vec::new();
let mut internal = Vec::new();
match tag {
1 => {
leaf.try_reserve_exact(capacity(item_count)?)
.map_err(|_| CanonicalDecodeError::AllocationFailed)?;
for _ in 0..item_count {
leaf.push(BlobChunkRef {
first_offset: decoder.u64()?,
end_offset: decoder.u64()?,
chunk: ObjectId {
kind: ObjectKind::BlobChunk,
digest: crate::foundation::Digest::from_bytes(decoder.fixed()?),
},
});
}
}
2 => {
internal
.try_reserve_exact(capacity(item_count)?)
.map_err(|_| CanonicalDecodeError::AllocationFailed)?;
for _ in 0..item_count {
internal.push(BlobChild {
first_offset: decoder.u64()?,
end_offset: decoder.u64()?,
page: ObjectId {
kind: ObjectKind::Blob,
digest: crate::foundation::Digest::from_bytes(decoder.fixed()?),
},
});
}
}
tag => {
return Err(CanonicalDecodeError::UnknownTag {
field: "blob_page",
tag,
});
}
}
decoder.finish()?;
let page = BlobPage {
first_offset,
end_offset,
node: if tag == 1 {
BlobNode::Leaf(leaf)
} else {
BlobNode::Internal(internal)
},
};
page.validate(limits.maximum_page_items)?;
Ok(page)
}
fn blob_page_decode_shape(
bytes: &[u8],
limits: DecodeLimits,
) -> Result<DecodedPageShape, CanonicalDecodeError> {
let mut decoder = Decoder::new(bytes, DOMAIN, VERSION, limits.maximum_page_object_bytes())?;
decoder.u64()?;
decoder.u64()?;
let tag = decoder.u8()?;
let item_count = decoder.u32()?;
if item_count > limits.maximum_page_items {
return Err(CanonicalDecodeError::FieldTooLarge {
observed: item_count,
maximum: limits.maximum_page_items,
});
}
let kind = match tag {
1 | 2 => {
for _ in 0..item_count {
decoder.u64()?;
decoder.u64()?;
let _: [u8; 32] = decoder.fixed()?;
}
if tag == 1 {
DecodedPageKind::Leaf
} else {
DecodedPageKind::Internal
}
}
tag => {
return Err(CanonicalDecodeError::UnknownTag {
field: "blob_page",
tag,
});
}
};
decoder.finish()?;
Ok(DecodedPageShape {
kind,
items: usize::try_from(item_count).map_err(|_| CanonicalDecodeError::LengthOverflow)?,
nested_bytes: 0,
})
}
pub fn blob_page_id(page: &BlobPage, maximum_items: u32) -> Result<ObjectId, CanonicalDecodeError> {
let bytes = encode_blob_page(page, maximum_items)?;
Ok(ObjectId {
kind: ObjectKind::Blob,
digest: object_digest(ObjectKind::Blob, &bytes),
})
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct BlobBuildOptions {
pub chunk_bytes: u32,
pub page_items: u32,
pub page_bytes: u32,
pub maximum_blob_bytes: u64,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BlobBuild {
pub root: ObjectId,
pub logical_bytes: u64,
pub work: WorkCounters,
}
pub trait AsyncBlobSource {
fn read<'a>(
&'a mut self,
destination: &'a mut [u8],
cancellation: &'a CancellationToken,
) -> impl Future<Output = std::io::Result<usize>>;
fn read_owned(
&mut self,
_maximum: usize,
_cancellation: &CancellationToken,
) -> impl Future<Output = std::io::Result<Option<Bytes>>> {
async { Ok(None) }
}
fn take_work(&mut self) -> Option<WorkCounters> {
None
}
}
impl<T: Read> AsyncBlobSource for T {
async fn read<'a>(
&'a mut self,
destination: &'a mut [u8],
cancellation: &'a CancellationToken,
) -> std::io::Result<usize> {
if cancellation.is_cancelled() {
return Err(std::io::Error::new(
std::io::ErrorKind::Interrupted,
"blob source read cancelled",
));
}
Read::read(self, destination)
}
}
pub fn build_blob<S: crate::ImmediateObjectStore, R: Read>(
store: &S,
source: &mut R,
options: BlobBuildOptions,
budget: WorkBudget,
) -> Result<BlobBuild, BlobBuildFailure> {
let cancellation = CancellationToken::new();
crate::async_storage::poll_immediate(build_blob_async(
store,
source,
options,
budget,
&cancellation,
))
}
pub async fn build_blob_async<S: AsyncObjectStore, R: AsyncBlobSource>(
store: &S,
source: &mut R,
options: BlobBuildOptions,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> Result<BlobBuild, BlobBuildFailure> {
let chunk_capacity = validate_blob_build(options, cancellation)?;
let (mut batching, mut work) = BlobBatchStore::new(store, budget)?;
let mut index = BlobIndexBuilder::new(options);
let mut logical_bytes = 0_u64;
loop {
let detection_bytes = options
.maximum_blob_bytes
.saturating_sub(logical_bytes)
.saturating_add(1);
let allocation = chunk_capacity.min(usize::try_from(detection_bytes).unwrap_or(usize::MAX));
let retained = batching.retained_bytes();
let owned = source.read_owned(allocation, cancellation).await;
let source_work = source.take_work();
if let Some(measured) = source_work {
work = charge_source_work(
work,
measured,
retained.saturating_add(index.live_allocation_bytes),
budget,
)?;
}
let owned = owned.map_err(|error| build_failed(BlobBuildError::Source(error), work))?;
if let Some(chunk) = owned {
if chunk.is_empty() {
break;
}
let filled = chunk.len();
(logical_bytes, work) = accept_owned_blob_chunk(
&mut batching,
&mut index,
chunk,
allocation,
retained,
source_work.is_none(),
logical_bytes,
options.maximum_blob_bytes,
budget,
work,
cancellation,
)
.await?;
if filled < allocation {
break;
}
continue;
}
let (bytes, filled, retained_capacity, simultaneous, next_work) = read_blob_chunk(
source,
allocation,
index.live_allocation_bytes,
retained,
work,
budget,
cancellation,
)
.await?;
work = next_work;
if filled == 0 {
break;
}
let bytes = if filled == bytes.len() {
Bytes::from(bytes)
} else {
Bytes::copy_from_slice(
bytes
.get(..filled)
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, work))?,
)
};
(logical_bytes, work) = accept_owned_blob_chunk(
&mut batching,
&mut index,
bytes,
allocation,
retained,
false,
logical_bytes,
options.maximum_blob_bytes,
budget,
work,
cancellation,
)
.await?;
let _ = (retained_capacity, simultaneous);
if filled < allocation {
break;
}
}
let root = index
.finish(&mut batching, budget, &mut work, cancellation)
.await?;
work = finish_blob_batch(&mut batching, work, budget, cancellation).await?;
Ok(BlobBuild {
root,
logical_bytes,
work,
})
}
#[allow(clippy::too_many_arguments)]
async fn read_blob_chunk<R: AsyncBlobSource>(
source: &mut R,
allocation: usize,
index_bytes: u64,
retained: u64,
mut work: WorkCounters,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> Result<(Vec<u8>, usize, u64, u64, WorkCounters), BlobBuildFailure> {
let (mut bytes, prospective, simultaneous, retained_capacity) =
allocate_chunk_buffer(allocation, index_bytes, retained, work, budget)?;
work = prospective;
let mut filled = 0_usize;
while filled < bytes.len() {
if cancellation.is_cancelled() {
return Err(build_failed(BlobBuildError::Cancelled, work));
}
#[allow(clippy::indexing_slicing, reason = "bounded by while loop above")]
let read = AsyncBlobSource::read(source, &mut bytes[filled..], cancellation).await;
let measured = source.take_work();
if let Some(measured) = measured {
work = charge_source_work(work, measured, simultaneous, budget)?;
}
match read {
Ok(0) => break,
Ok(count) => {
if count > bytes.len() - filled {
return Err(build_failed(BlobBuildError::TooLarge, work));
}
filled = filled.saturating_add(count);
if measured.is_none() {
work = build_add(
work,
WorkCounters {
source_bytes_read: u64::try_from(count).unwrap_or(u64::MAX),
..WorkCounters::default()
},
)?;
}
build_verify(work, budget)?;
}
Err(error) => return Err(build_failed(BlobBuildError::Source(error), work)),
}
}
Ok((bytes, filled, retained_capacity, simultaneous, work))
}
fn charge_source_work(
work: WorkCounters,
mut measured: WorkCounters,
live_bytes: u64,
budget: WorkBudget,
) -> Result<WorkCounters, BlobBuildFailure> {
let simultaneous = live_bytes
.checked_add(measured.peak_allocation_bytes)
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, work))?;
measured.peak_allocation_bytes = 0;
let mut combined = build_add(work, measured)?;
combined.peak_allocation_bytes = combined.peak_allocation_bytes.max(simultaneous);
build_verify(combined, budget)?;
Ok(combined)
}
#[allow(clippy::too_many_arguments)]
async fn accept_owned_blob_chunk<S: AsyncObjectStore>(
batching: &mut BlobBatchStore<'_, S>,
index: &mut BlobIndexBuilder,
chunk: Bytes,
allocation: usize,
retained: u64,
owned_source: bool,
logical_bytes: u64,
maximum_blob_bytes: u64,
budget: WorkBudget,
work: WorkCounters,
cancellation: &CancellationToken,
) -> Result<(u64, WorkCounters), BlobBuildFailure> {
if chunk.len() > allocation {
return Err(build_failed(BlobBuildError::TooLarge, work));
}
let retained_capacity = chunk.len() as u64;
let prospective = build_add(
work,
WorkCounters {
source_bytes_read: if owned_source { retained_capacity } else { 0 },
peak_allocation_bytes: index
.live_allocation_bytes
.saturating_add(retained)
.saturating_add(retained_capacity),
..WorkCounters::default()
},
)?;
build_verify(prospective, budget)?;
let logical_bytes = logical_bytes
.checked_add(retained_capacity)
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, prospective))?;
if logical_bytes > maximum_blob_bytes {
return Err(build_failed(BlobBuildError::TooLarge, prospective));
}
let first_offset = logical_bytes - retained_capacity;
let chunk_id = ObjectId {
kind: ObjectKind::BlobChunk,
digest: object_digest(ObjectKind::BlobChunk, &chunk),
};
let mut work = build_put(
batching,
chunk_id,
chunk,
PutAllocation {
retained: retained_capacity,
live: index
.live_allocation_bytes
.saturating_add(retained_capacity),
},
budget,
prospective,
cancellation,
)
.await?;
index
.push_chunk(
batching,
BlobChunkRef {
first_offset,
end_offset: logical_bytes,
chunk: chunk_id,
},
budget,
&mut work,
cancellation,
)
.await?;
Ok((logical_bytes, work))
}
fn validate_blob_build(
options: BlobBuildOptions,
cancellation: &CancellationToken,
) -> Result<usize, BlobBuildFailure> {
if cancellation.is_cancelled() {
return Err(OperationFailure::before_work(BlobBuildError::Cancelled));
}
if options.chunk_bytes == 0
|| effective_page_width(options).is_none()
|| options.maximum_blob_bytes == 0
{
return Err(build_failed(
BlobBuildError::InvalidOptions,
WorkCounters::default(),
));
}
usize::try_from(options.chunk_bytes)
.map_err(|_| build_failed(BlobBuildError::InvalidOptions, WorkCounters::default()))
}
async fn finish_blob_batch<S: AsyncObjectStore>(
batching: &mut BlobBatchStore<'_, S>,
work: WorkCounters,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> Result<WorkCounters, BlobBuildFailure> {
let retained = batching.retained_bytes();
let mut remaining = work
.remaining(budget)
.map_err(|error| build_failed(BlobBuildError::Work(error), work))?;
remaining.peak_allocation_bytes = remaining
.peak_allocation_bytes
.checked_sub(retained)
.ok_or_else(|| build_failed(BlobBuildError::Work(WorkError::Overflow), work))?;
let receipt = batching
.finish(remaining, cancellation)
.await
.map_err(|failure| failure.map_with_prior_work(work, BlobBuildError::Storage))?;
let backend_peak = receipt.work.peak_allocation_bytes;
let mut backend_work = receipt.work;
backend_work.peak_allocation_bytes = 0;
let mut combined = build_add(work, backend_work)?;
let simultaneous = retained
.checked_add(backend_peak)
.ok_or_else(|| build_failed(BlobBuildError::Work(WorkError::Overflow), work))?;
combined.peak_allocation_bytes = combined.peak_allocation_bytes.max(simultaneous);
build_verify(combined, budget)?;
Ok(combined)
}
const BLOB_BATCH_OBJECTS: usize = 4;
const BLOB_BATCH_BYTES: u64 = 4 * 1024 * 1024;
struct BlobBatchState {
writes: Vec<crate::storage::ObjectWrite>,
bytes: u64,
}
struct BlobBatchStore<'a, S> {
inner: &'a S,
state: BlobBatchState,
vector_bytes: u64,
}
impl<'a, S: AsyncObjectStore> BlobBatchStore<'a, S> {
fn new(inner: &'a S, budget: WorkBudget) -> Result<(Self, WorkCounters), BlobBuildFailure> {
let minimum = WorkCounters {
allocation_operations: 1,
peak_allocation_bytes: u64::try_from(BLOB_BATCH_OBJECTS)
.unwrap_or(u64::MAX)
.saturating_mul(
u64::try_from(size_of::<crate::storage::ObjectWrite>()).unwrap_or(u64::MAX),
),
..WorkCounters::default()
};
minimum
.verify(budget)
.map_err(|error| build_failed(BlobBuildError::Work(error), WorkCounters::default()))?;
let mut writes = Vec::new();
writes
.try_reserve_exact(BLOB_BATCH_OBJECTS)
.map_err(|_| build_failed(BlobBuildError::AllocationFailed, WorkCounters::default()))?;
let vector_bytes = u64::try_from(writes.capacity())
.unwrap_or(u64::MAX)
.saturating_mul(
u64::try_from(size_of::<crate::storage::ObjectWrite>()).unwrap_or(u64::MAX),
);
let admitted = WorkCounters {
allocation_operations: 1,
peak_allocation_bytes: vector_bytes,
..WorkCounters::default()
};
admitted
.verify(budget)
.map_err(|error| build_failed(BlobBuildError::Work(error), admitted))?;
Ok((
Self {
inner,
state: BlobBatchState { writes, bytes: 0 },
vector_bytes,
},
admitted,
))
}
fn retained_bytes(&self) -> u64 {
self.vector_bytes.saturating_add(self.state.bytes)
}
async fn flush_locked(
&mut self,
backend_budget: WorkBudget,
cancellation: &CancellationToken,
) -> crate::storage::ObjectResult<()> {
if self.state.writes.is_empty() {
return Ok(crate::storage::ObjectReceipt {
value: (),
work: WorkCounters::default(),
});
}
let receipt = self
.inner
.put_many(&self.state.writes, backend_budget, cancellation)
.await?;
self.state.writes.clear();
self.state.bytes = 0;
Ok(crate::storage::ObjectReceipt {
value: (),
work: receipt.work,
})
}
async fn finish(
&mut self,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> crate::storage::ObjectResult<()> {
self.flush_locked(budget, cancellation).await
}
async fn put_with_retained(
&mut self,
object_id: ObjectId,
bytes: Bytes,
retained_bytes: u64,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> crate::storage::ObjectResult<()> {
cancellation.check().map_err(|_| {
crate::storage::ObjectFailure::before_work(crate::storage::ObjectStoreError::Cancelled)
})?;
self.state.bytes = self
.state
.bytes
.checked_add(retained_bytes)
.ok_or_else(|| {
crate::storage::ObjectFailure::before_work(
crate::storage::ObjectStoreError::Rejected(
"object write batch length overflowed".to_owned(),
),
)
})?;
self.state
.writes
.push(crate::storage::ObjectWrite { object_id, bytes });
if self.state.writes.len() >= BLOB_BATCH_OBJECTS || self.state.bytes >= BLOB_BATCH_BYTES {
return self.flush_locked(budget, cancellation).await;
}
Ok(crate::storage::ObjectReceipt {
value: (),
work: WorkCounters::default(),
})
}
}
fn allocate_chunk_buffer(
allocation: usize,
index_live_bytes: u64,
retained_batch_bytes: u64,
work: WorkCounters,
budget: WorkBudget,
) -> Result<(Vec<u8>, WorkCounters, u64, u64), BlobBuildFailure> {
let allocation_bytes = u64::try_from(allocation).unwrap_or(u64::MAX);
let simultaneous = index_live_bytes
.checked_add(allocation_bytes)
.and_then(|bytes| bytes.checked_add(retained_batch_bytes))
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, work))?;
let mut prospective = build_add(
work,
WorkCounters {
allocation_operations: 1,
..WorkCounters::default()
},
)?;
prospective.peak_allocation_bytes = prospective.peak_allocation_bytes.max(simultaneous);
build_verify(prospective, budget)?;
let mut bytes = Vec::new();
if bytes.try_reserve_exact(allocation).is_err() {
return Err(build_failed(BlobBuildError::AllocationFailed, work));
}
bytes.resize(allocation, 0);
let retained_capacity = u64::try_from(bytes.capacity()).unwrap_or(u64::MAX);
let actual_simultaneous = index_live_bytes
.checked_add(retained_capacity)
.and_then(|value| value.checked_add(retained_batch_bytes))
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, prospective))?;
prospective.peak_allocation_bytes = prospective.peak_allocation_bytes.max(actual_simultaneous);
build_verify(prospective, budget)?;
Ok((bytes, prospective, actual_simultaneous, retained_capacity))
}
struct BlobIndexBuilder {
options: BlobBuildOptions,
width: usize,
leaf: Vec<BlobChunkRef>,
levels: Vec<Vec<BlobChild>>,
live_allocation_bytes: u64,
}
impl BlobIndexBuilder {
fn new(options: BlobBuildOptions) -> Self {
Self {
options,
width: effective_page_width(options).unwrap_or(0),
leaf: Vec::new(),
levels: Vec::new(),
live_allocation_bytes: 0,
}
}
async fn push_chunk<S: AsyncObjectStore>(
&mut self,
store: &mut BlobBatchStore<'_, S>,
chunk: BlobChunkRef,
budget: WorkBudget,
work: &mut WorkCounters,
cancellation: &CancellationToken,
) -> Result<(), BlobBuildFailure> {
reserve_page_items(
&mut self.leaf,
self.width,
&mut self.live_allocation_bytes,
work,
budget,
)?;
self.leaf.push(chunk);
if self.leaf.len() == self.width {
self.flush_leaf(store, budget, work, cancellation).await?;
}
Ok(())
}
async fn flush_leaf<S: AsyncObjectStore>(
&mut self,
store: &mut BlobBatchStore<'_, S>,
budget: WorkBudget,
work: &mut WorkCounters,
cancellation: &CancellationToken,
) -> Result<(), BlobBuildFailure> {
let mut items = std::mem::take(&mut self.leaf);
if items.is_empty() {
return Ok(());
}
let (first_offset, end_offset) = match items.as_slice() {
[only] => (only.first_offset, only.end_offset),
[first, .., last] => (first.first_offset, last.end_offset),
[] => unreachable!("checked items.is_empty() above"),
};
let page = BlobPage {
first_offset,
end_offset,
node: BlobNode::Leaf(items),
};
let page_id = put_blob_page(
store,
&page,
self.options,
self.live_allocation_bytes,
budget,
work,
cancellation,
)
.await?;
let child = BlobChild {
first_offset: page.first_offset,
end_offset: page.end_offset,
page: page_id,
};
let BlobNode::Leaf(reusable) = page.node else {
unreachable!();
};
items = reusable;
items.clear();
self.leaf = items;
self.push_child(store, 0, child, budget, work, cancellation)
.await
}
#[allow(
clippy::indexing_slicing,
reason = "`self.ensure_level(level, ...)` runs immediately above every use of \
`self.levels[level]` in this loop iteration; it loops `while self.levels.len() \
<= level` pushing new levels, so afterwards `level < self.levels.len()` and \
`self.levels` is not resized again before this iteration's accesses"
)]
async fn push_child<S: AsyncObjectStore>(
&mut self,
store: &mut BlobBatchStore<'_, S>,
mut level: usize,
mut child: BlobChild,
budget: WorkBudget,
work: &mut WorkCounters,
cancellation: &CancellationToken,
) -> Result<(), BlobBuildFailure> {
loop {
self.ensure_level(level, work, budget)?;
reserve_page_items(
&mut self.levels[level],
self.width,
&mut self.live_allocation_bytes,
work,
budget,
)?;
self.levels[level].push(child);
if self.levels[level].len() < self.width {
return Ok(());
}
child = self
.flush_internal(store, level, budget, work, cancellation)
.await?;
level = level
.checked_add(1)
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, *work))?;
}
}
#[allow(
clippy::indexing_slicing,
reason = "`flush_internal` is only called from `push_child` (right after \
`self.ensure_level(level, ...)` established `level < self.levels.len()`, with \
no intervening resize) and from `finish` (with `level` taken from \
`self.levels.iter().position(...)`, which is inherently `< self.levels.len()`)"
)]
async fn flush_internal<S: AsyncObjectStore>(
&mut self,
store: &mut BlobBatchStore<'_, S>,
level: usize,
budget: WorkBudget,
work: &mut WorkCounters,
cancellation: &CancellationToken,
) -> Result<BlobChild, BlobBuildFailure> {
let items = std::mem::take(&mut self.levels[level]);
let (first_offset, end_offset) = match items.as_slice() {
[only] => (only.first_offset, only.end_offset),
[first, .., last] => (first.first_offset, last.end_offset),
[] => unreachable!(
"flush_internal is only called once self.levels[level] is non-empty: from \
push_child at self.width (>= 2) items, or from finish() for a level chosen via \
self.levels.iter().position(|items| !items.is_empty())"
),
};
let page = BlobPage {
first_offset,
end_offset,
node: BlobNode::Internal(items),
};
let page_id = put_blob_page(
store,
&page,
self.options,
self.live_allocation_bytes,
budget,
work,
cancellation,
)
.await?;
let child = BlobChild {
first_offset: page.first_offset,
end_offset: page.end_offset,
page: page_id,
};
let BlobNode::Internal(mut reusable) = page.node else {
unreachable!();
};
reusable.clear();
self.levels[level] = reusable;
Ok(child)
}
fn ensure_level(
&mut self,
level: usize,
work: &mut WorkCounters,
budget: WorkBudget,
) -> Result<(), BlobBuildFailure> {
while self.levels.len() <= level {
let bytes = u64::try_from(size_of::<Vec<BlobChild>>()).unwrap_or(u64::MAX);
let next_live = self
.live_allocation_bytes
.checked_add(bytes)
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, *work))?;
let mut prospective = work
.checked_add(WorkCounters {
allocation_operations: 1,
..WorkCounters::default()
})
.map_err(|error| build_failed(error.into(), *work))?;
prospective.peak_allocation_bytes = prospective.peak_allocation_bytes.max(next_live);
build_verify(prospective, budget)?;
if self.levels.try_reserve_exact(1).is_err() {
return Err(build_failed(BlobBuildError::AllocationFailed, *work));
}
self.levels.push(Vec::new());
self.live_allocation_bytes = next_live;
*work = prospective;
}
Ok(())
}
async fn finish<S: AsyncObjectStore>(
mut self,
store: &mut BlobBatchStore<'_, S>,
budget: WorkBudget,
work: &mut WorkCounters,
cancellation: &CancellationToken,
) -> Result<ObjectId, BlobBuildFailure> {
if self.leaf.is_empty() && self.levels.is_empty() {
return put_blob_page(
store,
&BlobPage {
first_offset: 0,
end_offset: 0,
node: BlobNode::Leaf(Vec::new()),
},
self.options,
self.live_allocation_bytes,
budget,
work,
cancellation,
)
.await;
}
self.flush_leaf(store, budget, work, cancellation).await?;
loop {
let Some(level) = self.levels.iter().position(|items| !items.is_empty()) else {
return Err(build_failed(BlobBuildError::InvalidIndexState, *work));
};
#[allow(
clippy::indexing_slicing,
reason = "`level` comes from `self.levels.iter().position(...)` above, so `level \
< self.levels.len()` and `level + 1 <= self.levels.len()`, which is \
in-bounds for a `RangeFrom` slice"
)]
let higher_is_empty = self.levels[level + 1..].iter().all(Vec::is_empty);
#[allow(
clippy::indexing_slicing,
reason = "`level < self.levels.len()` from `position(...)` above; the \
`levels[level][0]` access is additionally guarded by the \
`levels[level].len() == 1` check in this same condition"
)]
if higher_is_empty && self.levels[level].len() == 1 {
return Ok(self.levels[level][0].page);
}
let child = self
.flush_internal(store, level, budget, work, cancellation)
.await?;
self.push_child(store, level + 1, child, budget, work, cancellation)
.await?;
}
}
}
fn effective_page_width(options: BlobBuildOptions) -> Option<usize> {
let item_bound = usize::try_from(options.page_items).ok()?;
let byte_bound = usize::try_from(options.page_bytes)
.ok()?
.checked_sub(BLOB_PAGE_HEADER_BYTES)?
/ BLOB_PAGE_ITEM_BYTES;
let width = item_bound.min(byte_bound);
(width >= 2).then_some(width)
}
fn reserve_page_items<T>(
items: &mut Vec<T>,
width: usize,
live_allocation_bytes: &mut u64,
work: &mut WorkCounters,
budget: WorkBudget,
) -> Result<(), BlobBuildFailure> {
if items.capacity() >= width {
return Ok(());
}
let bytes = width
.checked_mul(size_of::<T>())
.map(crate::foundation::usize_to_u64)
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, *work))?;
let next_live = live_allocation_bytes
.checked_add(bytes)
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, *work))?;
let mut prospective = work
.checked_add(WorkCounters {
allocation_operations: 1,
..WorkCounters::default()
})
.map_err(|error| build_failed(error.into(), *work))?;
prospective.peak_allocation_bytes = prospective.peak_allocation_bytes.max(next_live);
build_verify(prospective, budget)?;
if items.try_reserve_exact(width - items.capacity()).is_err() {
return Err(build_failed(BlobBuildError::AllocationFailed, *work));
}
*live_allocation_bytes = next_live;
*work = prospective;
Ok(())
}
async fn put_blob_page<S: AsyncObjectStore>(
store: &mut BlobBatchStore<'_, S>,
page: &BlobPage,
options: BlobBuildOptions,
live_allocation_bytes: u64,
budget: WorkBudget,
work: &mut WorkCounters,
cancellation: &CancellationToken,
) -> Result<ObjectId, BlobBuildFailure> {
let retained = store.retained_bytes();
let encoded = encode_blob_page(page, options.page_items)
.map_err(|error| build_failed(BlobBuildError::Encode(error), *work))?;
if encoded.len() > usize::try_from(options.page_bytes).unwrap_or(usize::MAX) {
return Err(build_failed(
BlobBuildError::PageTooLarge {
observed: u64::try_from(encoded.len()).unwrap_or(u64::MAX),
maximum: u64::from(options.page_bytes),
},
*work,
));
}
let page_id = ObjectId {
kind: ObjectKind::Blob,
digest: object_digest(ObjectKind::Blob, &encoded),
};
let encoded_bytes = u64::try_from(encoded.capacity()).unwrap_or(u64::MAX);
let simultaneous = live_allocation_bytes
.checked_add(encoded_bytes)
.and_then(|bytes| bytes.checked_add(retained))
.ok_or_else(|| build_failed(BlobBuildError::TooLarge, *work))?;
let mut encoded_work = build_add(
*work,
WorkCounters {
page_writes: 1,
bytes_encoded: u64::try_from(encoded.len()).unwrap_or(u64::MAX),
allocation_operations: 1,
..WorkCounters::default()
},
)?;
encoded_work.peak_allocation_bytes = encoded_work.peak_allocation_bytes.max(simultaneous);
build_verify(encoded_work, budget)?;
*work = build_put(
store,
page_id,
Bytes::from(encoded),
PutAllocation {
retained: encoded_bytes,
live: simultaneous,
},
budget,
encoded_work,
cancellation,
)
.await?;
Ok(page_id)
}
#[derive(Clone, Copy)]
struct PutAllocation {
retained: u64,
live: u64,
}
async fn build_put<S: AsyncObjectStore>(
store: &mut BlobBatchStore<'_, S>,
object: ObjectId,
bytes: Bytes,
allocation: PutAllocation,
budget: WorkBudget,
work: WorkCounters,
cancellation: &CancellationToken,
) -> Result<WorkCounters, BlobBuildFailure> {
let mut remaining = work
.remaining(budget)
.map_err(|error| build_failed(BlobBuildError::Work(error), work))?;
remaining.peak_allocation_bytes = budget
.peak_allocation_bytes
.checked_sub(allocation.live)
.ok_or_else(|| build_failed(BlobBuildError::Work(WorkError::Overflow), work))?;
match store
.put_with_retained(object, bytes, allocation.retained, remaining, cancellation)
.await
{
Ok(receipt) => {
let backend_peak = allocation
.live
.checked_add(receipt.work.peak_allocation_bytes)
.ok_or_else(|| build_failed(BlobBuildError::Work(WorkError::Overflow), work))?;
let mut backend_work = receipt.work;
backend_work.peak_allocation_bytes = 0;
let mut combined = build_add(work, backend_work)?;
combined.peak_allocation_bytes = combined.peak_allocation_bytes.max(backend_peak);
build_verify(combined, budget)?;
Ok(combined)
}
Err(failure) => Err(failure.map_with_prior_work(work, BlobBuildError::Storage)),
}
}
fn build_add(left: WorkCounters, right: WorkCounters) -> Result<WorkCounters, BlobBuildFailure> {
left.checked_add(right)
.map_err(|error| build_failed(BlobBuildError::Work(error), left))
}
fn build_verify(work: WorkCounters, budget: WorkBudget) -> Result<(), BlobBuildFailure> {
work.verify(budget)
.map_err(|error| build_failed(BlobBuildError::Work(error), work))
}
pub type BlobBuildFailure = OperationFailure<BlobBuildError>;
#[derive(Debug, Error)]
pub enum BlobBuildError {
#[error("blob build was cancelled")]
Cancelled,
#[error("blob build options are invalid")]
InvalidOptions,
#[error("blob index page has {observed} bytes; maximum is {maximum}")]
PageTooLarge {
observed: u64,
maximum: u64,
},
#[error("blob build scratch allocation failed")]
AllocationFailed,
#[error("blob build index state is invalid")]
InvalidIndexState,
#[error("blob source exceeds its admitted logical size")]
TooLarge,
#[error("blob source failed: {0}")]
Source(#[source] std::io::Error),
#[error(transparent)]
Encode(#[from] CanonicalDecodeError),
#[error(transparent)]
Storage(#[from] ObjectStoreError),
#[error(transparent)]
Work(#[from] WorkError),
}
fn build_failed(error: BlobBuildError, work: WorkCounters) -> BlobBuildFailure {
OperationFailure::new(error, work)
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BlobRead {
pub bytes: Bytes,
pub work: WorkCounters,
pub next_residency: Option<ResidencyHint>,
}
pub fn read_blob_range<S: ObjectStore>(
store: &S,
root: ObjectId,
range: ByteRange,
limits: DecodeLimits,
budget: WorkBudget,
) -> Result<BlobRead, BlobReadFailure> {
let mut machine = BlobRangeMachine::new(root, range, limits, budget)?;
frontier::drive_sync(store, &mut machine)
}
pub async fn read_blob_range_async<S: AsyncObjectStore>(
store: &S,
root: ObjectId,
range: ByteRange,
limits: DecodeLimits,
budget: WorkBudget,
cancellation: &CancellationToken,
) -> Result<BlobRead, BlobReadFailure> {
if cancellation.is_cancelled() {
return Err(failed(BlobReadError::Cancelled, WorkCounters::default()));
}
let mut machine = BlobRangeMachine::new(root, range, limits, budget)?;
drive_blob_async(store, &mut machine, cancellation).await
}
async fn drive_blob_async<S: AsyncObjectStore>(
store: &S,
machine: &mut BlobRangeMachine,
cancellation: &CancellationToken,
) -> Result<BlobRead, BlobReadFailure> {
loop {
if cancellation.is_cancelled() {
return Err(machine.cancelled());
}
if let Some(output) = machine.complete()? {
return Ok(output);
}
if let Some(batch) = machine.prepare_chunk_batch()? {
let receipt =
AsyncObjectStore::read_many(store, &batch.requests, batch.remaining, cancellation)
.await
.map_err(|failure| machine.batch_storage_failure(batch.prospective, failure))?;
machine.accept_chunk_batch(batch, receipt)?;
continue;
}
let request = machine.prepare_read()?;
let key = (request.page.kind == ObjectKind::Blob)
.then(|| DecodedCacheKey::new::<BlobPage>(request.page, machine.limits));
if let Some(key) = key {
let cached = store
.decoded_cache_get(key)
.map_err(|error| failed(error.into(), machine.work))?;
if let Some(cached) = cached {
machine.accept_cached_page(cached)?;
continue;
}
}
let receipt = AsyncObjectStore::read(
store,
request.page,
request.maximum_bytes,
request.remaining,
cancellation,
)
.await
.map_err(|failure| machine.storage_failure(request.prospective, failure))?;
machine.accept_read_with_cache(store, key, request.prospective, &receipt)?;
}
}
struct PendingChunkBatch {
chunks: Vec<BlobChunkRef>,
requests: Vec<ObjectReadRequest>,
scratch_bytes: u64,
remaining: WorkBudget,
prospective: WorkCounters,
}
#[derive(Clone, Copy)]
enum PendingBlobRead {
Page {
page: ObjectId,
expected_first: u64,
expected_end: Option<u64>,
height: u16,
},
Chunk(BlobChunkRef),
}
struct BlobRangeMachine {
range: ByteRange,
range_end: u64,
limits: DecodeLimits,
budget: WorkBudget,
allocation: usize,
maximum_pending: usize,
pending: Vec<PendingBlobRead>,
pending_capacity: LogicalVecCapacity,
awaiting: Option<PendingBlobRead>,
visited: VisitedObjectSet,
output: Vec<u8>,
work: WorkCounters,
allocations: AllocationLedger,
successor: Option<(u64, ResidencyHint)>,
}
impl BlobRangeMachine {
fn new(
root: ObjectId,
range: ByteRange,
limits: DecodeLimits,
budget: WorkBudget,
) -> Result<Self, BlobReadFailure> {
if root.kind != ObjectKind::Blob {
return Err(failed(
BlobReadError::WrongRootKind,
WorkCounters::default(),
));
}
let range_end = range
.offset
.checked_add(range.length)
.ok_or_else(|| failed(BlobReadError::InvalidRange, WorkCounters::default()))?;
let allocation = usize::try_from(range.length)
.map_err(|_| failed(BlobReadError::InvalidRange, WorkCounters::default()))?;
if !limits.page_limits_valid(1) {
return Err(failed(
BlobReadError::InvalidLimits,
WorkCounters::default(),
));
}
let maximum_visited = usize::try_from(limits.maximum_visited_pages)
.map_err(|_| failed(BlobReadError::InvalidLimits, WorkCounters::default()))?;
let maximum_page_items = usize::try_from(limits.maximum_page_items)
.map_err(|_| failed(BlobReadError::InvalidLimits, WorkCounters::default()))?;
let maximum_pending = maximum_visited
.checked_add(maximum_page_items)
.ok_or_else(|| failed(BlobReadError::InvalidLimits, WorkCounters::default()))?;
let mut allocations = AllocationLedger::default();
let mut work = WorkCounters::default();
allocations
.claim_bytes(range.length, u64::from(allocation != 0), &mut work, budget)
.map_err(|error| failed(error.into(), work))?;
let visited = VisitedObjectSet::new(maximum_visited, &mut allocations, &mut work, budget)
.map_err(|error| failed(error.into(), work))?;
let mut pending = Vec::new();
let mut pending_capacity = LogicalVecCapacity::default();
pending_capacity
.ensure_for_push(
&mut pending,
maximum_pending,
&mut allocations,
&mut work,
budget,
)
.map_err(|error| failed(error.into(), work))?;
pending.push(PendingBlobRead::Page {
page: root,
expected_first: 0,
expected_end: None,
height: 1,
});
let mut output = Vec::new();
if output.try_reserve_exact(allocation).is_err() {
return Err(failed(BlobReadError::AllocationFailed, work));
}
Ok(Self {
range,
range_end,
limits,
budget,
allocation,
maximum_pending,
pending,
pending_capacity,
awaiting: None,
visited,
output,
work,
allocations,
successor: None,
})
}
fn finish(&mut self) -> Result<BlobRead, BlobReadFailure> {
if self.output.len() != self.allocation {
return Err(failed(BlobReadError::IncompleteCoverage, self.work));
}
Ok(BlobRead {
bytes: Bytes::from(std::mem::take(&mut self.output)),
work: self.work,
next_residency: self.successor.map(|(_, hint)| hint),
})
}
fn prepare_read(&mut self) -> Result<frontier::ReadRequest, BlobReadFailure> {
let pending = self
.pending
.pop()
.ok_or_else(|| failed(BlobReadError::TraversalState, self.work))?;
let (object, maximum_bytes) = match pending {
PendingBlobRead::Page { page, height, .. } => {
if height > self.limits.maximum_page_height {
return Err(failed(BlobReadError::CycleOrHeight, self.work));
}
let visited = self
.visited
.insert(page, &mut self.allocations, &mut self.work, self.budget)
.map_err(|error| failed(error.into(), self.work))?;
if !visited.inserted {
return Err(failed(BlobReadError::CycleOrHeight, self.work));
}
(page, self.limits.maximum_page_object_bytes())
}
PendingBlobRead::Chunk(chunk) => (chunk.chunk, chunk.end_offset - chunk.first_offset),
};
let mut remaining = self
.work
.remaining(self.budget)
.map_err(|error| failed(error.into(), self.work))?;
remaining.peak_allocation_bytes = self
.budget
.peak_allocation_bytes
.checked_sub(self.allocations.live_bytes())
.ok_or_else(|| failed(BlobReadError::Work(WorkError::Overflow), self.work))?;
self.awaiting = Some(pending);
Ok(frontier::ReadRequest {
page: object,
maximum_bytes,
remaining,
prospective: self.work,
})
}
fn prepare_chunk_batch(&mut self) -> Result<Option<PendingChunkBatch>, BlobReadFailure> {
let chunk_count = self
.pending
.iter()
.rev()
.take_while(|pending| matches!(pending, PendingBlobRead::Chunk(_)))
.count();
if chunk_count < 2 {
return Ok(None);
}
let scratch_bytes = u64::try_from(chunk_count)
.unwrap_or(u64::MAX)
.checked_mul(
u64::try_from(size_of::<BlobChunkRef>() + size_of::<ObjectReadRequest>())
.unwrap_or(u64::MAX),
)
.ok_or_else(|| failed(BlobReadError::Work(WorkError::Overflow), self.work))?;
let live_bytes = self
.allocations
.live_bytes()
.checked_add(scratch_bytes)
.ok_or_else(|| failed(BlobReadError::Work(WorkError::Overflow), self.work))?;
let mut prospective = self.work;
prospective.allocation_operations = prospective
.allocation_operations
.checked_add(2)
.ok_or_else(|| failed(BlobReadError::Work(WorkError::Overflow), self.work))?;
prospective.peak_allocation_bytes = prospective.peak_allocation_bytes.max(live_bytes);
verify_work(prospective, self.budget)?;
let mut chunks = Vec::new();
let mut requests = Vec::new();
chunks
.try_reserve_exact(chunk_count)
.map_err(|_| failed(BlobReadError::AllocationFailed, self.work))?;
requests
.try_reserve_exact(chunk_count)
.map_err(|_| failed(BlobReadError::AllocationFailed, self.work))?;
for pending in self.pending.iter().rev().take(chunk_count) {
let PendingBlobRead::Chunk(chunk) = pending else {
return Err(failed(BlobReadError::TraversalState, self.work));
};
requests.push(ObjectReadRequest {
object_id: chunk.chunk,
maximum_bytes: chunk.end_offset - chunk.first_offset,
});
chunks.push(*chunk);
}
let mut remaining = prospective
.remaining(self.budget)
.map_err(|error| failed(error.into(), prospective))?;
remaining.peak_allocation_bytes = self
.budget
.peak_allocation_bytes
.checked_sub(live_bytes)
.ok_or_else(|| failed(BlobReadError::Work(WorkError::Overflow), prospective))?;
Ok(Some(PendingChunkBatch {
chunks,
requests,
scratch_bytes,
remaining,
prospective,
}))
}
fn accept_chunk_batch(
&mut self,
batch: PendingChunkBatch,
receipt: ObjectReceipt<Vec<ObjectRead>>,
) -> Result<(), BlobReadFailure> {
if receipt.value.len() != batch.chunks.len() {
return Err(failed(BlobReadError::TraversalState, batch.prospective));
}
let retained = self
.pending
.len()
.checked_sub(batch.chunks.len())
.ok_or_else(|| failed(BlobReadError::TraversalState, batch.prospective))?;
if self.pending.get(retained..).is_none_or(|pending| {
!pending
.iter()
.rev()
.zip(&batch.chunks)
.all(|(pending, chunk)| matches!(pending, PendingBlobRead::Chunk(found) if found == chunk))
}) {
return Err(failed(BlobReadError::TraversalState, batch.prospective));
}
self.pending.truncate(retained);
self.work = merge_blob_backend_work(
batch.prospective,
receipt.work,
self.allocations.live_bytes(),
)
.map_err(|error| failed(error.into(), batch.prospective))?;
let retained_bytes = receipt.value.iter().try_fold(0_u64, |total, value| {
total.checked_add(match value.retention {
ObjectReadRetention::Shared => 0,
ObjectReadRetention::Owned { logical_bytes } => logical_bytes,
})
});
let retained_bytes = retained_bytes
.ok_or_else(|| failed(BlobReadError::Work(WorkError::Overflow), self.work))?;
let simultaneous = self
.allocations
.live_bytes()
.checked_add(batch.scratch_bytes)
.and_then(|bytes| bytes.checked_add(retained_bytes))
.ok_or_else(|| failed(BlobReadError::Work(WorkError::Overflow), self.work))?;
self.work.peak_allocation_bytes = self.work.peak_allocation_bytes.max(simultaneous);
verify_work(self.work, self.budget)?;
for (chunk, value) in batch.chunks.into_iter().zip(receipt.value) {
self.accept_chunk(
&ObjectReceipt {
value,
work: WorkCounters::default(),
},
chunk,
)?;
}
Ok(())
}
fn batch_storage_failure(
&self,
prospective: WorkCounters,
failure: crate::storage::ObjectFailure,
) -> BlobReadFailure {
match merge_blob_backend_work(prospective, *failure.work, self.allocations.live_bytes()) {
Ok(combined) => failed(failure.error.into(), combined),
Err(error) => failed(error.into(), prospective),
}
}
fn accept_read(
&mut self,
prospective: WorkCounters,
receipt: &ObjectReceipt<ObjectRead>,
) -> Result<Option<BlobRead>, BlobReadFailure> {
let pending = self
.awaiting
.take()
.ok_or_else(|| failed(BlobReadError::TraversalState, self.work))?;
self.work =
merge_blob_backend_work(prospective, receipt.work, self.allocations.live_bytes())
.map_err(|error| failed(error.into(), prospective))?;
match pending {
PendingBlobRead::Page {
expected_first,
expected_end,
height,
..
} => self.accept_page(receipt, expected_first, expected_end, height)?,
PendingBlobRead::Chunk(chunk) => self.accept_chunk(receipt, chunk)?,
}
Ok(None)
}
fn accept_read_with_cache<S: AsyncObjectStore>(
&mut self,
store: &S,
key: Option<DecodedCacheKey>,
prospective: WorkCounters,
receipt: &ObjectReceipt<ObjectRead>,
) -> Result<(), BlobReadFailure> {
let pending = self
.awaiting
.take()
.ok_or_else(|| failed(BlobReadError::TraversalState, self.work))?;
self.work =
merge_blob_backend_work(prospective, receipt.work, self.allocations.live_bytes())
.map_err(|error| failed(error.into(), prospective))?;
match pending {
PendingBlobRead::Page {
expected_first,
expected_end,
height,
..
} => self.accept_page_with_cache(
store,
key.ok_or_else(|| failed(BlobReadError::TraversalState, self.work))?,
receipt,
expected_first,
expected_end,
height,
),
PendingBlobRead::Chunk(chunk) => self.accept_chunk(receipt, chunk),
}
}
fn add_items(&mut self, count: u64) -> Result<(), BlobReadFailure> {
self.work = add_work(
self.work,
WorkCounters {
items_examined: count,
..WorkCounters::default()
},
)?;
verify_work(self.work, self.budget)
}
fn accept_page(
&mut self,
receipt: &ObjectReceipt<ObjectRead>,
expected_first: u64,
expected_end: Option<u64>,
height: u16,
) -> Result<(), BlobReadFailure> {
let (page, decoded_bytes, retained_bytes) = self.decode_page(receipt)?;
let result = self.accept_decoded_page(&page, expected_first, expected_end, height);
self.allocations
.release(decoded_bytes)
.map_err(|error| failed(error.into(), self.work))?;
self.allocations
.release(retained_bytes)
.map_err(|error| failed(error.into(), self.work))?;
result
}
fn decode_page(
&mut self,
receipt: &ObjectReceipt<ObjectRead>,
) -> Result<(BlobPage, u64, u64), BlobReadFailure> {
self.work = add_work(
self.work,
WorkCounters {
page_reads: 1,
..WorkCounters::default()
},
)?;
verify_work(self.work, self.budget)?;
let retained_bytes = match receipt.value.retention {
ObjectReadRetention::Shared => 0,
ObjectReadRetention::Owned { logical_bytes } => logical_bytes,
};
self.allocations
.claim_bytes(retained_bytes, 0, &mut self.work, self.budget)
.map_err(|error| failed(error.into(), self.work))?;
let shape = blob_page_decode_shape(&receipt.value, self.limits)
.map_err(|error| failed(error.into(), self.work))?;
self.add_items(u64::try_from(shape.items).unwrap_or(u64::MAX))?;
let item_bytes = match shape.kind {
DecodedPageKind::Leaf => size_of::<BlobChunkRef>(),
DecodedPageKind::Internal => size_of::<BlobChild>(),
};
let decoded_bytes = shape
.items
.checked_mul(item_bytes)
.map(crate::foundation::usize_to_u64)
.ok_or_else(|| failed(BlobReadError::AllocationFailed, self.work))?;
self.allocations
.claim_bytes(
decoded_bytes,
u64::from(decoded_bytes != 0),
&mut self.work,
self.budget,
)
.map_err(|error| failed(error.into(), self.work))?;
let page = decode_blob_page(&receipt.value, self.limits)
.map_err(|error| failed(BlobReadError::Decode(error), self.work))?;
Ok((page, decoded_bytes, retained_bytes))
}
fn accept_cached_page(&mut self, cached: DecodedCacheValue) -> Result<(), BlobReadFailure> {
let pending = self
.awaiting
.take()
.ok_or_else(|| failed(BlobReadError::TraversalState, self.work))?;
let PendingBlobRead::Page {
expected_first,
expected_end,
height,
..
} = pending
else {
return Err(failed(BlobReadError::TraversalState, self.work));
};
self.work = add_work(
self.work,
WorkCounters {
page_reads: 1,
..WorkCounters::default()
},
)?;
verify_work(self.work, self.budget)?;
let page = cached
.value
.downcast::<BlobPage>()
.map_err(|_| failed(ObjectStoreError::Corrupt.into(), self.work))?;
self.accept_decoded_page(&page, expected_first, expected_end, height)
}
fn accept_page_with_cache<S: AsyncObjectStore>(
&mut self,
store: &S,
key: DecodedCacheKey,
receipt: &ObjectReceipt<ObjectRead>,
expected_first: u64,
expected_end: Option<u64>,
height: u16,
) -> Result<(), BlobReadFailure> {
let (page, decoded_bytes, retained_bytes) = self.decode_page(receipt)?;
let admitted = store.decoded_cache_admit(
key,
DecodedCacheValue {
value: Arc::new(page),
logical_bytes: decoded_bytes,
},
);
let result = match admitted {
Ok(DecodedCacheAdmission::Shared(value)) => {
self.allocations
.release(decoded_bytes)
.map_err(|error| failed(error.into(), self.work))?;
let page = value
.value
.downcast::<BlobPage>()
.map_err(|_| failed(ObjectStoreError::Corrupt.into(), self.work))?;
self.accept_decoded_page(&page, expected_first, expected_end, height)
}
Ok(DecodedCacheAdmission::Uncached(value)) => {
let page = value
.value
.downcast::<BlobPage>()
.map_err(|_| failed(ObjectStoreError::Corrupt.into(), self.work))?;
let result = self.accept_decoded_page(&page, expected_first, expected_end, height);
self.allocations
.release(decoded_bytes)
.map_err(|error| failed(error.into(), self.work))?;
result
}
Err(error) => {
self.allocations
.release(decoded_bytes)
.map_err(|release| failed(release.into(), self.work))?;
Err(failed(error.into(), self.work))
}
};
self.allocations
.release(retained_bytes)
.map_err(|error| failed(error.into(), self.work))?;
result
}
fn accept_decoded_page(
&mut self,
page: &BlobPage,
expected_first: u64,
expected_end: Option<u64>,
height: u16,
) -> Result<(), BlobReadFailure> {
if page.first_offset != expected_first
|| expected_end.is_some_and(|end| page.end_offset != end)
{
return Err(failed(BlobReadError::ChildBoundsMismatch, self.work));
}
if expected_end.is_none() && self.range_end > page.end_offset {
return Err(failed(BlobReadError::InvalidRange, self.work));
}
match &page.node {
BlobNode::Internal(children) => self.push_children(children, height)?,
BlobNode::Leaf(chunks) => self.push_chunks(chunks)?,
}
Ok(())
}
fn push_children(
&mut self,
children: &[BlobChild],
height: u16,
) -> Result<(), BlobReadFailure> {
let child_height = height
.checked_add(1)
.ok_or_else(|| failed(BlobReadError::CycleOrHeight, self.work))?;
for child in children.iter().rev() {
self.examine_item()?;
if child.first_offset < self.range_end && child.end_offset > self.range.offset {
self.pending_capacity
.ensure_for_push(
&mut self.pending,
self.maximum_pending,
&mut self.allocations,
&mut self.work,
self.budget,
)
.map_err(|error| failed(error.into(), self.work))?;
self.pending.push(PendingBlobRead::Page {
page: child.page,
expected_first: child.first_offset,
expected_end: Some(child.end_offset),
height: child_height,
});
} else if self.range.length != 0 && child.first_offset >= self.range_end {
self.retain_successor(
child.first_offset,
ObjectReadRequest {
object_id: child.page,
maximum_bytes: self.limits.maximum_page_object_bytes(),
},
);
}
}
Ok(())
}
fn push_chunks(&mut self, chunks: &[BlobChunkRef]) -> Result<(), BlobReadFailure> {
for chunk in chunks.iter().rev() {
self.examine_item()?;
if chunk.first_offset < self.range_end && chunk.end_offset > self.range.offset {
self.pending_capacity
.ensure_for_push(
&mut self.pending,
self.maximum_pending,
&mut self.allocations,
&mut self.work,
self.budget,
)
.map_err(|error| failed(error.into(), self.work))?;
self.pending.push(PendingBlobRead::Chunk(*chunk));
} else if self.range.length != 0 && chunk.first_offset >= self.range_end {
self.retain_successor(
chunk.first_offset,
ObjectReadRequest {
object_id: chunk.chunk,
maximum_bytes: chunk.end_offset - chunk.first_offset,
},
);
}
}
Ok(())
}
fn retain_successor(&mut self, offset: u64, request: ObjectReadRequest) {
if self
.successor
.is_none_or(|(current_offset, _)| offset < current_offset)
{
self.successor = Some((
offset,
ResidencyHint {
request,
reason: ResidencyReason::SequentialRange,
},
));
}
}
fn examine_item(&mut self) -> Result<(), BlobReadFailure> {
self.work = add_work(
self.work,
WorkCounters {
items_examined: 1,
..WorkCounters::default()
},
)?;
verify_work(self.work, self.budget)
}
fn accept_chunk(
&mut self,
receipt: &ObjectReceipt<ObjectRead>,
chunk: BlobChunkRef,
) -> Result<(), BlobReadFailure> {
let retained_bytes = match receipt.value.retention {
ObjectReadRetention::Shared => 0,
ObjectReadRetention::Owned { logical_bytes } => logical_bytes,
};
self.allocations
.claim_bytes(retained_bytes, 0, &mut self.work, self.budget)
.map_err(|error| failed(error.into(), self.work))?;
let chunk_length = chunk.end_offset - chunk.first_offset;
if u64::try_from(receipt.value.len()).unwrap_or(u64::MAX) != chunk_length {
return Err(failed(BlobReadError::ChunkLengthMismatch, self.work));
}
let start = self.range.offset.max(chunk.first_offset) - chunk.first_offset;
let end = self.range_end.min(chunk.end_offset) - chunk.first_offset;
let start =
usize::try_from(start).map_err(|_| failed(BlobReadError::InvalidRange, self.work))?;
let end =
usize::try_from(end).map_err(|_| failed(BlobReadError::InvalidRange, self.work))?;
#[allow(
clippy::indexing_slicing,
reason = "start <= end <= receipt.value.len() by construction: `BlobPage::validate` \
(run by `decode_blob_page` before any chunk reaches this machine) requires \
`contiguous()` to hold, so `chunk.first_offset < chunk.end_offset`; only \
chunks with `chunk.first_offset < self.range_end && chunk.end_offset > \
self.range.offset` are ever queued (see the frontier filter above); \
`self.range_end` is `range.offset + range.length` via `checked_add`, so \
`range.offset <= self.range_end`; together these bound \
`max(range.offset, chunk.first_offset) <= min(range_end, \
chunk.end_offset)`, i.e. start <= end; and the length check just above \
establishes `end <= chunk_length == receipt.value.len()`"
)]
self.output.extend_from_slice(&receipt.value[start..end]);
self.work = add_work(
self.work,
WorkCounters {
bytes_copied: u64::try_from(end - start).unwrap_or(u64::MAX),
output_bytes: u64::try_from(end - start).unwrap_or(u64::MAX),
items_returned: 1,
..WorkCounters::default()
},
)?;
verify_work(self.work, self.budget)?;
self.allocations
.release(retained_bytes)
.map_err(|error| failed(error.into(), self.work))
}
}
impl frontier::Machine for BlobRangeMachine {
type Output = BlobRead;
type Failure = BlobReadFailure;
fn complete(&mut self) -> Result<Option<Self::Output>, Self::Failure> {
if self.pending.is_empty() && self.awaiting.is_none() {
return self.finish().map(Some);
}
Ok(None)
}
fn prepare_read(&mut self) -> Result<frontier::ReadRequest, Self::Failure> {
BlobRangeMachine::prepare_read(self)
}
fn accept(
&mut self,
prospective: WorkCounters,
receipt: &ObjectReceipt<ObjectRead>,
) -> Result<Option<Self::Output>, Self::Failure> {
self.accept_read(prospective, receipt)
}
fn storage_failure(
&self,
prospective: WorkCounters,
failure: crate::storage::ObjectFailure,
) -> Self::Failure {
match merge_blob_backend_work(prospective, *failure.work, self.allocations.live_bytes()) {
Ok(combined) => failed(failure.error.into(), combined),
Err(error) => failed(error.into(), prospective),
}
}
fn cancelled(&self) -> Self::Failure {
failed(BlobReadError::Cancelled, self.work)
}
}
fn add_work(left: WorkCounters, right: WorkCounters) -> Result<WorkCounters, BlobReadFailure> {
left.checked_add(right)
.map_err(|error| failed(BlobReadError::Work(error), left))
}
fn verify_work(work: WorkCounters, budget: WorkBudget) -> Result<(), BlobReadFailure> {
work.verify(budget)
.map_err(|error| failed(BlobReadError::Work(error), work))
}
fn merge_blob_backend_work(
prior: WorkCounters,
mut backend: WorkCounters,
live_bytes: u64,
) -> Result<WorkCounters, WorkError> {
let simultaneous_peak = live_bytes
.checked_add(backend.peak_allocation_bytes)
.ok_or(WorkError::Overflow)?;
backend.peak_allocation_bytes = 0;
let mut merged = prior.checked_add(backend)?;
merged.peak_allocation_bytes = merged.peak_allocation_bytes.max(simultaneous_peak);
Ok(merged)
}
pub type BlobReadFailure = OperationFailure<BlobReadError>;
fn failed(error: BlobReadError, work: WorkCounters) -> BlobReadFailure {
OperationFailure::new(error, work)
}
#[derive(Debug, Error)]
pub enum BlobReadError {
#[error("blob range reading was cancelled")]
Cancelled,
#[error("blob range transition state is invalid")]
TraversalState,
#[error("blob root is not a blob-index page")]
WrongRootKind,
#[error("requested blob range is invalid")]
InvalidRange,
#[error("blob read limits are invalid")]
InvalidLimits,
#[error("blob read scratch allocation failed")]
AllocationFailed,
#[error("blob index contains a cycle or exceeds its height bound")]
CycleOrHeight,
#[error("blob child bounds do not match its page")]
ChildBoundsMismatch,
#[error("blob chunks do not cover the requested range exactly")]
IncompleteCoverage,
#[error("blob chunk length does not match its index span")]
ChunkLengthMismatch,
#[error(transparent)]
Storage(#[from] ObjectStoreError),
#[error(transparent)]
Decode(#[from] CanonicalDecodeError),
#[error(transparent)]
Work(#[from] WorkError),
}
impl From<AllocationError> for BlobReadError {
fn from(error: AllocationError) -> Self {
match error {
AllocationError::Work(error) => Self::Work(error),
AllocationError::Overflow | AllocationError::ReleaseInvariant => {
Self::Work(WorkError::Overflow)
}
AllocationError::InvalidCapacity
| AllocationError::CapacityExceeded
| AllocationError::AllocationFailed => Self::AllocationFailed,
}
}
}
fn count(value: usize) -> Result<u32, CanonicalDecodeError> {
u32::try_from(value).map_err(|_| CanonicalDecodeError::LengthOverflow)
}
fn capacity(value: u32) -> Result<usize, CanonicalDecodeError> {
usize::try_from(value).map_err(|_| CanonicalDecodeError::LengthOverflow)
}
fn invariant(message: &str) -> CanonicalDecodeError {
CanonicalDecodeError::Invariant(message.to_owned())
}
#[cfg(all(test, feature = "memory"))]
#[path = "tests/blob.rs"]
mod tests;