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 = "amqp")]
12pub use apalis_amqp;
13#[cfg(feature = "board")]
14pub use apalis_board;
15#[cfg(feature = "sql-mysql")]
16pub use apalis_mysql;
17#[cfg(feature = "sql-postgres")]
18pub use apalis_postgres;
19#[cfg(feature = "redis")]
20pub use apalis_redis;
21#[cfg(feature = "sql-sqlite")]
22pub use apalis_sqlite;
23#[cfg(any(feature = "sql-postgres", feature = "sql-sqlite", feature = "sql-mysql"))]
24pub use spring_sqlx;
25
26pub struct ApalisPlugin;
27
28pub type WorkerRegister = fn(&mut AppBuilder, Monitor) -> Monitor;
29
30#[async_trait]
31impl Plugin for ApalisPlugin {
32    async fn build(&self, app: &mut AppBuilder) {
33        if let Some(builders) = app.get_component::<Vec<WorkerRegister>>() {
34            let mut monitor = Monitor::new();
35            for build_fn in &builders {
36                monitor = build_fn(app, monitor);
37            }
38            if !builders.is_empty() {
39                app.add_scheduler(move |_app| Box::new(Self::schedule(monitor)));
40            }
41        }
42    }
43
44    #[cfg(feature = "redis")]
45    fn dependencies(&self) -> Vec<&str> {
46        vec![std::any::type_name::<spring_redis::RedisPlugin>()]
47    }
48
49    #[cfg(all(
50        any(feature = "sql-postgres", feature = "sql-sqlite", feature = "sql-mysql"),
51        not(feature = "redis")
52    ))]
53    fn dependencies(&self) -> Vec<&str> {
54        vec![std::any::type_name::<spring_sqlx::SqlxPlugin>()]
55    }
56}
57
58impl ApalisPlugin {
59    async fn schedule(monitor: Monitor) -> Result<String> {
60        let _ = monitor.run_with_signal(shutdown_signal()).await;
61        Ok("apalis scheduled finished".to_string())
62    }
63}
64
65async fn shutdown_signal() -> std::io::Result<()> {
66    let _ = signal::shutdown_signal("apalis").await;
67    Ok(())
68}
69
70pub trait ApalisConfigurator {
71    fn add_worker(&mut self, worker_register: WorkerRegister) -> &mut Self;
72}
73
74impl ApalisConfigurator for AppBuilder {
75    fn add_worker(&mut self, worker_register: WorkerRegister) -> &mut Self {
76        if let Some(workers) = self.get_component_ref::<Vec<WorkerRegister>>() {
77            unsafe {
78                let raw_ptr = ComponentRef::into_raw(workers);
79                let workers = &mut *(raw_ptr as *mut Vec<WorkerRegister>);
80                workers.push(worker_register);
81            }
82            self
83        } else {
84            self.add_component(vec![worker_register])
85        }
86    }
87}