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 if execute_remaining_stages(
209 layout,
210 &workflow,
211 completed,
212 retained,
213 stage_limit.get(),
214 driver,
215 )
216 .await?
217 {
218 Ok(ExecutionWorkflowAdvance::Complete(
219 read_complete_execution_workflow(layout, expected_workflow)?,
220 ))
221 } else {
222 Ok(ExecutionWorkflowAdvance::Paused)
223 }
224}
225
226fn admit_retained_stages<E: std::error::Error + 'static>(
227 layout: &BackupLayoutGuard,
228 expected_workflow: &ArtifactChecksumRecord,
229 retained: &BTreeMap<u64, RetainedStageEvidence>,
230) -> Result<
231 (
232 ExecutionWorkflowRecord,
233 BTreeMap<u64, ExecutionStagePredecessorRecord>,
234 ),
235 ExecutionWorkflowRunError<E>,
236> {
237 let workflow = read_execution_workflow(layout, expected_workflow)?;
238 let occupied = read_execution_stage_occupancy(layout, expected_workflow)?;
239 if !retained.keys().copied().eq(occupied.iter().copied()) {
240 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
241 }
242 let mut completed: BTreeMap<u64, ExecutionStagePredecessorRecord> = BTreeMap::new();
243 for (&sequence, evidence) in retained {
244 for successor in evidence.successor_evidence.keys() {
245 let Ok(node) = workflow.allocation().graph().node(*successor) else {
246 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
247 };
248 if !node.depends_on().contains(&sequence) {
249 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
250 }
251 }
252 let stage =
253 ExecutionStageGuard::open(layout, expected_workflow, sequence, &evidence.binding)?;
254 for predecessor in stage.binding().predecessors() {
255 let Some(original) = retained.get(&predecessor.stage_sequence()) else {
256 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
257 };
258 if original.learned_for_successor(sequence) != predecessor.learned_evidence() {
259 return Err(ExecutionWorkflowRunError::RetainedLearnedMismatch(sequence));
260 }
261 }
262 let checkpoint = stage.replay_checkpoint(evidence.learned_evidence.clone())?;
263 completed.insert(sequence, checkpoint);
264 }
265 if read_execution_stage_occupancy(layout, expected_workflow)? != occupied {
268 return Err(ExecutionWorkflowRunError::RetainedStageSetMismatch);
269 }
270 Ok((workflow, completed))
271}
272
273async fn execute_remaining_stages<D: ExecutionWorkflowStageDriver>(
276 layout: &BackupLayoutGuard,
277 workflow: &ExecutionWorkflowRecord,
278 mut completed: BTreeMap<u64, ExecutionStagePredecessorRecord>,
279 retained: &BTreeMap<u64, RetainedStageEvidence>,
280 stage_limit: usize,
281 driver: &mut D,
282) -> Result<bool, ExecutionWorkflowRunError<D::Error>> {
283 let mut advanced = 0;
284 for node in workflow.allocation().graph().ordered_nodes() {
285 let sequence = node.operation_sequence();
286 if completed.contains_key(&sequence) {
287 continue;
288 }
289 if advanced == stage_limit {
290 return Ok(false);
291 }
292 let predecessors = node
293 .depends_on()
294 .iter()
295 .map(|dependency| {
296 let original =
297 completed.get(dependency).ok_or(
298 ExecutionWorkflowRunError::<D::Error>::MissingPredecessor(*dependency),
299 )?;
300 let learned = if let Some(evidence) = retained.get(dependency) {
301 evidence.learned_for_successor(sequence).clone()
302 } else {
303 driver
304 .learned_for_successor(original, sequence)
305 .map_err(ExecutionWorkflowRunError::Driver)?
306 };
307 Ok(ExecutionStagePredecessorRecord::new(
308 *dependency,
309 original.binding().clone(),
310 original.settlement().clone(),
311 learned,
312 ))
313 })
314 .collect::<Result<Vec<_>, ExecutionWorkflowRunError<D::Error>>>()?;
315 let checkpoint =
316 execute_absent_stage(layout, workflow, sequence, predecessors, driver).await?;
317 completed.insert(sequence, checkpoint);
318 advanced += 1;
319 }
320 Ok(true)
321}
322
323async fn execute_absent_stage<D: ExecutionWorkflowStageDriver>(
324 layout: &BackupLayoutGuard,
325 workflow: &ExecutionWorkflowRecord,
326 sequence: u64,
327 predecessors: Vec<ExecutionStagePredecessorRecord>,
328 driver: &mut D,
329) -> Result<ExecutionStagePredecessorRecord, ExecutionWorkflowRunError<D::Error>> {
330 let plan = driver
331 .bind_stage(workflow, sequence, &predecessors)
332 .map_err(ExecutionWorkflowRunError::Driver)?;
333 let binding = ExecutionStageBindingRecord::new(workflow, sequence, &plan, predecessors)?;
334 let stage = ExecutionStageGuard::prepare(layout, binding, plan)?;
335 let learned = driver
336 .execute_stage(&stage)
337 .await
338 .map_err(ExecutionWorkflowRunError::Driver)?;
339 Ok(stage.reconcile_checkpoint(learned)?)
340}
341
342#[derive(Debug, Error)]
344pub enum ExecutionWorkflowRunError<E: std::error::Error + 'static> {
345 #[error("original workflow predecessor {0} did not complete")]
347 MissingPredecessor(u64),
348 #[error("retained stage identities do not match occupied original stages")]
350 RetainedStageSetMismatch,
351 #[error("retained stage {0} has a different learned predecessor commitment")]
354 RetainedLearnedMismatch(u64),
355 #[error(transparent)]
357 Workflow(#[from] ExecutionWorkflowPersistenceError),
358 #[error(transparent)]
360 Binding(#[from] ExecutionWorkflowError),
361 #[error(transparent)]
363 Preparation(#[from] ExecutionStagePreparationError),
364 #[error("original workflow stage driver failed: {0}")]
366 Driver(#[source] E),
367 #[error(transparent)]
369 Checkpoint(#[from] ExecutionStageCheckpointError),
370}
371
372#[cfg(test)]
373mod tests;