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
9pub 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
47pub 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 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 #[must_use]
80 pub fn len(&self) -> usize {
81 self.step_count
82 }
83
84 #[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}