Skip to main content

quokka_job/
lib.rs

1use std::{future::Future, pin::Pin, sync::Arc};
2
3pub use error::*;
4
5use quokka_config::{Config, TryFromConfig};
6use quokka_pouch::Pouch;
7use quokka_state::ProvideStateRef;
8
9mod error;
10
11pub struct JobPouch {}
12
13pub trait Job<S> {
14    type Error: std::error::Error;
15
16    fn handle(
17        &self,
18        state: &S,
19    ) -> impl Future<Output = std::result::Result<(), Self::Error>> + Send;
20}
21
22#[derive(Clone, Default)]
23pub struct Jobs<S> {
24    jobs: Vec<Arc<dyn JobBuilder<S> + Send + Sync>>,
25}
26
27impl<S> TryFromConfig for Jobs<S> {
28    type Error = crate::Error;
29
30    async fn try_from_config(_: &Config) -> crate::Result<Self>
31    where
32        Self: Sized,
33    {
34        Ok(Self { jobs: vec![] })
35    }
36}
37
38pub trait JobStateExt {
39    fn jobs(&mut self) -> &mut Jobs<Self>
40    where
41        Self: Sized;
42
43    fn register_job<J: Job<Self> + Send + Sync + 'static>(&mut self, job: J)
44    where
45        Self: Sized;
46}
47
48trait JobBuilder<S> {
49    fn run_job(&self, state: &S) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>;
50
51    fn name(&self) -> &'static str;
52}
53
54struct WrapJob<J> {
55    job: Arc<J>,
56}
57
58impl<S> Pouch<S> for JobPouch
59where
60    S: ProvideStateRef<Jobs<S>> + Send + Sync + Clone + 'static,
61{
62    fn from_config(_: &Config) -> quokka_pouch::Result<Self>
63    where
64        Self: Sized,
65    {
66        Ok(Self {})
67    }
68
69    #[tracing::instrument(skip_all)]
70    fn run_setup_job(
71        &self,
72        state: &S,
73    ) -> Pin<Box<dyn Future<Output = quokka_pouch::Result<()>> + Send>> {
74        let state = state.clone();
75
76        Box::pin(async move {
77            let jobs: &Jobs<S> = state.provide_ref();
78
79            for job in &jobs.jobs {
80                if let Err(error) = job.run_job(&state).await {
81                    tracing::error!(?error, job = job.name(), "Unable to run job");
82                }
83            }
84
85            Ok(())
86        })
87    }
88}
89
90impl<S: Send + Sync + Clone + 'static> Jobs<S> {
91    pub fn register_job<J: Job<S> + Send + Sync + 'static>(&mut self, job: J) {
92        self.jobs
93            .push(Arc::new(WrapJob::<J> { job: Arc::new(job) }));
94    }
95}
96
97impl<J, S> JobBuilder<S> for WrapJob<J>
98where
99    J: Job<S> + Send + Sync + 'static,
100    S: Send + Sync + Clone + 'static,
101{
102    fn run_job(&self, state: &S) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
103        let state = state.clone();
104        let job = self.job.clone();
105
106        Box::pin(async move {
107            job.handle(&state)
108                .await
109                .inspect_err(|error| {
110                    tracing::error!(
111                        ?error,
112                        job = std::any::type_name::<J>(),
113                        "Unable to run job"
114                    )
115                })
116                .map_err(|_| Error::JobRunError(std::any::type_name::<J>().to_string()))?;
117
118            Ok(())
119        })
120    }
121
122    fn name(&self) -> &'static str {
123        std::any::type_name::<J>()
124    }
125}
126
127impl<S: Send + Sync + Clone + ProvideStateRef<Jobs<S>> + 'static> JobStateExt for S {
128    fn jobs(&mut self) -> &mut Jobs<Self>
129    where
130        Self: Sized,
131    {
132        self.provide_mut()
133    }
134
135    fn register_job<J: Job<S> + Send + Sync + 'static>(&mut self, job: J) {
136        self.jobs().register_job(job);
137    }
138}