use super::{ic_snapshot_data::IcSnapshotDataReply, ic_snapshot_metadata::IcSnapshotMetadataReply};
use ic_management_canister_types::SnapshotDataKind;
use std::fmt;
use thiserror::Error;
pub struct IcSnapshotDataCoverage<'metadata> {
metadata: &'metadata IcSnapshotMetadataReply<'metadata>,
regions: [u64; 3],
chunks: Vec<bool>,
}
impl<'metadata> IcSnapshotDataCoverage<'metadata> {
#[must_use]
pub fn new(metadata: &'metadata IcSnapshotMetadataReply<'metadata>) -> Self {
Self {
metadata,
regions: [0; 3],
chunks: vec![false; metadata.metadata().wasm_chunk_store.len()],
}
}
pub fn admit(
&mut self,
reply: &IcSnapshotDataReply<'_, '_>,
) -> Result<(), IcSnapshotDataCoverageError> {
if reply.request().metadata().digest() != self.metadata.digest() {
return Err(IcSnapshotDataCoverageError::MetadataMismatch);
}
let (index, offset, size) = match reply.request().kind() {
SnapshotDataKind::WasmModule { offset, size } => (0, *offset, *size),
SnapshotDataKind::WasmMemory { offset, size } => (1, *offset, *size),
SnapshotDataKind::StableMemory { offset, size } => (2, *offset, *size),
SnapshotDataKind::WasmChunk { hash } => {
let index = self
.metadata
.metadata()
.wasm_chunk_store
.iter()
.position(|chunk| chunk.hash == *hash)
.ok_or(IcSnapshotDataCoverageError::MetadataMismatch)?;
if self.chunks[index] {
return Err(IcSnapshotDataCoverageError::DuplicateChunk);
}
self.chunks[index] = true;
return Ok(());
}
};
if offset != self.regions[index] {
return Err(IcSnapshotDataCoverageError::NoncontiguousRange);
}
let end = offset
.checked_add(size)
.ok_or(IcSnapshotDataCoverageError::NoncontiguousRange)?;
self.regions[index] = end;
Ok(())
}
#[must_use]
pub const fn metadata(&self) -> &'metadata IcSnapshotMetadataReply<'metadata> {
self.metadata
}
#[must_use]
pub const fn covered_region_bytes(&self) -> [u64; 3] {
self.regions
}
#[must_use]
pub fn covered_chunks(&self) -> usize {
self.chunks.iter().filter(|covered| **covered).count()
}
#[must_use]
pub fn complete(&self) -> Option<IcSnapshotDataCoverageView<'_, 'metadata>> {
let values = self.metadata.metadata();
let sizes = [
values.wasm_module_size,
values.wasm_memory_size,
values.stable_memory_size,
];
(self.regions == sizes && self.chunks.iter().all(|covered| *covered))
.then_some(IcSnapshotDataCoverageView { coverage: self })
}
}
impl fmt::Debug for IcSnapshotDataCoverage<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IcSnapshotDataCoverage")
.field("covered_region_bytes", &self.regions)
.field("covered_chunks", &self.covered_chunks())
.finish_non_exhaustive()
}
}
#[derive(Debug)]
pub struct IcSnapshotDataCoverageView<'coverage, 'metadata> {
coverage: &'coverage IcSnapshotDataCoverage<'metadata>,
}
impl<'metadata> IcSnapshotDataCoverageView<'_, 'metadata> {
#[must_use]
pub const fn metadata(&self) -> &'metadata IcSnapshotMetadataReply<'metadata> {
self.coverage.metadata()
}
}
#[derive(Debug, Error, Eq, PartialEq)]
pub enum IcSnapshotDataCoverageError {
#[error("snapshot data coverage metadata mismatch")]
MetadataMismatch,
#[error("snapshot data coverage requires the next contiguous range")]
NoncontiguousRange,
#[error("snapshot data coverage already contains this chunk")]
DuplicateChunk,
}
#[cfg(test)]
mod tests;