Skip to main content

rskit_chain/
builder.rs

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
8/// Fluent builder for typed sequential chains.
9pub 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    /// Create a new empty builder.
20    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    /// Add a typed step to the chain.
51    #[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    /// Build the typed chain executor.
138    #[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}