From 4ce64b536c38fd97f8ce6fe5a827a796314c26cd Mon Sep 17 00:00:00 2001
From: lemonsh <lemon@lemonsh.moe>
Date: Wed, 6 Apr 2022 09:02:54 +0200
Subject: [PATCH] Fix cleanup task shutdown

---
 src/main.rs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 9d37265..a52734d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -74,7 +74,7 @@ async fn main() -> anyhow::Result<()> {
     let session_store = PostgresSessionStore::from_client(database.pool()).with_table_name("sessions");
     session_store.migrate().await?;
 
-    let cleanup_task = spawn_session_cleanup_task(&session_store, Duration::from_secs(600));
+    let cleanup_task = spawn_session_cleanup_task(&session_store, Duration::from_secs(600), ctx.subscribe());
     info!("Started session cleanup task");
 
     let web_app = App::new(cfg.clone(), database.clone())?;
@@ -82,6 +82,7 @@ async fn main() -> anyhow::Result<()> {
     info!("Started the web app at http://{}", cfg.listen_addr);
 
     terminate_signal().await;
+    ctx.send(()).unwrap();
 
     cleanup_task
         .await
@@ -94,14 +95,19 @@ async fn main() -> anyhow::Result<()> {
     Ok(())
 }
 
-fn spawn_session_cleanup_task(store: &PostgresSessionStore, period: Duration) -> JoinHandle<()> {
+fn spawn_session_cleanup_task(store: &PostgresSessionStore, period: Duration, mut cancel: broadcast::Receiver<()>) -> JoinHandle<()> {
     let store = store.clone();
     tokio::spawn(async move {
         loop {
-            sleep(period).await;
-            if let Err(error) = store.cleanup().await {
-                error!("cleanup error: {}", error);
+            tokio::select! {
+                _ = sleep(period) => {
+                    if let Err(error) = store.cleanup().await {
+                        error!("Error in cleanup task: {}", error);
+                    }
+                }
+                _ = cancel.recv() => break
             }
         }
+        info!("Cleanup task has been shut down");
     })
 }