ic_backup/model/ic_snapshot_upload/plan/
mod.rs1use super::{
4 IcSnapshotUploadError, IcSnapshotUploadReply, IcSnapshotUploadReplyKind,
5 IcSnapshotUploadRequest,
6};
7use crate::model::{
8 artifacts::ArtifactChecksumRecord,
9 attempt_journal::AttemptJournalRecordError,
10 data_stage_plan::{self, DataStagePlanError, LearnedDataStage},
11 effect_graph::EffectGraphError,
12 execution_workflow::{
13 ExecutionStageBindingRecord, ExecutionStagePredecessorRecord, ExecutionWorkflowError,
14 ExecutionWorkflowRecord,
15 },
16 ic_snapshot_data::{ExtentPlanningError, planned_extents},
17 operation_plan::{OperationPlanError, OperationPlanRecord, PlannedOperationRecord},
18};
19use ic_platform::management_canister_types::SnapshotDataKind;
20use thiserror::Error;
21
22pub struct IcSnapshotDataUploadPlan<'workflow, 'request, 'source> {
30 binding: LearnedDataStage<'workflow>,
31 metadata: &'request IcSnapshotUploadRequest<'source>,
32 destination: Vec<u8>,
33 kinds: Vec<SnapshotDataKind>,
34 plan: Option<OperationPlanRecord>,
35}
36#[derive(Clone)]
38pub(crate) enum UploadCompletion {
39 Data,
40 Allocation {
41 sequence: u64,
42 evidence: ArtifactChecksumRecord,
43 },
44}
45
46pub(crate) struct DataUploadPreparation<'workflow, 'request, 'source> {
49 workflow: &'workflow ExecutionWorkflowRecord,
50 stage: &'workflow PlannedOperationRecord,
51 metadata: &'request IcSnapshotUploadRequest<'source>,
52 destination: Vec<u8>,
53 allocation_evidence: ArtifactChecksumRecord,
54 kinds: Vec<SnapshotDataKind>,
55}
56
57impl<'workflow, 'request, 'source> DataUploadPreparation<'workflow, 'request, 'source> {
58 pub(crate) fn new(
59 workflow: &'workflow ExecutionWorkflowRecord,
60 sequence: u64,
61 allocation: &IcSnapshotUploadReply<'request, 'source>,
62 chunk_bytes: u64,
63 ) -> Result<Self, IcSnapshotDataUploadPlanningError> {
64 let metadata = allocation.request();
65 let stage = workflow.stage(sequence)?;
66 let IcSnapshotUploadReplyKind::Metadata { snapshot_id } = allocation.kind() else {
67 return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
68 };
69 metadata.validate_data_destination(snapshot_id)?;
70 if stage.target() != metadata.target()
71 || !super::attempt::source_context_matches(workflow.allocation(), metadata)
72 {
73 return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
74 }
75 let kinds = planned_extents(metadata.source(), chunk_bytes, stage.budget().mutations())?;
76 Ok(Self {
77 workflow,
78 stage,
79 metadata,
80 destination: snapshot_id.clone(),
81 allocation_evidence: allocation.digest(),
82 kinds,
83 })
84 }
85
86 pub(crate) const fn metadata(&self) -> &'request IcSnapshotUploadRequest<'source> {
87 self.metadata
88 }
89 pub(crate) fn destination(&self) -> &[u8] {
90 &self.destination
91 }
92 pub(crate) fn kinds(&self) -> &[SnapshotDataKind] {
93 &self.kinds
94 }
95
96 pub(crate) fn finish(
97 self,
98 bindings: &[ArtifactChecksumRecord],
99 ) -> Result<
100 IcSnapshotDataUploadPlan<'workflow, 'request, 'source>,
101 IcSnapshotDataUploadPlanningError,
102 > {
103 if bindings.len() != self.kinds.len() {
104 return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
105 }
106 let plan = data_stage_plan::build(
107 self.workflow.allocation(),
108 self.stage,
109 bindings.iter().cloned(),
110 )?;
111 Ok(IcSnapshotDataUploadPlan {
112 binding: LearnedDataStage {
113 workflow: self.workflow,
114 sequence: self.stage.operation_sequence(),
115 evidence: self.allocation_evidence,
116 },
117 metadata: self.metadata,
118 destination: self.destination,
119 kinds: self.kinds,
120 plan,
121 })
122 }
123}
124
125impl<'request, 'source> IcSnapshotDataUploadPlan<'_, 'request, 'source> {
126 pub(crate) fn workflow_digest(&self) -> ArtifactChecksumRecord {
127 self.binding.workflow.digest()
128 }
129 pub(crate) fn validate_completion(
130 &self,
131 binding: &ExecutionStageBindingRecord,
132 plan: &OperationPlanRecord,
133 learned: &ArtifactChecksumRecord,
134 ) -> Result<UploadCompletion, IcSnapshotDataUploadPlanningError> {
135 if let Some(expected) = self.plan() {
136 self.validate_binding(binding)?;
137 if plan != expected {
138 return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
139 }
140 Ok(UploadCompletion::Data)
141 } else {
142 binding.validate(self.binding.workflow, plan)?;
143 let [operation] = plan.operations() else {
144 return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
145 };
146 if operation.target() != self.metadata.target()
147 || operation.request() != self.metadata.binding_digest().hash()
148 || learned != &self.binding.evidence
149 {
150 return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
151 }
152 Ok(UploadCompletion::Allocation {
153 sequence: operation.operation_sequence(),
154 evidence: self.binding.evidence.clone(),
155 })
156 }
157 }
158 #[must_use]
160 pub fn kinds(&self) -> &[SnapshotDataKind] {
161 &self.kinds
162 }
163 #[must_use]
165 pub const fn metadata(&self) -> &'request IcSnapshotUploadRequest<'source> {
166 self.metadata
167 }
168 #[must_use]
170 pub fn destination(&self) -> &[u8] {
171 &self.destination
172 }
173 #[must_use]
175 pub const fn plan(&self) -> Option<&OperationPlanRecord> {
176 self.plan.as_ref()
177 }
178
179 pub fn bind(
187 &self,
188 allocation_binding: &ExecutionStageBindingRecord,
189 allocation_plan: &OperationPlanRecord,
190 predecessors: Vec<ExecutionStagePredecessorRecord>,
191 ) -> Result<ExecutionStageBindingRecord, IcSnapshotDataUploadPlanningError> {
192 let plan = self
193 .plan
194 .as_ref()
195 .ok_or(IcSnapshotDataUploadPlanningError::NoDataWrites)?;
196 self.binding
197 .bind(
198 plan,
199 allocation_binding,
200 allocation_plan,
201 self.metadata.binding_digest().hash(),
202 predecessors,
203 )?
204 .ok_or(IcSnapshotDataUploadPlanningError::OriginalMismatch)
205 }
206 pub(crate) fn validate_binding(
207 &self,
208 binding: &ExecutionStageBindingRecord,
209 ) -> Result<(), IcSnapshotDataUploadPlanningError> {
210 let plan = self
211 .plan
212 .as_ref()
213 .ok_or(IcSnapshotDataUploadPlanningError::NoDataWrites)?;
214 if !self.binding.validate(binding, plan)? {
215 return Err(IcSnapshotDataUploadPlanningError::OriginalMismatch);
216 }
217 Ok(())
218 }
219}
220
221#[derive(Debug, Error)]
223pub enum IcSnapshotDataUploadPlanningError {
224 #[error("invalid snapshot data upload chunk size")]
226 InvalidChunkSize,
227 #[error("snapshot data upload count overflow")]
229 CountOverflow,
230 #[error("snapshot data upload needs {required} updates; original allocation is {original}")]
232 InsufficientAllowance {
233 required: u64,
235 original: u32,
237 },
238 #[error("snapshot data upload originals differ")]
240 OriginalMismatch,
241 #[error("snapshot data upload has no data writes to bind")]
243 NoDataWrites,
244 #[error(transparent)]
246 Plan(#[from] OperationPlanError),
247 #[error(transparent)]
249 Graph(#[from] EffectGraphError),
250 #[error(transparent)]
252 Workflow(#[from] ExecutionWorkflowError),
253 #[error(transparent)]
255 Upload(#[from] IcSnapshotUploadError),
256 #[error(transparent)]
258 Budget(#[from] AttemptJournalRecordError),
259}
260impl From<ExtentPlanningError> for IcSnapshotDataUploadPlanningError {
261 fn from(error: ExtentPlanningError) -> Self {
262 match error {
263 ExtentPlanningError::InvalidChunkSize => Self::InvalidChunkSize,
264 ExtentPlanningError::CountOverflow => Self::CountOverflow,
265 ExtentPlanningError::InsufficientAllowance { required, original } => {
266 Self::InsufficientAllowance { required, original }
267 }
268 }
269 }
270}
271
272impl From<DataStagePlanError> for IcSnapshotDataUploadPlanningError {
273 fn from(error: DataStagePlanError) -> Self {
274 match error {
275 DataStagePlanError::CountOverflow => Self::CountOverflow,
276 DataStagePlanError::Graph(error) => Self::Graph(error),
277 DataStagePlanError::Plan(error) => Self::Plan(error),
278 DataStagePlanError::Budget(error) => Self::Budget(error),
279 }
280 }
281}
282
283#[cfg(test)]
284mod tests;