1use apalis::prelude::Monitor;
2use spring::{
3 app::AppBuilder,
4 async_trait,
5 error::Result,
6 plugin::{component::ComponentRef, ComponentRegistry, MutableComponentRegistry, Plugin},
7 signal,
8};
9
10pub use apalis;
11#[cfg(feature = "redis")]
12pub use apalis_redis;
13#[cfg(any(
14 feature = "sql-postgres",
15 feature = "sql-sqlite",
16 feature = "sql-mysql"
17))]
18pub use apalis_sql;
19
20pub struct ApalisPlugin;
21
22pub type WorkerRegister = fn(&mut AppBuilder, Monitor) -> Monitor;
23
24#[async_trait]
25impl Plugin for ApalisPlugin {
26 async fn build(&self, app: &mut AppBuilder) {
27 if let Some(builders) = app.get_component::<Vec<WorkerRegister>>() {
28 let mut monitor = Monitor::new();
29 for build_fn in &builders {
30 monitor = build_fn(app, monitor);
31 }
32 if !builders.is_empty() {
33 app.add_scheduler(move |_app| Box::new(Self::schedule(monitor)));
34 }
35 }
36 }
37}
38
39impl ApalisPlugin {
40 async fn schedule(monitor: Monitor) -> Result<String> {
41 let _ = monitor.run_with_signal(shutdown_signal()).await;
42 Ok("apalis scheduled finished".to_string())
43 }
44}
45
46async fn shutdown_signal() -> std::io::Result<()> {
47 let _ = signal::shutdown_signal().await;
48 Ok(())
49}
50
51pub trait ApalisConfigurator {
52 fn add_worker(&mut self, worker_register: WorkerRegister) -> &mut Self;
53}
54
55impl ApalisConfigurator for AppBuilder {
56 fn add_worker(&mut self, worker_register: WorkerRegister) -> &mut Self {
57 if let Some(workers) = self.get_component_ref::<Vec<WorkerRegister>>() {
58 unsafe {
59 let raw_ptr = ComponentRef::into_raw(workers);
60 let workers = &mut *(raw_ptr as *mut Vec<WorkerRegister>);
61 workers.push(worker_register);
62 }
63 self
64 } else {
65 self.add_component(vec![worker_register])
66 }
67 }
68}