spring_apalis/
lib.rs

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}