New job queue: worker registration and leader election

Cargo.toml

@@ -373,3 +373,10 @@ rayon.opt-level = 3
 regalloc2.opt-level = 3
 sha2.opt-level = 3
 sqlx-macros.opt-level = 3
+
+[patch.crates-io]
+sqlx = { git = "https://github.com/launchbadge/sqlx.git", branch = "main" }
+sqlx-core = { git = "https://github.com/launchbadge/sqlx.git", branch = "main" }
+sqlx-macros = { git = "https://github.com/launchbadge/sqlx.git", branch = "main" }
+sqlx-macros-core = { git = "https://github.com/launchbadge/sqlx.git", branch = "main" }
+sqlx-postgres = { git = "https://github.com/launchbadge/sqlx.git", branch = "main" }

crates/cli/src/commands/server.rs

@@ -172,6 +172,8 @@ impl Options {
                 &mailer,
                 homeserver_connection.clone(),
                 url_builder.clone(),
+                shutdown.soft_shutdown_token(),
+                shutdown.task_tracker(),
             )
             .await?;
 

crates/cli/src/commands/worker.rs

@@ -17,15 +17,20 @@ use rand::{
 };
 use tracing::{info, info_span};
 
-use crate::util::{
-    database_pool_from_config, mailer_from_config, site_config_from_config, templates_from_config,
+use crate::{
+    shutdown::ShutdownManager,
+    util::{
+        database_pool_from_config, mailer_from_config, site_config_from_config,
+        templates_from_config,
+    },
 };
 
 #[derive(Parser, Debug, Default)]
 pub(super) struct Options {}
 
 impl Options {
     pub async fn run(self, figment: &Figment) -> anyhow::Result<ExitCode> {
+        let shutdown = ShutdownManager::new()?;
         let span = info_span!("cli.worker.init").entered();
         let config = AppConfig::extract(figment)?;
 
@@ -71,11 +76,35 @@ impl Options {
         let worker_name = Alphanumeric.sample_string(&mut rng, 10);
 
         info!(worker_name, "Starting task scheduler");
-        let monitor = mas_tasks::init(&worker_name, &pool, &mailer, conn, url_builder).await?;
-
+        let monitor = mas_tasks::init(
+            &worker_name,
+            &pool,
+            &mailer,
+            conn,
+            url_builder,
+            shutdown.soft_shutdown_token(),
+            shutdown.task_tracker(),
+        )
+        .await?;
+
+        // XXX: The monitor from apalis is a bit annoying to use for graceful shutdowns,
+        // ideally we'd just give it a cancellation token
+        let shutdown_future = shutdown.soft_shutdown_token().cancelled_owned();
+        shutdown.task_tracker().spawn(async move {
+            if let Err(e) = monitor
+                .run_with_signal(async move {
+                    shutdown_future.await;
+                    Ok(())
+                })
+                .await
+            {
+                tracing::error!(error = &e as &dyn std::error::Error, "Task worker failed");
+            }
+        });
         span.exit();
 
-        monitor.run().await?;
+        shutdown.run().await;
+
         Ok(ExitCode::SUCCESS)
     }
 }