demonstrate deletion task cleanup (#5)

This extends the examples which use a deletion task for cleanup such that the task is aborted as part of a graceful shutdown.
maxcountryman · Feb 1, 2024 · 51cc758 · 51cc758
1 parent 6ebc561
commit 51cc758
Show file tree

Hide file tree

Showing 7 changed files with 261 additions and 12 deletions.
diff --git a/mongodb-store/README.md b/mongodb-store/README.md
@@ -14,7 +14,9 @@ use std::net::SocketAddr;
 use axum::{response::IntoResponse, routing::get, Router};
 use serde::{Deserialize, Serialize};
 use time::Duration;
+use tokio::{signal, task::AbortHandle};
 use tower_sessions::{Expiry, Session, SessionManagerLayer};
+use tower_sessions_core::ExpiredDeletion;
 use tower_sessions_mongodb_store::{mongodb::Client, MongoDBStore};
 
 const COUNTER_KEY: &str = "counter";
@@ -33,6 +35,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let database_url = std::option_env!("DATABASE_URL").expect("Missing DATABASE_URL.");
     let client = Client::with_uri_str(database_url).await?;
     let session_store = MongoDBStore::new(client, "tower-sessions".to_string());
+
+    let deletion_task = tokio::task::spawn(
+        session_store
+            .clone()
+            .continuously_delete_expired(tokio::time::Duration::from_secs(60)),
+    );
+
     let session_layer = SessionManagerLayer::new(session_store)
         .with_secure(false)
         .with_expiry(Expiry::OnInactivity(Duration::seconds(10)));
@@ -41,8 +50,38 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let listener = tokio::net::TcpListener::bind(&addr).await?;
-    axum::serve(listener, app.into_make_service()).await?;
+
+    // Ensure we use a shutdown signal to abort the deletion task.
+    axum::serve(listener, app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
+
+    deletion_task.await??;
 
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}
 ```
diff --git a/mongodb-store/examples/mongodb.rs b/mongodb-store/examples/mongodb.rs
@@ -3,7 +3,9 @@ use std::net::SocketAddr;
 use axum::{response::IntoResponse, routing::get, Router};
 use serde::{Deserialize, Serialize};
 use time::Duration;
+use tokio::{signal, task::AbortHandle};
 use tower_sessions::{Expiry, Session, SessionManagerLayer};
+use tower_sessions_core::ExpiredDeletion;
 use tower_sessions_mongodb_store::{mongodb::Client, MongoDBStore};
 
 const COUNTER_KEY: &str = "counter";
@@ -22,6 +24,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let database_url = std::option_env!("DATABASE_URL").expect("Missing DATABASE_URL.");
     let client = Client::with_uri_str(database_url).await?;
     let session_store = MongoDBStore::new(client, "tower-sessions".to_string());
+
+    let deletion_task = tokio::task::spawn(
+        session_store
+            .clone()
+            .continuously_delete_expired(tokio::time::Duration::from_secs(60)),
+    );
+
     let session_layer = SessionManagerLayer::new(session_store)
         .with_secure(false)
         .with_expiry(Expiry::OnInactivity(Duration::seconds(10)));
@@ -30,7 +39,37 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let listener = tokio::net::TcpListener::bind(&addr).await?;
-    axum::serve(listener, app.into_make_service()).await?;
+
+    // Ensure we use a shutdown signal to abort the deletion task.
+    axum::serve(listener, app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
+
+    deletion_task.await??;
 
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}
diff --git a/sqlx-store/README.md b/sqlx-store/README.md
@@ -16,6 +16,7 @@ use std::net::SocketAddr;
 use axum::{response::IntoResponse, routing::get, Router};
 use serde::{Deserialize, Serialize};
 use time::Duration;
+use tokio::{signal, task::AbortHandle};
 use tower_sessions::{session_store::ExpiredDeletion, Expiry, Session, SessionManagerLayer};
 use tower_sessions_sqlx_store::{sqlx::MySqlPool, MySqlStore};
 
@@ -51,12 +52,40 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let listener = tokio::net::TcpListener::bind(&addr).await?;
-    axum::serve(listener, app.into_make_service()).await?;
+
+    // Ensure we use a shutdown signal to abort the deletion task.
+    axum::serve(listener, app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
 
     deletion_task.await??;
 
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}
 ```
 
 ### Postgres example
@@ -67,6 +96,7 @@ use std::net::SocketAddr;
 use axum::{response::IntoResponse, routing::get, Router};
 use serde::{Deserialize, Serialize};
 use time::Duration;
+use tokio::{signal, task::AbortHandle};
 use tower_sessions::{session_store::ExpiredDeletion, Expiry, Session, SessionManagerLayer};
 use tower_sessions_sqlx_store::{sqlx::PgPool, PostgresStore};
 
@@ -102,12 +132,40 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let listener = tokio::net::TcpListener::bind(&addr).await?;
-    axum::serve(listener, app.into_make_service()).await?;
+
+    // Ensure we use a shutdown signal to abort the deletion task.
+    axum::serve(listener, app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
 
     deletion_task.await??;
 
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}
 ```
 
 ### SQLite example
@@ -118,6 +176,7 @@ use std::net::SocketAddr;
 use axum::{response::IntoResponse, routing::get, Router};
 use serde::{Deserialize, Serialize};
 use time::Duration;
+use tokio::{signal, task::AbortHandle};
 use tower_sessions::{session_store::ExpiredDeletion, Expiry, Session, SessionManagerLayer};
 use tower_sessions_sqlx_store::{sqlx::SqlitePool, SqliteStore};
 
@@ -152,10 +211,38 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let listener = tokio::net::TcpListener::bind(&addr).await?;
-    axum::serve(listener, app.into_make_service()).await?;
+
+    // Ensure we use a shutdown signal to abort the deletion task.
+    axum::serve(listener, app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
 
     deletion_task.await??;
 
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}
 ```
diff --git a/sqlx-store/examples/mysql.rs b/sqlx-store/examples/mysql.rs
@@ -3,6 +3,7 @@ use std::net::SocketAddr;
 use axum::{response::IntoResponse, routing::get, Router};
 use serde::{Deserialize, Serialize};
 use time::Duration;
+use tokio::{signal, task::AbortHandle};
 use tower_sessions::{session_store::ExpiredDeletion, Expiry, Session, SessionManagerLayer};
 use tower_sessions_sqlx_store::{sqlx::MySqlPool, MySqlStore};
 
@@ -38,9 +39,37 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let listener = tokio::net::TcpListener::bind(&addr).await?;
-    axum::serve(listener, app.into_make_service()).await?;
+
+    // Ensure we use a shutdown signal to abort the deletion task.
+    axum::serve(listener, app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
 
     deletion_task.await??;
 
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}
diff --git a/sqlx-store/examples/postgres.rs b/sqlx-store/examples/postgres.rs
@@ -3,6 +3,7 @@ use std::net::SocketAddr;
 use axum::{response::IntoResponse, routing::get, Router};
 use serde::{Deserialize, Serialize};
 use time::Duration;
+use tokio::{signal, task::AbortHandle};
 use tower_sessions::{session_store::ExpiredDeletion, Expiry, Session, SessionManagerLayer};
 use tower_sessions_sqlx_store::{sqlx::PgPool, PostgresStore};
 
@@ -38,9 +39,37 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let listener = tokio::net::TcpListener::bind(&addr).await?;
-    axum::serve(listener, app.into_make_service()).await?;
+
+    // Ensure we use a shutdown signal to abort the deletion task.
+    axum::serve(listener, app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
 
     deletion_task.await??;
 
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}