use super::*;
use crate::background::CheckpointAllocationOwner;
use crate::scan::CheckpointRangeReadError;
struct CheckpointDescriptor {
descriptor: CanonicalSegmentDescriptor,
_memory: CheckpointAllocationOwner,
}
impl std::ops::Deref for CheckpointDescriptor {
type Target = CanonicalSegmentDescriptor;
fn deref(&self) -> &Self::Target {
&self.descriptor
}
}
impl CanonicalSegmentReader {
pub(crate) fn checkpoint_node(
&self,
id: NodeId,
work: &CheckpointWorkContext,
) -> Result<Option<CheckpointRecord<NodeRecord>>, CanonicalSegmentError> {
self.checkpoint_record(CanonicalSegmentKind::Nodes, id.0, work, |id, payload| {
checkpoint_decode::node(
id,
payload,
self.property_spills.as_ref(),
self.property_keys(),
work,
)
})
}
pub(crate) fn checkpoint_relationship(
&self,
id: RelId,
work: &CheckpointWorkContext,
) -> Result<Option<CheckpointRecord<RelRecord>>, CanonicalSegmentError> {
self.checkpoint_record(
CanonicalSegmentKind::Relationships,
id.0,
work,
|id, payload| {
checkpoint_decode::relationship(
id,
payload,
self.property_spills.as_ref(),
self.property_keys(),
work,
)
},
)
}
fn checkpoint_record<T>(
&self,
kind: CanonicalSegmentKind,
id: u64,
work: &CheckpointWorkContext,
mut decode: impl FnMut(u64, &[u8]) -> Result<T, CanonicalSegmentError>,
) -> Result<Option<T>, CanonicalSegmentError> {
self.ensure_healthy()?;
let result = (|| {
work.checkpoint()?;
let descriptor = self.checkpoint_descriptor_for_id(kind, id, work)?;
work.checkpoint()?;
let Some(descriptor) = descriptor else {
return Ok(None);
};
self.validate_selected_descriptor(&descriptor)?;
work.checkpoint()?;
let range = SegmentReadRange::new(
self.manifest.artifact_id,
descriptor.segment_id,
descriptor.offset,
descriptor.length,
)
.with_content_digest(descriptor.content_digest);
let bytes = self
.range_reader
.checkpoint_range_with_work_context(&range, work)
.map_err(|error| match error {
CheckpointRangeReadError::Read(error) => CanonicalSegmentError::Read(error),
CheckpointRangeReadError::Work(error) => CanonicalSegmentError::Work(error),
})?;
let mut found = None;
decode_segment_records_control(
&bytes,
self.manifest.generation,
&descriptor,
|record_id, payload| {
let unit = work.start_unit()?;
let ordering = record_id.cmp(&id);
unit.finish();
work.checkpoint()?;
if ordering.is_lt() {
return Ok(CanonicalScanControl::Continue);
}
if ordering.is_eq() {
found = Some(decode(record_id, payload)?);
}
Ok(CanonicalScanControl::Stop)
},
)?;
work.checkpoint()?;
Ok(found)
})();
self.poison_on_physical_failure(&result);
result
}
fn checkpoint_descriptor_for_id(
&self,
kind: CanonicalSegmentKind,
id: u64,
work: &CheckpointWorkContext,
) -> Result<Option<CheckpointDescriptor>, CanonicalSegmentError> {
let mut lower_bound = [0; 9];
lower_bound[0] = kind.tag();
lower_bound[1..].copy_from_slice(&id.to_be_bytes());
let mut selected = None;
let mut decode_error = None;
let mut memory = CheckpointAllocationOwner::default();
let read = self.descriptor_reader.checkpoint_scan_from(
&lower_bound,
GraphDescriptorTreeReadLimits {
max_descriptors: NonZeroU64::MIN,
..Default::default()
},
work,
|key, value| {
memory.reserve(value.len(), work)?;
match CanonicalSegmentDescriptor::decode_descriptor_tree_entry(key, value) {
Ok(descriptor) => selected = Some(descriptor),
Err(error) => decode_error = Some(error),
}
Ok(GraphDescriptorTreeScanControl::Stop)
},
);
if let Some(error) = decode_error {
return Err(error);
}
read.map_err(|error| match error {
GraphDescriptorTreeError::Work(error) => CanonicalSegmentError::Work(error),
error => CanonicalSegmentError::DescriptorTree(error),
})?;
match selected {
Some(descriptor) if descriptor.kind < kind => Err(CanonicalSegmentError::Corrupt(
"canonical checkpoint descriptor seek moved backwards across kinds".into(),
)),
Some(descriptor)
if descriptor.kind == kind
&& descriptor.min_record_id <= id
&& id <= descriptor.max_record_id =>
{
Ok(Some(CheckpointDescriptor {
descriptor,
_memory: memory,
}))
}
Some(_) | None => Ok(None),
}
}
}