make db maintenance raii

Author: john-z-yang

src/store/inflight_activation.rs

@@ -23,13 +23,15 @@ use sqlx::{
     },
 };
 use tokio::{fs, select, task::JoinSet, time};
+use tokio_util::sync::{CancellationToken, DropGuard};
 use tracing::{info, instrument};
 
 use crate::config::Config;
 
 pub struct InflightActivationStore {
     config: InflightActivationStoreConfig,
     shards: Vec<Arc<InflightActivationShard>>,
+    _maintenance_shutdown: DropGuard,
 }
 
 impl InflightActivationStore {
@@ -73,31 +75,37 @@ impl InflightActivationStore {
             shards
         };
 
+        let maintenance_shutdown = CancellationToken::new();
+
         for (i, shard) in shards.iter().cloned().enumerate() {
+            let cancellation = maintenance_shutdown.clone();
             tokio::spawn(async move {
-                let guard = elegant_departure::get_shutdown_guard().shutdown_on_drop();
-
                 let mut timer = time::interval(Duration::from_millis(config.vacuum_interval_ms));
 
                 timer.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
                 loop {
                     select! {
                         _ = timer.tick() => {
                             shard.vacuum_db().await.unwrap_or_else(|_| {
+                                drop(elegant_departure::get_shutdown_guard().shutdown_on_drop());
                                 panic!("Failed to run maintenance vacuum on shard {:}", i)
                             });
                             info!("ran maintenance vacuum on shard {:}", i);
                         }
 
-                        _ = guard.wait() => {
+                        _ = cancellation.cancelled() => {
                             break;
                         }
                     }
                 }
             });
         }
 
-        Ok(Self { config, shards })
+        Ok(Self {
+            config,
+            shards,
+            _maintenance_shutdown: maintenance_shutdown.drop_guard(),
+        })
     }
 
     fn route(&self, id: &str) -> usize {