1use crate::executor::{Chain, ChainContext, ChainRunner, ChainState, CleanupAction, run_cleanups};
2use crate::operation::{Step, StepContext};
3use crate::types::{StepProgress, StepStatus};
4use rskit_errors::{AppError, ErrorCode};
5use std::marker::PhantomData;
6use std::sync::Arc;
7
8pub struct ChainBuilder<I, O> {
10 step_count: usize,
11 runner: ChainRunner<I, O>,
12 _types: PhantomData<fn(I) -> O>,
13}
14
15impl<T> ChainBuilder<T, T>
16where
17 T: Send + 'static,
18{
19 pub fn new() -> Self {
21 Self {
22 step_count: 0,
23 runner: Arc::new(|input, _context| {
24 Box::pin(async move {
25 Ok(ChainState {
26 output: input,
27 cleanups: Vec::new(),
28 })
29 })
30 }),
31 _types: PhantomData,
32 }
33 }
34}
35
36impl<T> Default for ChainBuilder<T, T>
37where
38 T: Send + 'static,
39{
40 fn default() -> Self {
41 Self::new()
42 }
43}
44
45impl<I, O> ChainBuilder<I, O>
46where
47 I: Send + 'static,
48 O: Send + 'static,
49{
50 #[must_use]
52 pub fn step<N>(self, step: Step<O, N>) -> ChainBuilder<I, N>
53 where
54 N: Send + 'static,
55 {
56 let previous = self.runner;
57 let step = Arc::new(step);
58 let step_index = self.step_count;
59 let runner: ChainRunner<I, N> = Arc::new(move |input, context: ChainContext| {
60 let previous = previous.clone();
61 let step = step.clone();
62 Box::pin(async move {
63 let mut state = previous(input, context.clone()).await?;
64 if context.cancel.is_cancelled() {
65 let mut error = cancelled_error(step.id());
66 if let Err(cleanup_error) = run_cleanups(state.cleanups).await {
67 error = error.context(format!("{cleanup_error}"));
68 }
69 return Err(error);
70 }
71
72 emit_progress(
73 &context,
74 step_index,
75 step.id(),
76 StepStatus::Running,
77 0,
78 None,
79 );
80 let step_context = StepContext::new(
81 context.cancel.clone(),
82 context.progress.as_ref().map(|progress| {
83 let progress = progress.clone();
84 let step_id = step.id().to_string();
85 Arc::new(move |percent, message| {
86 progress(StepProgress {
87 step_index,
88 step_id: step_id.clone(),
89 status: StepStatus::Running,
90 progress_percent: percent,
91 message,
92 });
93 }) as Arc<dyn Fn(u8, Option<String>) + Send + Sync>
94 }),
95 );
96
97 let output = match step.execute(state.output, step_context).await {
98 Ok(output) => output,
99 Err(error) => {
100 let mut error = error.context(format!("chain step `{}` failed", step.id()));
101 if let Err(cleanup_error) = run_cleanups(state.cleanups).await {
102 error = error.context(format!("{cleanup_error}"));
103 }
104 return Err(error);
105 }
106 };
107
108 emit_progress(
109 &context,
110 step_index,
111 step.id(),
112 StepStatus::Completed,
113 100,
114 None,
115 );
116
117 if let Some(cleanup) = step.cleanup() {
118 state
119 .cleanups
120 .push(Box::new(move || cleanup()) as CleanupAction);
121 }
122
123 Ok(ChainState {
124 output,
125 cleanups: state.cleanups,
126 })
127 })
128 });
129
130 ChainBuilder {
131 step_count: self.step_count + 1,
132 runner,
133 _types: PhantomData,
134 }
135 }
136
137 #[must_use]
139 pub fn build(self) -> Chain<I, O> {
140 Chain::new(self.step_count, self.runner)
141 }
142}
143
144fn emit_progress(
145 context: &ChainContext,
146 step_index: usize,
147 step_id: &str,
148 status: StepStatus,
149 progress_percent: u8,
150 message: Option<String>,
151) {
152 if let Some(progress) = &context.progress {
153 progress(StepProgress {
154 step_index,
155 step_id: step_id.to_string(),
156 status,
157 progress_percent,
158 message,
159 });
160 }
161}
162
163fn cancelled_error(step_id: &str) -> AppError {
164 AppError::new(ErrorCode::Cancelled, "chain cancelled").with_detail("step", step_id.to_string())
165}