rskit_chain/
operation.rs

1use rskit_errors::AppResult;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5use tokio_util::sync::CancellationToken;
6
7/// Boxed future returned by chain steps.
8pub type StepFuture<T> = Pin<Box<dyn Future<Output = AppResult<T>> + Send + 'static>>;
9
10type StepProgressFn = Arc<dyn Fn(u8, Option<String>) + Send + Sync>;
11type ExecuteFn<I, O> = dyn Fn(I, StepContext) -> StepFuture<O> + Send + Sync;
12type CleanupFn = dyn Fn() -> StepFuture<()> + Send + Sync;
13
14/// Per-step execution context.
15#[derive(Clone)]
16pub struct StepContext {
17    cancel: CancellationToken,
18    progress: Option<StepProgressFn>,
19}
20
21impl StepContext {
22    pub(crate) fn new(cancel: CancellationToken, progress: Option<StepProgressFn>) -> Self {
23        Self { cancel, progress }
24    }
25
26    /// Return the cancellation token shared by the chain execution.
27    #[must_use]
28    pub fn cancellation_token(&self) -> &CancellationToken {
29        &self.cancel
30    }
31
32    /// Whether cancellation has been requested.
33    #[must_use]
34    pub fn is_cancelled(&self) -> bool {
35        self.cancel.is_cancelled()
36    }
37
38    /// Report step-local progress.
39    pub fn progress(&self, percent: u8, message: Option<String>) {
40        if let Some(progress) = &self.progress {
41            progress(percent.min(100), message);
42        }
43    }
44}
45
46/// A typed operation in a sequential chain.
47pub struct Step<I, O> {
48    id: String,
49    name: String,
50    execute: Arc<ExecuteFn<I, O>>,
51    cleanup: Option<Arc<CleanupFn>>,
52}
53
54impl<I, O> Step<I, O>
55where
56    I: Send + 'static,
57    O: Send + 'static,
58{
59    /// Create a typed chain step.
60    #[must_use]
61    pub fn new<F, Fut>(id: impl Into<String>, name: impl Into<String>, execute: F) -> Self
62    where
63        F: Fn(I, StepContext) -> Fut + Send + Sync + 'static,
64        Fut: Future<Output = AppResult<O>> + Send + 'static,
65    {
66        Self {
67            id: id.into(),
68            name: name.into(),
69            execute: Arc::new(move |input, context| Box::pin(execute(input, context))),
70            cleanup: None,
71        }
72    }
73
74    /// Create a typed chain step using `id` as the display name.
75    #[must_use]
76    pub fn from_fn<F, Fut>(id: impl Into<String>, execute: F) -> Self
77    where
78        F: Fn(I, StepContext) -> Fut + Send + Sync + 'static,
79        Fut: Future<Output = AppResult<O>> + Send + 'static,
80    {
81        let id = id.into();
82        Self::new(id.clone(), id, execute)
83    }
84
85    /// Register cleanup to run only if a later step fails or cancellation interrupts the chain.
86    #[must_use]
87    pub fn with_cleanup<F, Fut>(mut self, cleanup: F) -> Self
88    where
89        F: Fn() -> Fut + Send + Sync + 'static,
90        Fut: Future<Output = AppResult<()>> + Send + 'static,
91    {
92        self.cleanup = Some(Arc::new(move || Box::pin(cleanup())));
93        self
94    }
95
96    /// Unique step identifier.
97    #[must_use]
98    pub fn id(&self) -> &str {
99        &self.id
100    }
101
102    /// Human-readable step name.
103    #[must_use]
104    pub fn name(&self) -> &str {
105        &self.name
106    }
107
108    pub(crate) fn execute(&self, input: I, context: StepContext) -> StepFuture<O> {
109        (self.execute)(input, context)
110    }
111
112    pub(crate) fn cleanup(&self) -> Option<Arc<CleanupFn>> {
113        self.cleanup.clone()
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use parking_lot::Mutex;
121    use std::sync::Arc;
122
123    #[tokio::test]
124    async fn step_context_reports_cancellation_and_clamps_progress() {
125        let cancel = CancellationToken::new();
126        let events = Arc::new(Mutex::new(Vec::new()));
127        let captured = events.clone();
128        let context = StepContext::new(
129            cancel.clone(),
130            Some(Arc::new(move |percent, message| {
131                captured.lock().push((percent, message));
132            })),
133        );
134
135        assert!(!context.is_cancelled());
136        assert!(!context.cancellation_token().is_cancelled());
137        context.progress(150, Some("too high".into()));
138        cancel.cancel();
139        assert!(context.is_cancelled());
140
141        let events = events.lock();
142        assert_eq!(events.as_slice(), &[(100, Some("too high".to_string()))]);
143    }
144
145    #[tokio::test]
146    async fn step_accessors_execute_and_cleanup_work() {
147        let step = Step::new("id", "display", |value: u32, _context| async move {
148            Ok(value + 1)
149        })
150        .with_cleanup(|| async { Ok(()) });
151
152        assert_eq!(step.id(), "id");
153        assert_eq!(step.name(), "display");
154        let output = step
155            .execute(1, StepContext::new(CancellationToken::new(), None))
156            .await
157            .expect("step should run");
158        assert_eq!(output, 2);
159        step.cleanup().expect("cleanup should be registered")()
160            .await
161            .expect("cleanup should succeed");
162    }
163}
rskit_chain/operation.rs

rskit_chain/
operation.rs
