1 files changed, 40 insertions, 19 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 239b91b..ba654c0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -18,6 +18,7 @@ use rocket::{
     response::Redirect,
     tokio::{
         spawn,
+        task::JoinHandle,
         time::{interval_at, Instant, MissedTickBehavior},
     },
     Request, State,
@@ -166,7 +167,7 @@ fn oid(cfg: &State<Config>) -> Value {
     })
 }
 
-fn spawn_periodic<A, P, F>(context: &'static str, t: StdDuration, p: P, f: A)
+fn spawn_periodic<A, P, F>(context: &'static str, t: StdDuration, p: P, f: A) -> JoinHandle<()>
 where
     A: Fn(P) -> F + Send + Sync + Sized + 'static,
     P: Clone + Send + Sync + 'static,
@@ -183,7 +184,7 @@ where
                 error!("periodic {context} failed: {e}");
             }
         }
-    });
+    })
 }
 
 async fn ensure_invite_admin(db: &Db, cfg: &Config) -> anyhow::Result<()> {
@@ -204,6 +205,28 @@ async fn ensure_invite_admin(db: &Db, cfg: &Config) -> anyhow::Result<()> {
     }
 }
 
+fn periodic_code_prune(db: &Arc<Db>) -> JoinHandle<()> {
+    spawn_periodic("verify code prune", StdDuration::from_secs(5 * 60), Arc::clone(&db), {
+        |db| async move {
+            let tx = db.begin().await?;
+            tx.prune_expired_verify_codes().await?;
+            tx.commit().await?;
+            Ok(())
+        }
+    })
+}
+
+fn periodic_token_prune(config: &Config, db: &Arc<Db>) -> JoinHandle<()> {
+    spawn_periodic("expired token prune", config.prune_expired_interval, Arc::clone(&db), {
+        |db| async move {
+            let tx = db.begin().await?;
+            tx.prune_expired_tokens().await?;
+            tx.commit().await?;
+            Ok(())
+        }
+    })
+}
+
 pub async fn build(
     base: rocket::Rocket<rocket::Build>,
 ) -> anyhow::Result<rocket::Rocket<rocket::Build>> {
@@ -231,28 +254,26 @@ pub async fn build(
         )
         .context("setting up mail notifications")?,
     );
-    spawn_periodic("verify code prune", StdDuration::from_secs(5 * 60), Arc::clone(&db), {
-        |db| async move {
-            let tx = db.begin().await?;
-            tx.prune_expired_verify_codes().await?;
-            tx.commit().await?;
-            Ok(())
-        }
-    });
-    spawn_periodic("expired token prune", config.prune_expired_interval, Arc::clone(&db), {
-        |db| async move {
-            let tx = db.begin().await?;
-            tx.prune_expired_tokens().await?;
-            tx.commit().await?;
-            Ok(())
-        }
-    });
     let rocket = base
         .manage(config)
         .manage(push)
         .manage(mailer)
-        .attach(db)
+        .attach(Arc::clone(&db))
         .attach(DeferredActions)
+        .attach(AdHoc::on_ignite("periodic actions", |rocket| async move {
+            struct AbortOnDrop<T>(JoinHandle<T>);
+
+            impl<T> Drop for AbortOnDrop<T> {
+                fn drop(&mut self) {
+                    self.0.abort()
+                }
+            }
+
+            let config = rocket.state::<Config>().unwrap();
+            let code_prune = periodic_code_prune(&db);
+            let token_prune = periodic_token_prune(&config, &db);
+            rocket.manage((AbortOnDrop(code_prune), AbortOnDrop(token_prune)))
+        }))
         .mount("/", routes![root, settings, oid, auth_auth, force_auth, fxa_client_configuration,])
         .register("/auth/v1", catchers![api::auth::catch_all,])
         .mount(