1use crate::{
4 model::{
5 artifacts::ArtifactChecksumRecord,
6 execution_workflow::{
7 ExecutionStageBindingRecord, ExecutionStagePredecessorRecord, ExecutionWorkflowError,
8 ExecutionWorkflowRecord,
9 },
10 operation_plan::OperationPlanRecord,
11 },
12 ops::persistence::{
13 BackupLayoutGuard, ExecutionStageCheckpointError, ExecutionStageGuard,
14 ExecutionStagePreparationError, ExecutionWorkflowPersistenceError,
15 ExecutionWorkflowSettlementView, create_execution_workflow,
16 read_complete_execution_workflow, read_execution_stage_occupancy, read_execution_workflow,
17 },
18};
19use std::{collections::BTreeMap, future::Future, num::NonZeroUsize};
20use thiserror::Error;
21
22pub trait ExecutionWorkflowStageDriver {
30 type Error: std::error::Error + 'static;
32
33 fn bind_stage(
37 &mut self,
38 workflow: &ExecutionWorkflowRecord,
39 stage_sequence: u64,
40 predecessors: &[ExecutionStagePredecessorRecord],
41 ) -> Result<OperationPlanRecord, Self::Error>;
42
43 fn execute_stage(
45 &mut self,
46 stage: &ExecutionStageGuard<'_>,
47 ) -> impl Future<Output = Result<ArtifactChecksumRecord, Self::Error>>;
48
49 fn learned_for_successor(
58 &mut self,
59 predecessor: &ExecutionStagePredecessorRecord,
60 _successor_sequence: u64,
61 ) -> Result<ArtifactChecksumRecord, Self::Error> {
62 Ok(predecessor.learned_evidence().clone())
63 }
64}
65
66#[derive(Clone, Debug, Eq, PartialEq)]
72pub struct RetainedStageEvidence {
73 binding: ArtifactChecksumRecord,
74 learned_evidence: ArtifactChecksumRecord,
75 successor_evidence: BTreeMap<u64, ArtifactChecksumRecord>,
76}
77impl RetainedStageEvidence {
78 #[must_use]
80 pub const fn new(
81 binding: ArtifactChecksumRecord,
82 learned_evidence: ArtifactChecksumRecord,
83 ) -> Self {
84 Self {
85 binding,
86 learned_evidence,
87 successor_evidence: BTreeMap::new(),
88 }
89 }
90
91 #[must_use]
95 pub fn with_successor_evidence(
96 mut self,
97 successor_evidence: BTreeMap<u64, ArtifactChecksumRecord>,
98 ) -> Self {
99 self.successor_evidence = successor_evidence;
100 self
101 }
102
103 fn learned_for_successor(&self, sequence: u64) -> &ArtifactChecksumRecord {
104 self.successor_evidence
105 .get(&sequence)
106 .unwrap_or(&self.learned_evidence)
107 }
108}
109
110pub async fn run_new_execution_workflow<D: ExecutionWorkflowStageDriver>(
129 layout: &BackupLayoutGuard,
130 workflow: &ExecutionWorkflowRecord,
131 driver: &mut D,
132) -> Result<ExecutionWorkflowSettlementView, ExecutionWorkflowRunError<D::Error>> {
133 create_execution_workflow(layout, workflow)?;
134 execute_remaining_stages(layout, workflow, BTreeMap::new(), usize::MAX, driver).await?;
135 Ok(read_complete_execution_workflow(
136 layout,
137 &workflow.digest(),
138 )?)
139}
140
141pub async fn run_retained_execution_workflow<D: ExecutionWorkflowStageDriver>(
156 layout: &BackupLayoutGuard,
157 expected_workflow: &ArtifactChecksumRecord,
158 retained: &BTreeMap<u64, RetainedStageEvidence>,
159 driver: &mut D,
160) -> Result<ExecutionWorkflowSettlementView, ExecutionWorkflowRunError<D::Error>> {
161 let (workflow, completed) = admit_retained_stages(layout, expected_workflow, retained)?;
162 execute_remaining_stages(layout, &workflow, completed, usize::MAX, driver).await?;
163 Ok(read_complete_execution_workflow(layout, expected_workflow)?)
164}
165
166#[derive(Clone, Debug, Eq, PartialEq)]
168pub enum ExecutionWorkflowAdvance {
169 Paused,
171 Complete(ExecutionWorkflowSettlementView),
174}
175
176pub async fn advance_execution_workflow<D: ExecutionWorkflowStageDriver>(
193 layout: &BackupLayoutGuard,
194 expected_workflow: &ArtifactChecksumRecord,
195 retained: &BTreeMap<u64, RetainedStageEvidence>,
196 stage_limit: NonZeroUsize,
197 driver: &mut D,
198) -> Result<ExecutionWorkflowAdvance, ExecutionWorkflowRunError<D::Error>> {
199 let (workflow, completed) = admit_retained_stages(layout, expected_workflow, retained)?;
200 match execute_remaining_stages(layout, &workflow, completed, stage_limit.get(), driver).await? {
201 StageExecutionOutcome::Complete => Ok(ExecutionWorkflowAdvance::Complete(
202 read_complete_execution_workflow(layout, expected_workflow)?,
203 )),
204 StageExecutionOutcome::Paused => Ok(ExecutionWorkflowAdvance::Paused),
205 }
206}
207
208struct CompletedStage<'a> {
210 checkpoint: ExecutionStagePredecessorRecord,
211 retained: Option<&'a RetainedStageEvidence>,
212}
213
214fn admit_retained_stages<'a, E: std::error::Error + 'static>(
215 layout: &BackupLayoutGuard,
216 expected_workflow: &ArtifactChecksumRecord,
217 retained: &'a BTreeMap<u64, RetainedStageEvidence>,
218) -> Result<
219 (ExecutionWorkflowRecord, BTreeMap<u64, CompletedStage<'a>>),
220 ExecutionWorkflowRunError<E>,
221> {
222 let workflow = read_execution_workflow(layout, expected_workflow)?;
223 let occupied = read_execution_stage_occupancy(layout, expected_workflow)?;
224 if !retained.keys().copied().eq(occupied.iter().copied()) {
225 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
226 }
227 let mut completed = BTreeMap::new();
228 for (&sequence, evidence) in retained {
229 for successor in evidence.successor_evidence.keys() {
230 let Ok(node) = workflow.allocation().graph().node(*successor) else {
231 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
232 };
233 if !node.depends_on().contains(&sequence) {
234 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
235 }
236 }
237 let stage =
238 ExecutionStageGuard::open(layout, expected_workflow, sequence, &evidence.binding)?;
239 for predecessor in stage.binding().predecessors() {
240 let Some(original) = retained.get(&predecessor.stage_sequence()) else {
241 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
242 };
243 if original.learned_for_successor(sequence) != predecessor.learned_evidence() {
244 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
245 }
246 }
247 let checkpoint = stage.replay_checkpoint(evidence.learned_evidence.clone())?;
248 completed.insert(
249 sequence,
250 CompletedStage {
251 checkpoint,
252 retained: Some(evidence),
253 },
254 );
255 }
256 if read_execution_stage_occupancy(layout, expected_workflow)? != occupied {
259 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
260 }
261 Ok((workflow, completed))
262}
263
264enum StageExecutionOutcome {
267 Complete,
268 Paused,
269}
270
271async fn execute_remaining_stages<D: ExecutionWorkflowStageDriver>(
272 layout: &BackupLayoutGuard,
273 workflow: &ExecutionWorkflowRecord,
274 mut completed: BTreeMap<u64, CompletedStage<'_>>,
275 stage_limit: usize,
276 driver: &mut D,
277) -> Result<StageExecutionOutcome, ExecutionWorkflowRunError<D::Error>> {
278 let mut advanced = 0;
279 for node in workflow.allocation().graph().ordered_nodes() {
280 let sequence = node.operation_sequence();
281 if completed.contains_key(&sequence) {
282 continue;
283 }
284 if advanced == stage_limit {
285 return Ok(StageExecutionOutcome::Paused);
286 }
287 let predecessors = node
288 .depends_on()
289 .iter()
290 .map(|dependency| {
291 let original =
292 completed.get(dependency).ok_or(
293 ExecutionWorkflowRunError::<D::Error>::MissingPredecessor(*dependency),
294 )?;
295 let learned = if let Some(evidence) = original.retained {
296 evidence.learned_for_successor(sequence).clone()
297 } else {
298 driver
299 .learned_for_successor(&original.checkpoint, sequence)
300 .map_err(ExecutionWorkflowRunError::Driver)?
301 };
302 Ok(ExecutionStagePredecessorRecord::new(
303 *dependency,
304 original.checkpoint.binding().clone(),
305 original.checkpoint.settlement().clone(),
306 learned,
307 ))
308 })
309 .collect::<Result<Vec<_>, ExecutionWorkflowRunError<D::Error>>>()?;
310 let checkpoint =
311 execute_absent_stage(layout, workflow, sequence, predecessors, driver).await?;
312 completed.insert(
313 sequence,
314 CompletedStage {
315 checkpoint,
316 retained: None,
317 },
318 );
319 advanced += 1;
320 }
321 Ok(StageExecutionOutcome::Complete)
322}
323
324async fn execute_absent_stage<D: ExecutionWorkflowStageDriver>(
325 layout: &BackupLayoutGuard,
326 workflow: &ExecutionWorkflowRecord,
327 sequence: u64,
328 predecessors: Vec<ExecutionStagePredecessorRecord>,
329 driver: &mut D,
330) -> Result<ExecutionStagePredecessorRecord, ExecutionWorkflowRunError<D::Error>> {
331 let plan = driver
332 .bind_stage(workflow, sequence, &predecessors)
333 .map_err(ExecutionWorkflowRunError::Driver)?;
334 let binding = ExecutionStageBindingRecord::new(workflow, sequence, &plan, predecessors)?;
335 let stage = ExecutionStageGuard::prepare(layout, binding, plan)?;
336 let learned = driver
337 .execute_stage(&stage)
338 .await
339 .map_err(ExecutionWorkflowRunError::Driver)?;
340 Ok(stage.reconcile_checkpoint(learned)?)
341}
342
343#[derive(Debug, Error)]
345pub enum ExecutionWorkflowRunError<E: std::error::Error + 'static> {
346 #[error("original workflow predecessor {0} did not complete")]
348 MissingPredecessor(u64),
349 #[error("retained stage identities do not match occupied original stages")]
351 RetainedStageSetMismatch,
352 #[error("retained stage {0} has a different learned predecessor commitment")]
355 RetainedLearnedMismatch(u64),
356 #[error(transparent)]
358 Workflow(#[from] ExecutionWorkflowPersistenceError),
359 #[error(transparent)]
361 Binding(#[from] ExecutionWorkflowError),
362 #[error(transparent)]
364 Preparation(#[from] ExecutionStagePreparationError),
365 #[error("original workflow stage driver failed: {0}")]
367 Driver(#[source] E),
368 #[error(transparent)]
370 Checkpoint(#[from] ExecutionStageCheckpointError),
371}
372
373#[cfg(test)]
374mod tests;