diff options
| -rw-r--r-- | src/lib.rs | 59 |
1 files changed, 40 insertions, 19 deletions
@@ -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( |