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(
135 layout,
136 workflow,
137 BTreeMap::new(),
138 &BTreeMap::new(),
139 usize::MAX,
140 driver,
141 )
142 .await?;
143 Ok(read_complete_execution_workflow(
144 layout,
145 &workflow.digest(),
146 )?)
147}
148
149pub async fn run_retained_execution_workflow<D: ExecutionWorkflowStageDriver>(
164 layout: &BackupLayoutGuard,
165 expected_workflow: &ArtifactChecksumRecord,
166 retained: &BTreeMap<u64, RetainedStageEvidence>,
167 driver: &mut D,
168) -> Result<ExecutionWorkflowSettlementView, ExecutionWorkflowRunError<D::Error>> {
169 let (workflow, completed) = admit_retained_stages(layout, expected_workflow, retained)?;
170 execute_remaining_stages(layout, &workflow, completed, retained, usize::MAX, driver).await?;
171 Ok(read_complete_execution_workflow(layout, expected_workflow)?)
172}
173
174#[derive(Clone, Debug, Eq, PartialEq)]
176pub enum ExecutionWorkflowAdvance {
177 Paused,
179 Complete(ExecutionWorkflowSettlementView),
182}
183
184pub async fn advance_execution_workflow<D: ExecutionWorkflowStageDriver>(
201 layout: &BackupLayoutGuard,
202 expected_workflow: &ArtifactChecksumRecord,
203 retained: &BTreeMap<u64, RetainedStageEvidence>,
204 stage_limit: NonZeroUsize,
205 driver: &mut D,
206) -> Result<ExecutionWorkflowAdvance, ExecutionWorkflowRunError<D::Error>> {
207 let (workflow, completed) = admit_retained_stages(layout, expected_workflow, retained)?;
208 match execute_remaining_stages(
209 layout,
210 &workflow,
211 completed,
212 retained,
213 stage_limit.get(),
214 driver,
215 )
216 .await?
217 {
218 StageExecutionOutcome::Complete => Ok(ExecutionWorkflowAdvance::Complete(
219 read_complete_execution_workflow(layout, expected_workflow)?,
220 )),
221 StageExecutionOutcome::Paused => Ok(ExecutionWorkflowAdvance::Paused),
222 }
223}
224
225fn admit_retained_stages<E: std::error::Error + 'static>(
226 layout: &BackupLayoutGuard,
227 expected_workflow: &ArtifactChecksumRecord,
228 retained: &BTreeMap<u64, RetainedStageEvidence>,
229) -> Result<
230 (
231 ExecutionWorkflowRecord,
232 BTreeMap<u64, ExecutionStagePredecessorRecord>,
233 ),
234 ExecutionWorkflowRunError<E>,
235> {
236 let workflow = read_execution_workflow(layout, expected_workflow)?;
237 let occupied = read_execution_stage_occupancy(layout, expected_workflow)?;
238 if !retained.keys().copied().eq(occupied.iter().copied()) {
239 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
240 }
241 let mut completed: BTreeMap<u64, ExecutionStagePredecessorRecord> = BTreeMap::new();
242 for (&sequence, evidence) in retained {
243 for successor in evidence.successor_evidence.keys() {
244 let Ok(node) = workflow.allocation().graph().node(*successor) else {
245 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
246 };
247 if !node.depends_on().contains(&sequence) {
248 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
249 }
250 }
251 let stage =
252 ExecutionStageGuard::open(layout, expected_workflow, sequence, &evidence.binding)?;
253 for predecessor in stage.binding().predecessors() {
254 let Some(original) = retained.get(&predecessor.stage_sequence()) else {
255 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
256 };
257 if original.learned_for_successor(sequence) != predecessor.learned_evidence() {
258 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
259 }
260 }
261 let checkpoint = stage.replay_checkpoint(evidence.learned_evidence.clone())?;
262 completed.insert(sequence, checkpoint);
263 }
264 if read_execution_stage_occupancy(layout, expected_workflow)? != occupied {
267 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
268 }
269 Ok((workflow, completed))
270}
271
272enum StageExecutionOutcome {
275 Complete,
276 Paused,
277}
278
279async fn execute_remaining_stages<D: ExecutionWorkflowStageDriver>(
280 layout: &BackupLayoutGuard,
281 workflow: &ExecutionWorkflowRecord,
282 mut completed: BTreeMap<u64, ExecutionStagePredecessorRecord>,
283 retained: &BTreeMap<u64, RetainedStageEvidence>,
284 stage_limit: usize,
285 driver: &mut D,
286) -> Result<StageExecutionOutcome, ExecutionWorkflowRunError<D::Error>> {
287 let mut advanced = 0;
288 for node in workflow.allocation().graph().ordered_nodes() {
289 let sequence = node.operation_sequence();
290 if completed.contains_key(&sequence) {
291 continue;
292 }
293 if advanced == stage_limit {
294 return Ok(StageExecutionOutcome::Paused);
295 }
296 let predecessors = node
297 .depends_on()
298 .iter()
299 .map(|dependency| {
300 let original =
301 completed.get(dependency).ok_or(
302 ExecutionWorkflowRunError::<D::Error>::MissingPredecessor(*dependency),
303 )?;
304 let learned = if let Some(evidence) = retained.get(dependency) {
305 evidence.learned_for_successor(sequence).clone()
306 } else {
307 driver
308 .learned_for_successor(original, sequence)
309 .map_err(ExecutionWorkflowRunError::Driver)?
310 };
311 Ok(ExecutionStagePredecessorRecord::new(
312 *dependency,
313 original.binding().clone(),
314 original.settlement().clone(),
315 learned,
316 ))
317 })
318 .collect::<Result<Vec<_>, ExecutionWorkflowRunError<D::Error>>>()?;
319 let checkpoint =
320 execute_absent_stage(layout, workflow, sequence, predecessors, driver).await?;
321 completed.insert(sequence, checkpoint);
322 advanced += 1;
323 }
324 Ok(StageExecutionOutcome::Complete)
325}
326
327async fn execute_absent_stage<D: ExecutionWorkflowStageDriver>(
328 layout: &BackupLayoutGuard,
329 workflow: &ExecutionWorkflowRecord,
330 sequence: u64,
331 predecessors: Vec<ExecutionStagePredecessorRecord>,
332 driver: &mut D,
333) -> Result<ExecutionStagePredecessorRecord, ExecutionWorkflowRunError<D::Error>> {
334 let plan = driver
335 .bind_stage(workflow, sequence, &predecessors)
336 .map_err(ExecutionWorkflowRunError::Driver)?;
337 let binding = ExecutionStageBindingRecord::new(workflow, sequence, &plan, predecessors)?;
338 let stage = ExecutionStageGuard::prepare(layout, binding, plan)?;
339 let learned = driver
340 .execute_stage(&stage)
341 .await
342 .map_err(ExecutionWorkflowRunError::Driver)?;
343 Ok(stage.reconcile_checkpoint(learned)?)
344}
345
346#[derive(Debug, Error)]
348pub enum ExecutionWorkflowRunError<E: std::error::Error + 'static> {
349 #[error("original workflow predecessor {0} did not complete")]
351 MissingPredecessor(u64),
352 #[error("retained stage identities do not match occupied original stages")]
354 RetainedStageSetMismatch,
355 #[error("retained stage {0} has a different learned predecessor commitment")]
358 RetainedLearnedMismatch(u64),
359 #[error(transparent)]
361 Workflow(#[from] ExecutionWorkflowPersistenceError),
362 #[error(transparent)]
364 Binding(#[from] ExecutionWorkflowError),
365 #[error(transparent)]
367 Preparation(#[from] ExecutionStagePreparationError),
368 #[error("original workflow stage driver failed: {0}")]
370 Driver(#[source] E),
371 #[error(transparent)]
373 Checkpoint(#[from] ExecutionStageCheckpointError),
374}
375
376#[cfg(test)]
377mod tests;