use rskit_errors::AppResult;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
pub type StepFuture<T> = Pin<Box<dyn Future<Output = AppResult<T>> + Send + 'static>>;
type StepProgressFn = Arc<dyn Fn(u8, Option<String>) + Send + Sync>;
type ExecuteFn<I, O> = dyn Fn(I, StepContext) -> StepFuture<O> + Send + Sync;
type CleanupFn = dyn Fn() -> StepFuture<()> + Send + Sync;
#[derive(Clone)]
pub struct StepContext {
cancel: CancellationToken,
progress: Option<StepProgressFn>,
}
impl StepContext {
pub(crate) fn new(cancel: CancellationToken, progress: Option<StepProgressFn>) -> Self {
Self { cancel, progress }
}
#[must_use]
pub fn cancellation_token(&self) -> &CancellationToken {
&self.cancel
}
#[must_use]
pub fn is_cancelled(&self) -> bool {
self.cancel.is_cancelled()
}
pub fn progress(&self, percent: u8, message: Option<String>) {
if let Some(progress) = &self.progress {
progress(percent.min(100), message);
}
}
}
pub struct Step<I, O> {
id: String,
name: String,
execute: Arc<ExecuteFn<I, O>>,
cleanup: Option<Arc<CleanupFn>>,
}
impl<I, O> Step<I, O>
where
I: Send + 'static,
O: Send + 'static,
{
#[must_use]
pub fn new<F, Fut>(id: impl Into<String>, name: impl Into<String>, execute: F) -> Self
where
F: Fn(I, StepContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = AppResult<O>> + Send + 'static,
{
Self {
id: id.into(),
name: name.into(),
execute: Arc::new(move |input, context| Box::pin(execute(input, context))),
cleanup: None,
}
}
#[must_use]
pub fn from_fn<F, Fut>(id: impl Into<String>, execute: F) -> Self
where
F: Fn(I, StepContext) -> Fut + Send + Sync + 'static,
Fut: Future<Output = AppResult<O>> + Send + 'static,
{
let id = id.into();
Self::new(id.clone(), id, execute)
}
#[must_use]
pub fn with_cleanup<F, Fut>(mut self, cleanup: F) -> Self
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = AppResult<()>> + Send + 'static,
{
self.cleanup = Some(Arc::new(move || Box::pin(cleanup())));
self
}
#[must_use]
pub fn id(&self) -> &str {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
pub(crate) fn execute(&self, input: I, context: StepContext) -> StepFuture<O> {
(self.execute)(input, context)
}
pub(crate) fn cleanup(&self) -> Option<Arc<CleanupFn>> {
self.cleanup.clone()
}
}