use crate::model::{
data_stage_plan::{self, LearnedDataStage, SnapshotDataPlanningError},
execution_workflow::{
ExecutionStageBindingRecord, ExecutionStagePredecessorRecord, ExecutionWorkflowRecord,
},
ic_snapshot_data::{IcSnapshotDataRequest, planned_extents},
ic_snapshot_metadata::IcSnapshotMetadataReply,
operation_plan::OperationPlanRecord,
};
pub struct IcSnapshotDownloadPlan<'workflow, 'metadata> {
binding: LearnedDataStage<'workflow>,
metadata: &'metadata IcSnapshotMetadataReply<'metadata>,
reads: Option<DataReads<'metadata>>,
}
struct DataReads<'metadata> {
requests: Vec<IcSnapshotDataRequest<'metadata>>,
plan: OperationPlanRecord,
}
impl<'workflow, 'metadata> IcSnapshotDownloadPlan<'workflow, 'metadata> {
pub fn new(
workflow: &'workflow ExecutionWorkflowRecord,
sequence: u64,
metadata: &'metadata IcSnapshotMetadataReply<'metadata>,
chunk_bytes: u64,
) -> Result<Self, SnapshotDataPlanningError> {
let stage = workflow.stage(sequence)?;
if stage.target() != metadata.request().target() {
return Err(SnapshotDataPlanningError::OriginalMismatch);
}
let kinds = planned_extents(metadata, chunk_bytes, stage.budget().mutations())?;
let requests = kinds
.into_iter()
.map(|kind| IcSnapshotDataRequest::new(metadata, kind))
.collect::<Result<Vec<_>, _>>()?;
let plan = data_stage_plan::build(
workflow.allocation(),
stage,
requests.iter().map(IcSnapshotDataRequest::digest),
)?;
Ok(Self {
binding: LearnedDataStage {
workflow,
sequence,
evidence: metadata.digest(),
},
metadata,
reads: plan.map(|plan| DataReads { requests, plan }),
})
}
#[must_use]
pub fn requests(&self) -> &[IcSnapshotDataRequest<'metadata>] {
self.reads
.as_ref()
.map_or(&[], |reads| reads.requests.as_slice())
}
pub(crate) fn matches_metadata(&self, metadata: &IcSnapshotMetadataReply<'_>) -> bool {
self.metadata.digest() == metadata.digest()
}
#[must_use]
pub const fn plan(&self) -> Option<&OperationPlanRecord> {
match &self.reads {
Some(reads) => Some(&reads.plan),
None => None,
}
}
pub(crate) fn validate_binding(
&self,
binding: &ExecutionStageBindingRecord,
) -> Result<(), SnapshotDataPlanningError> {
let plan = self
.plan()
.ok_or(SnapshotDataPlanningError::NoDataOperations)?;
self.binding.validate(binding, plan)
}
pub fn bind(
&self,
metadata_binding: &ExecutionStageBindingRecord,
metadata_plan: &OperationPlanRecord,
predecessors: Vec<ExecutionStagePredecessorRecord>,
) -> Result<ExecutionStageBindingRecord, SnapshotDataPlanningError> {
let plan = self
.plan()
.ok_or(SnapshotDataPlanningError::NoDataOperations)?;
self.binding.bind(
plan,
metadata_binding,
metadata_plan,
self.metadata.request().digest().hash(),
predecessors,
)
}
}
#[cfg(test)]
mod tests;