Skip to main content

rskit_chain/
executor.rs

1use crate::types::StepProgress;
2use rskit_errors::{AppError, AppResult};
3use std::future::Future;
4use std::marker::PhantomData;
5use std::pin::Pin;
6use std::sync::Arc;
7use tokio_util::sync::CancellationToken;
8
9/// Callback for chain-level progress updates.
10pub type ChainProgressFn = Arc<dyn Fn(StepProgress) + Send + Sync>;
11
12pub(crate) type CleanupAction =
13    Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = AppResult<()>> + Send + 'static>> + Send>;
14
15pub(crate) struct ChainState<O> {
16    pub(crate) output: O,
17    pub(crate) cleanups: Vec<CleanupAction>,
18}
19
20#[derive(Clone)]
21pub(crate) struct ChainContext {
22    pub(crate) progress: Option<ChainProgressFn>,
23    pub(crate) cancel: CancellationToken,
24}
25
26pub(crate) type ChainRunner<I, O> = Arc<
27    dyn Fn(I, ChainContext) -> Pin<Box<dyn Future<Output = AppResult<ChainState<O>>> + Send>>
28        + Send
29        + Sync,
30>;
31
32pub(crate) async fn run_cleanups(mut cleanups: Vec<CleanupAction>) -> AppResult<()> {
33    let mut cleanup_error: Option<AppError> = None;
34    while let Some(cleanup) = cleanups.pop() {
35        if let Err(error) = cleanup().await {
36            cleanup_error = Some(match cleanup_error {
37                Some(existing) => {
38                    existing.context(format!("additional chain cleanup failed: {error}"))
39                }
40                None => error.context("chain cleanup failed"),
41            });
42        }
43    }
44    cleanup_error.map_or(Ok(()), Err)
45}
46
47/// Executes a typed sequence of steps.
48pub struct Chain<I, O> {
49    step_count: usize,
50    runner: ChainRunner<I, O>,
51    _types: PhantomData<fn(I) -> O>,
52}
53
54impl<I, O> Chain<I, O>
55where
56    I: Send + 'static,
57    O: Send + 'static,
58{
59    pub(crate) fn new(step_count: usize, runner: ChainRunner<I, O>) -> Self {
60        Self {
61            step_count,
62            runner,
63            _types: PhantomData,
64        }
65    }
66
67    /// Execute the chain, short-circuiting on the first failed step.
68    pub async fn execute(
69        &self,
70        input: I,
71        progress: Option<ChainProgressFn>,
72        cancel: CancellationToken,
73    ) -> AppResult<O> {
74        let state = (self.runner)(input, ChainContext { progress, cancel }).await?;
75        Ok(state.output)
76    }
77
78    /// Number of steps in the chain.
79    #[must_use]
80    pub fn len(&self) -> usize {
81        self.step_count
82    }
83
84    /// Whether the chain has no steps.
85    #[must_use]
86    pub fn is_empty(&self) -> bool {
87        self.step_count == 0
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use rskit_errors::ErrorCode;
95
96    #[tokio::test]
97    async fn run_cleanups_collects_multiple_failures_in_lifo_order() {
98        let cleanups: Vec<CleanupAction> = vec![
99            Box::new(|| {
100                Box::pin(async { Err(AppError::new(ErrorCode::Internal, "cleanup one failed")) })
101            }),
102            Box::new(|| {
103                Box::pin(async { Err(AppError::new(ErrorCode::Internal, "cleanup two failed")) })
104            }),
105        ];
106
107        let error = run_cleanups(cleanups)
108            .await
109            .expect_err("cleanup failures should be reported");
110
111        assert_eq!(error.code(), ErrorCode::Internal);
112        assert!(error.message().contains("cleanup two failed"));
113        assert!(error.message().contains("cleanup one failed"));
114    }
115
116    #[tokio::test]
117    async fn run_cleanups_succeeds_when_all_actions_succeed() {
118        let cleanups: Vec<CleanupAction> = vec![Box::new(|| Box::pin(async { Ok(()) }))];
119
120        run_cleanups(cleanups)
121            .await
122            .expect("successful cleanups should pass");
123    }
124}