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}