use super::{
IcSnapshotUploadError, IcSnapshotUploadReply, IcSnapshotUploadReplyKind,
IcSnapshotUploadRequest,
};
use crate::model::{
artifacts::ArtifactChecksumRecord,
attempt_journal::AttemptJournalRecordError,
data_stage_plan::{self, DataStagePlanError, LearnedDataStage},
effect_graph::EffectGraphError,
execution_workflow::{
ExecutionStageBindingRecord, ExecutionStagePredecessorRecord, ExecutionWorkflowError,
ExecutionWorkflowRecord,
},
ic_snapshot_data::{ExtentPlanningError, planned_extents},
operation_plan::{OperationPlanError, OperationPlanRecord, PlannedOperationRecord},
};
use ic_platform::management_canister_types::SnapshotDataKind;
use thiserror::Error;
pub struct IcSnapshotDataUploadPlan<'workflow, 'request, 'source> {
binding: LearnedDataStage<'workflow>,
metadata: &'request IcSnapshotUploadRequest<'source>,
destination: Vec<u8>,
kinds: Vec<SnapshotDataKind>,
plan: Option<OperationPlanRecord>,
}
pub(crate) struct DataUploadPreparation<'workflow, 'request, 'source> {
workflow: &'workflow ExecutionWorkflowRecord,
stage: &'workflow PlannedOperationRecord,
metadata: &'request IcSnapshotUploadRequest<'source>,
destination: Vec<u8>,
allocation_evidence: ArtifactChecksumRecord,
kinds: Vec<SnapshotDataKind>,
}
impl<'workflow, 'request, 'source> DataUploadPreparation<'workflow, 'request, 'source> {
pub(crate) fn new(
workflow: &'workflow ExecutionWorkflowRecord,
sequence: u64,
allocation: &IcSnapshotUploadReply<'request, 'source>,
chunk_bytes: u64,
) -> Result<Self, IcSnapshotDataUploadPlanningError> {
let metadata = allocation.request();
let stage = workflow.stage(sequence)?;
let IcSnapshotUploadReplyKind::Metadata { snapshot_id } = allocation.kind() else {
return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
};
metadata.validate_data_destination(snapshot_id)?;
if stage.target() != metadata.target()
|| !super::attempt::source_context_matches(workflow.allocation(), metadata)
{
return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
}
let kinds = planned_extents(metadata.source(), chunk_bytes, stage.budget().mutations())?;
Ok(Self {
workflow,
stage,
metadata,
destination: snapshot_id.clone(),
allocation_evidence: allocation.digest(),
kinds,
})
}
pub(crate) const fn metadata(&self) -> &'request IcSnapshotUploadRequest<'source> {
self.metadata
}
pub(crate) fn destination(&self) -> &[u8] {
&self.destination
}
pub(crate) fn kinds(&self) -> &[SnapshotDataKind] {
&self.kinds
}
pub(crate) fn finish(
self,
bindings: &[ArtifactChecksumRecord],
) -> Result<
IcSnapshotDataUploadPlan<'workflow, 'request, 'source>,
IcSnapshotDataUploadPlanningError,
> {
if bindings.len() != self.kinds.len() {
return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
}
let plan = data_stage_plan::build(
self.workflow.allocation(),
self.stage,
bindings.iter().cloned(),
)?;
Ok(IcSnapshotDataUploadPlan {
binding: LearnedDataStage {
workflow: self.workflow,
sequence: self.stage.operation_sequence(),
evidence: self.allocation_evidence,
},
metadata: self.metadata,
destination: self.destination,
kinds: self.kinds,
plan,
})
}
}
impl<'request, 'source> IcSnapshotDataUploadPlan<'_, 'request, 'source> {
pub(crate) fn workflow_digest(&self) -> ArtifactChecksumRecord {
self.binding.workflow.digest()
}
pub(crate) fn validate_completion(
&self,
binding: &ExecutionStageBindingRecord,
plan: &OperationPlanRecord,
learned: &ArtifactChecksumRecord,
) -> Result<(), IcSnapshotDataUploadPlanningError> {
if let Some(expected) = self.plan() {
self.validate_binding(binding)?;
if plan != expected {
return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
}
} else {
binding.validate(self.binding.workflow, plan)?;
if plan.operations().len() != 1
|| plan.operations()[0].target() != self.metadata.target()
|| plan.operations()[0].request() != self.metadata.binding_digest().hash()
|| learned != &self.binding.evidence
{
return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
}
}
Ok(())
}
#[must_use]
pub fn kinds(&self) -> &[SnapshotDataKind] {
&self.kinds
}
#[must_use]
pub const fn metadata(&self) -> &'request IcSnapshotUploadRequest<'source> {
self.metadata
}
#[must_use]
pub fn destination(&self) -> &[u8] {
&self.destination
}
#[must_use]
pub const fn plan(&self) -> Option<&OperationPlanRecord> {
self.plan.as_ref()
}
pub fn bind(
&self,
allocation_binding: &ExecutionStageBindingRecord,
allocation_plan: &OperationPlanRecord,
predecessors: Vec<ExecutionStagePredecessorRecord>,
) -> Result<ExecutionStageBindingRecord, IcSnapshotDataUploadPlanningError> {
let plan = self
.plan
.as_ref()
.ok_or(IcSnapshotDataUploadPlanningError::NoDataWrites)?;
self.binding
.bind(
plan,
allocation_binding,
allocation_plan,
self.metadata.binding_digest().hash(),
predecessors,
)?
.ok_or(IcSnapshotDataUploadPlanningError::OriginalMismatch)
}
pub(crate) fn validate_binding(
&self,
binding: &ExecutionStageBindingRecord,
) -> Result<(), IcSnapshotDataUploadPlanningError> {
let plan = self
.plan
.as_ref()
.ok_or(IcSnapshotDataUploadPlanningError::NoDataWrites)?;
if !self.binding.validate(binding, plan)? {
return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
}
Ok(())
}
}
#[derive(Debug, Error)]
pub enum IcSnapshotDataUploadPlanningError {
#[error("invalid snapshot data upload chunk size")]
InvalidChunkSize,
#[error("snapshot data upload count overflow")]
CountOverflow,
#[error("snapshot data upload needs {required} updates; original allocation is {original}")]
InsufficientAllowance {
required: u64,
original: u32,
},
#[error("snapshot data upload originals differ")]
OriginalMismatch,
#[error("snapshot data upload has no data writes to bind")]
NoDataWrites,
#[error(transparent)]
Plan(#[from] OperationPlanError),
#[error(transparent)]
Graph(#[from] EffectGraphError),
#[error(transparent)]
Workflow(#[from] ExecutionWorkflowError),
#[error(transparent)]
Upload(#[from] IcSnapshotUploadError),
#[error(transparent)]
Budget(#[from] AttemptJournalRecordError),
}
impl From<ExtentPlanningError> for IcSnapshotDataUploadPlanningError {
fn from(error: ExtentPlanningError) -> Self {
match error {
ExtentPlanningError::InvalidChunkSize => Self::InvalidChunkSize,
ExtentPlanningError::CountOverflow => Self::CountOverflow,
ExtentPlanningError::InsufficientAllowance { required, original } => {
Self::InsufficientAllowance { required, original }
}
}
}
}
impl From<DataStagePlanError> for IcSnapshotDataUploadPlanningError {
fn from(error: DataStagePlanError) -> Self {
match error {
DataStagePlanError::CountOverflow => Self::CountOverflow,
DataStagePlanError::Graph(error) => Self::Graph(error),
DataStagePlanError::Plan(error) => Self::Plan(error),
DataStagePlanError::Budget(error) => Self::Budget(error),
}
}
}
#[cfg(test)]
mod tests;