diff --git a/slingshot/src/error.rs b/slingshot/src/error.rs index 1412142..17eb0ee 100644 --- a/slingshot/src/error.rs +++ b/slingshot/src/error.rs @@ -54,6 +54,8 @@ pub enum MainTaskError { ServerTaskError(#[from] ServerError), #[error(transparent)] IdentityTaskError(#[from] IdentityError), + #[error("firehose cache failed to close: {0}")] + FirehoseCacheCloseError(foyer::Error), } #[derive(Debug, Error)] diff --git a/slingshot/src/firehose_cache.rs b/slingshot/src/firehose_cache.rs index 0163dca..1ff33da 100644 --- a/slingshot/src/firehose_cache.rs +++ b/slingshot/src/firehose_cache.rs @@ -10,7 +10,11 @@ pub async fn firehose_cache( .memory(64 * 2_usize.pow(20)) .with_weighter(|k: &String, v| k.len() + std::mem::size_of_val(v)) .storage(Engine::large()) - .with_device_options(DirectFsDeviceOptions::new(cache_dir)) + .with_device_options( + DirectFsDeviceOptions::new(cache_dir) + .with_capacity(2_usize.pow(30)) // TODO: configurable (1GB to have something) + .with_file_size(16 * 2_usize.pow(20)), // note: this does limit the max cached item size, warning jumbo records + ) .build() .await .map_err(|e| format!("foyer setup error: {e:?}"))?; diff --git a/slingshot/src/identity.rs b/slingshot/src/identity.rs index 4550ce7..a55821c 100644 --- a/slingshot/src/identity.rs +++ b/slingshot/src/identity.rs @@ -13,6 +13,7 @@ use std::sync::Arc; /// 3. DID -> handle resolution: for bidirectional handle validation and in case we want to offer this use std::time::Duration; use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; use crate::error::IdentityError; use atrium_api::{ @@ -175,8 +176,12 @@ impl Identity { .with_name("identity") .memory(16 * 2_usize.pow(20)) .with_weighter(|k, v| std::mem::size_of_val(k) + std::mem::size_of_val(v)) - .storage(Engine::large()) - .with_device_options(DirectFsDeviceOptions::new(cache_dir)) + .storage(Engine::small()) + .with_device_options( + DirectFsDeviceOptions::new(cache_dir) + .with_capacity(2_usize.pow(30)) // TODO: configurable (1GB to have something) + .with_file_size(2_usize.pow(20)), // note: this does limit the max cached item size, warning jumbo records + ) .build() .await?; @@ -403,12 +408,21 @@ impl Identity { } /// run the refresh queue consumer - pub async fn run_refresher(&self) -> Result<(), IdentityError> { + pub async fn run_refresher(&self, shutdown: CancellationToken) -> Result<(), IdentityError> { let _guard = self .refresher .try_lock() .expect("there to only be one refresher running"); loop { + if shutdown.is_cancelled() { + log::info!("identity refresher: exiting for shutdown: closing cache..."); + if let Err(e) = self.cache.close().await { + log::error!("cache close errored: {e}"); + } else { + log::info!("identity cache closed.") + } + return Ok(()); + } let Some(task_key) = self.peek_refresh().await else { tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; continue; diff --git a/slingshot/src/main.rs b/slingshot/src/main.rs index 4cbd0a7..c427e2a 100644 --- a/slingshot/src/main.rs +++ b/slingshot/src/main.rs @@ -89,8 +89,9 @@ async fn main() -> Result<(), String> { .map_err(|e| format!("identity setup failed: {e:?}"))?; log::info!("identity service ready."); let identity_refresher = identity.clone(); + let identity_shutdown = shutdown.clone(); tasks.spawn(async move { - identity_refresher.run_refresher().await?; + identity_refresher.run_refresher(identity_shutdown).await?; Ok(()) }); @@ -113,13 +114,14 @@ async fn main() -> Result<(), String> { }); let consumer_shutdown = shutdown.clone(); + let consumer_cache = cache.clone(); tasks.spawn(async move { consume( args.jetstream, None, args.jetstream_no_zstd, consumer_shutdown, - cache, + consumer_cache, ) .await?; Ok(()) @@ -133,13 +135,20 @@ async fn main() -> Result<(), String> { } } + tasks.spawn(async move { + cache + .close() + .await + .map_err(MainTaskError::FirehoseCacheCloseError) + }); + tokio::select! { _ = async { while let Some(completed) = tasks.join_next().await { log::info!("shutdown: task completed: {completed:?}"); } } => {}, - _ = tokio::time::sleep(std::time::Duration::from_secs(3)) => { + _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => { log::info!("shutdown: not all tasks completed on time. aborting..."); tasks.shutdown().await; }, diff --git a/slingshot/src/server.rs b/slingshot/src/server.rs index 3c52bb6..26612d6 100644 --- a/slingshot/src/server.rs +++ b/slingshot/src/server.rs @@ -410,7 +410,7 @@ pub async fn serve( host: Option, acme_contact: Option, certs: Option, - _shutdown: CancellationToken, + shutdown: CancellationToken, ) -> Result<(), ServerError> { let repo = Arc::new(repo); let api_service = OpenApiService::new( @@ -452,13 +452,18 @@ pub async fn serve( } let auto_cert = auto_cert.build().map_err(ServerError::AcmeBuildError)?; - run(TcpListener::bind("0.0.0.0:443").acme(auto_cert), app).await + run( + TcpListener::bind("0.0.0.0:443").acme(auto_cert), + app, + shutdown, + ) + .await } else { - run(TcpListener::bind("127.0.0.1:3000"), app).await + run(TcpListener::bind("127.0.0.1:3000"), app, shutdown).await } } -async fn run(listener: L, app: Route) -> Result<(), ServerError> +async fn run(listener: L, app: Route, shutdown: CancellationToken) -> Result<(), ServerError> where L: Listener + 'static, { @@ -472,7 +477,8 @@ where .with(Tracing); Server::new(listener) .name("slingshot") - .run(app) + .run_with_graceful_shutdown(app, shutdown.cancelled(), None) .await .map_err(ServerError::ServerExited) + .inspect(|()| log::info!("server ended. goodbye.")) }