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}