diff --git a/MIGRATION.md b/MIGRATION.md index b6f5560..d8164b9 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -1,5 +1,6 @@ # Migration Guide +- [Unreleased: owned scheduled tasks](#unreleased-owned-scheduled-tasks) - [Unreleased: shared startup](#unreleased-shared-startup) - [Unreleased: runtime type bounds](#unreleased-runtime-type-bounds) - [Unreleased: HTTP response contracts](#unreleased-http-response-contracts) @@ -14,6 +15,68 @@ --- +# Unreleased: owned scheduled tasks + +**Task registration no longer starts background work.** The application owns its scheduled tasks +and shuts them down together with HTTP serving. + +## Registration and ownership + +`TaskScheduler::cron` and `interval` now take `&mut self` and only register work. Existing application +hooks already receive a mutable scheduler. Update helper functions that accepted `&TaskScheduler` +to accept `&mut TaskScheduler`. + +All registrations must succeed before serving starts any task. A registration error, panic, or +cancelled startup discards the pending tasks. Do not wait inside the registration hook for a +scheduled task to execute: it cannot start until the hook returns. + +Standalone callers must explicitly start the scheduler and retain the returned owner: + +```rust,ignore +let mut scheduler = TaskScheduler::new(context); +scheduler.interval("cleanup", Duration::from_secs(60), cleanup); +let running = scheduler.start(); +// Keep `running` alive for as long as the tasks should run. +running.shutdown(Duration::from_secs(10)).await; +``` + +`RunningTasks::shutdown(timeout)` stops subsequent executions, waits for current executions, and +aborts and joins unfinished tasks at the deadline. Dropping the owner or cancelling its shutdown +future requests immediate abortion; dropping cannot perform an asynchronous join. The same ownership +includes the current execution, which is no longer a separate detached task. + +## Application shutdown + +Builder `run` / `listen` / `listen_on` and `GotchaApp::run` now wait for Ctrl-C or Unix SIGTERM by +default, then use one cancellation signal to stop HTTP acceptance and new scheduled executions. +In-flight HTTP requests drain while scheduled tasks finish. Override the signal with a future: + +```rust,no_run +use gotcha::Gotcha; +let (stop, stopped) = tokio::sync::oneshot::channel::<()>(); +let app = Gotcha::new().shutdown_signal(async move { let _ = stopped.await; }); +// Resolve the future by sending through `stop`; `app.run().await` waits for shutdown to finish. +``` + +For `GotchaApp`, override `fn shutdown_signal(&self) -> impl Future + Send` instead. +That future may borrow the application. A custom signal replaces the default OS signal handling. + +Scheduled executions get a 30-second grace period by default. Configure it with the builder's +`.task_shutdown_timeout(duration)` or the trait's `fn task_shutdown_timeout(&self) -> Duration`. +The deadline starts when shutdown is requested, independently of HTTP draining. It does not limit +HTTP request duration. A serve error also stops and drains the scheduled tasks; cancelling the +serving future requests abortion of all scheduled work and signals HTTP shutdown. + +Cancellation requires async tasks to yield. Synchronous blocking work and tasks explicitly detached +by application code (including `Messager::spawn`) remain outside this ownership. Public preparation, +application cleanup hooks, and HTTP drain limits remain follow-up work in #82. + +Single-execution panics are logged and the schedule continues, including a panic while constructing +the execution future. Interval tasks retain their existing immediate first tick and non-overlapping +executions. Invalid cron expressions and zero intervals are logged and skipped during registration. + +--- + # Unreleased: shared startup **Builder and trait applications now share configuration policy, address resolution, binding, @@ -78,8 +141,8 @@ error releases the listener. The initialized context's server settings contain t address, including a dynamically assigned port, so state initialization and handlers agree. State initialization now runs while the port is reserved, before requests are accepted. Applications -that depended on initializing state before a bind attempt must account for the new order. Task -cancellation and waiting on shutdown are unchanged and tracked separately in #93. +that depended on initializing state before a bind attempt must account for the new order. See +[owned scheduled tasks](#unreleased-owned-scheduled-tasks) for task startup and shutdown behavior. --- diff --git a/README.md b/README.md index e3c316d..bf9fb3c 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,26 @@ Fallback applies only to startup configuration loading; eager `build_config()` a `GotchaApp::config()` calls still return errors. See the [migration guide](MIGRATION.md#unreleased-shared-startup) for behavior changes and the trait example. +### Shutdown + +Both application APIs stop gracefully on Ctrl-C or Unix SIGTERM. HTTP stops accepting connections, +existing requests drain, and scheduled tasks stop starting new executions. To provide your own +shutdown signal: + +```rust,no_run +use gotcha::Gotcha; + +let (stop, stopped) = tokio::sync::oneshot::channel::<()>(); +let app = Gotcha::new() + .shutdown_signal(async move { let _ = stopped.await; }); +// Send through `stop` to request shutdown; await `app.run()` to wait for it to finish. +``` + +The trait API offers the same `shutdown_signal(&self)` hook. With the `task` feature, in-flight +scheduled executions get 30 seconds to finish before being aborted and joined. Override that period +with builder `.task_shutdown_timeout(duration)` or trait `task_shutdown_timeout(&self)`. This timeout +applies to scheduled tasks; HTTP requests drain independently. See the [migration guide](MIGRATION.md#unreleased-owned-scheduled-tasks). + ### Advanced Trait API (For complex applications) ```rust,no_run @@ -397,7 +417,10 @@ Configuration supports: ### Task Scheduling -Requires the `task` feature. +Requires the `task` feature. Registration collects work; execution starts after all application +initialization succeeds. The application retains ownership until shutdown completes. Standalone +users call `scheduler.start()`, retain the returned `RunningTasks`, and call +`running.shutdown(timeout).await` to stop and join their tasks. Dropping the owner aborts its tasks. ```rust,ignore use gotcha::prelude::*; diff --git a/gotcha/Cargo.toml b/gotcha/Cargo.toml index 02b259d..597b955 100644 --- a/gotcha/Cargo.toml +++ b/gotcha/Cargo.toml @@ -30,7 +30,7 @@ openapi = ["dep:oas", "dep:gotcha_core", "gotcha_core/axum"] # dozen more, so a CORS-only application should not pay for the static file machinery. cors = ["dep:tower-http", "tower-http/cors"] static_files = ["dep:tower-http", "tower-http/fs"] -task = ["dep:cron", "tokio/time"] +task = ["dep:cron", "dep:futures-util", "tokio/time"] [dependencies] @@ -38,9 +38,10 @@ async-trait = "0.1.60" gotcha_macro = { version = "0.4.1", path = "../gotcha_macro" } gotcha_core = { version = "0.4.1", path = "../gotcha_core", optional = true } serde = {version = "1", features = ["derive"]} -# `rt` is needed by the message system, which is always available (it gated nothing but a -# 160-line module and this one tokio feature). -tokio = {version = "1", features = ["rt"]} +# `rt` also supports the always-available message system; macros and signals drive server shutdown. +tokio = {version = "1", features = ["rt", "macros", "signal"]} +tokio-util = { version = "0.7", features = ["rt"] } +futures-util = { version = "0.3", optional = true, default-features = false, features = ["std", "async-await"] } tracing = "0.1" tracing-subscriber = {version="0.3", features=["env-filter"]} log = "0.4" @@ -78,8 +79,8 @@ assert-json-diff = "2" # Non-optional so trybuild-generated test crates (which derive Schematic and thus # emit `::gotcha_core::` paths) have gotcha_core in their extern prelude. gotcha_core = { version = "0.4.1", path = "../gotcha_core" } -# So `#[tokio::main]` doc examples compile (the lib's own tokio dep is minimal). -tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +# Multiple-worker runtimes, paused-clock task tests, and socket-based lifecycle checks. +tokio = { version = "1", features = ["macros", "rt-multi-thread", "test-util", "io-util"] } futures-core = "0.3.33" futures-util = "0.3.33" tower = "0.5.3" diff --git a/gotcha/src/builder.rs b/gotcha/src/builder.rs index 80e816f..9025030 100644 --- a/gotcha/src/builder.rs +++ b/gotcha/src/builder.rs @@ -117,10 +117,13 @@ where router: GotchaRouter>, listen_options: ListenOptions, config_error_policy: ConfigErrorPolicy, + shutdown_signal: Option + Send>>>, state: StateSource, config: Configuration, #[cfg(feature = "task")] tasks: Vec>, + #[cfg(feature = "task")] + task_shutdown_timeout: std::time::Duration, } impl Default for Gotcha { @@ -267,10 +270,13 @@ where router: GotchaRouter::default(), listen_options: ListenOptions::default(), config_error_policy: ConfigErrorPolicy::Strict, + shutdown_signal: None, state, config, #[cfg(feature = "task")] tasks: Vec::new(), + #[cfg(feature = "task")] + task_shutdown_timeout: crate::task::DEFAULT_SHUTDOWN_TIMEOUT, } } @@ -715,7 +721,8 @@ where /// /// The closure receives a [`TaskScheduler`](crate::TaskScheduler) when the /// server starts, on which you can register `cron` / `interval` jobs. This - /// brings the builder to parity with `GotchaApp::tasks`. + /// brings the builder to parity with `GotchaApp::tasks`. Jobs start only after all + /// registrations succeed and are owned by the application until shutdown. /// Register module tasks here on the top-level application; `nest` and `merge` accept /// only route modules and cannot accept another application's task registrations. /// @@ -739,6 +746,24 @@ where self } + /// Replace the default Ctrl-C / Unix SIGTERM signal. When this future resolves, HTTP + /// stops accepting connections and scheduled tasks stop starting new executions. + pub fn shutdown_signal(mut self, signal: F) -> Self + where + F: std::future::Future + Send + 'static, + { + self.shutdown_signal = Some(Box::pin(signal)); + self + } + + /// Set the grace period for in-flight scheduled executions (default: 30 seconds). + /// Remaining tasks are aborted and joined at the deadline. HTTP draining has no time limit here. + #[cfg(feature = "task")] + pub fn task_shutdown_timeout(mut self, timeout: std::time::Duration) -> Self { + self.task_shutdown_timeout = timeout; + self + } + /// Add CORS support (requires "cors" feature) #[cfg(feature = "cors")] pub fn with_cors(self) -> Self { @@ -806,6 +831,18 @@ where std::mem::take(&mut self.config_error_policy) } + async fn shutdown_signal(&mut self) { + match self.shutdown_signal.take() { + Some(signal) => signal.await, + None => startup::shutdown_signal().await, + } + } + + #[cfg(feature = "task")] + fn task_shutdown_timeout(&self) -> std::time::Duration { + self.task_shutdown_timeout + } + async fn config(&mut self) -> GotchaResult> { self.config.resolve() } @@ -865,8 +902,8 @@ mod tests { panic!("state was already initialized"); } } - let app = Gotcha::with_state::().state(Initialized).config(ConfigWrapper::default()); - startup::prepare(app, Some("127.0.0.1:0".parse().unwrap())).await.unwrap(); + let mut app = Gotcha::with_state::().state(Initialized).config(ConfigWrapper::default()); + startup::prepare(&mut app, Some("127.0.0.1:0".parse().unwrap())).await.unwrap(); } #[tokio::test] @@ -892,7 +929,7 @@ mod tests { } let calls = Arc::new(AtomicUsize::new(41)); - let app = Gotcha::from_context(GotchaContext { + let mut app = Gotcha::from_context(GotchaContext { state: RuntimeState(calls.clone()), config: ConfigWrapper { app: RuntimeConfig { value: "parent".into() }, @@ -902,7 +939,7 @@ mod tests { .with_file_config("unused-required-file.toml") .nest("/api", GotchaRouter::default().get("/context", read)) .merge(GotchaRouter::default().get("/context", read)); - let prepared = startup::prepare(app, Some("127.0.0.1:0".parse().unwrap())).await.unwrap(); + let prepared = startup::prepare(&mut app, Some("127.0.0.1:0".parse().unwrap())).await.unwrap(); let router = prepared.router; for (path, expected) in [("/api/context", "parent:41"), ("/context", "parent:42")] { diff --git a/gotcha/src/lib.rs b/gotcha/src/lib.rs index 070d75e..f566acc 100644 --- a/gotcha/src/lib.rs +++ b/gotcha/src/lib.rs @@ -110,7 +110,7 @@ pub use crate::openapi::schematic::EnhancedSchema; pub use serde_json; #[cfg(feature = "task")] #[cfg_attr(docsrs, doc(cfg(feature = "task")))] -pub use task::TaskScheduler; +pub use task::{RunningTasks, TaskScheduler}; #[cfg(feature = "static_files")] #[cfg_attr(docsrs, doc(cfg(feature = "static_files")))] pub use tower_http::services::{ServeDir, ServeFile}; @@ -184,6 +184,21 @@ pub trait GotchaApp: Sized + Send + Sync { ConfigErrorPolicy::Strict } + /// Resolve to stop accepting HTTP connections and scheduling new task executions. + /// The default waits for Ctrl-C or, on Unix, SIGTERM. Existing HTTP requests drain; + /// scheduled executions wait for at most `task_shutdown_timeout()` (with the `task` feature). + fn shutdown_signal(&self) -> impl std::future::Future + Send { + crate::startup::shutdown_signal() + } + + /// How long to wait for in-flight scheduled executions after shutdown is requested. + /// Defaults to 30 seconds. At the deadline, remaining scheduled tasks are aborted and joined. + /// This does not limit HTTP request draining. + #[cfg(feature = "task")] + fn task_shutdown_timeout(&self) -> std::time::Duration { + crate::task::DEFAULT_SHUTDOWN_TIMEOUT + } + /// Install the tracing subscriber. The default reads `RUST_LOG`. fn logger(&self) -> GotchaResult<()> { tracing_subscriber::registry() @@ -206,7 +221,8 @@ pub trait GotchaApp: Sized + Send + Sync { fn state(&self, config: &ConfigWrapper) -> impl std::future::Future> + Send; #[cfg(feature = "task")] - /// Register background tasks. The default registers none. + /// Register background tasks without starting them. Returning an error discards all registrations. + /// Tasks start with serving after successful initialization. The default registers none. fn tasks(&self, _task_scheduler: &mut TaskScheduler) -> impl std::future::Future> + Send { async { Ok(()) } } diff --git a/gotcha/src/prelude.rs b/gotcha/src/prelude.rs index c2f29c8..e13d5d9 100644 --- a/gotcha/src/prelude.rs +++ b/gotcha/src/prelude.rs @@ -62,7 +62,7 @@ pub use crate::{api, Responsible, Schematic}; pub use crate::layers::CorsLayer; #[cfg(feature = "task")] -pub use crate::TaskScheduler; +pub use crate::{RunningTasks, TaskScheduler}; // Utility macros for common patterns #[macro_export] diff --git a/gotcha/src/startup.rs b/gotcha/src/startup.rs index 00b0e84..43cf0db 100644 --- a/gotcha/src/startup.rs +++ b/gotcha/src/startup.rs @@ -1,10 +1,11 @@ //! Shared application startup. API adapters provide the application-specific hooks. -use std::future::Future; +use std::future::{Future, IntoFuture}; use std::net::{IpAddr, SocketAddr}; use axum::Router; use tokio::net::TcpListener; +use tokio_util::sync::CancellationToken; use crate::{ConfigErrorPolicy, ConfigWrapper, GotchaApp, GotchaContext, GotchaError, GotchaResult, ServerConfig}; @@ -36,6 +37,9 @@ pub(crate) trait Application: Send { ListenOptions::default() } fn config_error_policy(&mut self) -> ConfigErrorPolicy; + fn shutdown_signal(&mut self) -> impl Future + Send; + #[cfg(feature = "task")] + fn task_shutdown_timeout(&self) -> std::time::Duration; fn config(&mut self) -> impl Future>> + Send; fn state(&mut self, config: &ConfigWrapper) -> impl Future> + Send; fn router(&mut self, context: GotchaContext) -> impl Future> + Send; @@ -46,9 +50,14 @@ pub(crate) trait Application: Send { pub(crate) struct PreparedServer { listener: TcpListener, pub(crate) router: Router, + shutdown: CancellationToken, + #[cfg(feature = "task")] + tasks: Vec, + #[cfg(feature = "task")] + task_shutdown_timeout: std::time::Duration, } -pub(crate) async fn prepare(mut app: A, explicit: Option) -> GotchaResult { +pub(crate) async fn prepare(app: &mut A, explicit: Option) -> GotchaResult { app.logger()?; let policy = app.config_error_policy(); let mut config = policy.apply(app.config().await)?; @@ -66,19 +75,82 @@ pub(crate) async fn prepare(mut app: A, explicit: Option(app: A, explicit: Option) -> GotchaResult<()> { - prepare(app, explicit).await?.serve().await +pub(crate) async fn run(mut app: A, explicit: Option) -> GotchaResult<()> { + let prepared = prepare(&mut app, explicit).await?; + prepared.serve(app.shutdown_signal()).await } impl PreparedServer { - async fn serve(self) -> GotchaResult<()> { + async fn serve(self, signal: impl Future + Send) -> GotchaResult<()> { tracing::info!("Server listening on http://{}", self.listener.local_addr().map_err(GotchaError::Io)?); - axum::serve(self.listener, self.router).await.map_err(GotchaError::Io) + // If the serving future is dropped, stop HTTP acceptance and abort owned scheduled tasks. + let _shutdown_on_drop = self.shutdown.clone().drop_guard(); + #[cfg(feature = "task")] + let tasks = crate::RunningTasks::start(self.tasks, self.shutdown.clone()); + let server = axum::serve(self.listener, self.router) + .with_graceful_shutdown(self.shutdown.clone().cancelled_owned()) + .into_future(); + tokio::pin!(server); + let serving = async { + let result = tokio::select! { + biased; + _ = signal => { + self.shutdown.cancel(); + server.await + }, + result = &mut server => result, + }; + self.shutdown.cancel(); + result.map_err(GotchaError::Io) + }; + #[cfg(feature = "task")] + { + let drain_tasks = async { + self.shutdown.cancelled().await; + tasks.shutdown(self.task_shutdown_timeout).await; + }; + // Start the task deadline when shutdown is requested, independently of HTTP draining. + let (result, ()) = tokio::join!(serving, drain_tasks); + result + } + #[cfg(not(feature = "task"))] + serving.await + } +} + +pub(crate) async fn shutdown_signal() { + #[cfg(unix)] + let result = { + match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) { + Ok(mut terminate) => tokio::select! { + result = tokio::signal::ctrl_c() => result, + _ = terminate.recv() => Ok(()), + }, + Err(error) => Err(error), + } + }; + #[cfg(not(unix))] + let result = tokio::signal::ctrl_c().await; + if let Err(error) = result { + tracing::error!("failed to install shutdown signal handler: {error}; shutting down"); } } @@ -101,6 +173,15 @@ where A::config_error_policy(self) } + async fn shutdown_signal(&mut self) { + A::shutdown_signal(self).await; + } + + #[cfg(feature = "task")] + fn task_shutdown_timeout(&self) -> std::time::Duration { + A::task_shutdown_timeout(self) + } + async fn config(&mut self) -> GotchaResult> { A::config(self).await } diff --git a/gotcha/src/startup/tests.rs b/gotcha/src/startup/tests.rs index 58adf89..32486c9 100644 --- a/gotcha/src/startup/tests.rs +++ b/gotcha/src/startup/tests.rs @@ -134,7 +134,7 @@ impl GotchaApp for TraitApp { async fn assert_http(prepared: PreparedServer, label: &str) { let address = prepared.listener.local_addr().unwrap(); - let server = tokio::spawn(prepared.serve()); + let server = tokio::spawn(prepared.serve(std::future::pending())); let response = tokio::task::spawn_blocking(move || { use std::io::{Read, Write}; let mut stream = std::net::TcpStream::connect(address).unwrap(); @@ -158,10 +158,10 @@ async fn both_apis_serve_the_loaded_config_and_actual_bound_address() { let path = dir.path().join("config.toml"); std::fs::write(&path, "label = 'loaded'\n[server]\nhost = '127.0.0.1'\nport = 0").unwrap(); let app = TraitApp::new(&path); - let builder = Gotcha::with_config::().with_file_config(&path).get("/config", read); - let prepared = prepare(builder, None).await.unwrap(); + let mut builder = Gotcha::with_config::().with_file_config(&path).get("/config", read); + let prepared = prepare(&mut builder, None).await.unwrap(); assert_http(prepared, "loaded").await; - let prepared = prepare(&app, None).await.unwrap(); + let prepared = prepare(&mut &app, None).await.unwrap(); assert_eq!(*app.bound.lock().unwrap(), Some(prepared.listener.local_addr().unwrap())); assert_http(prepared, "loaded").await; let mut expected = vec!["logger", "config", "state", "router"]; @@ -190,8 +190,8 @@ async fn run_uses_configured_address_and_does_not_initialize_after_bind_failure( } assert_eq!(*events.lock().unwrap(), ["logger", "config"]); // An explicit builder override can use the same configuration on an available port. - let builder = Gotcha::with_config::().with_file_config(&path).port(0).get("/config", read); - assert_http(prepare(builder, None).await.unwrap(), "occupied").await; + let mut builder = Gotcha::with_config::().with_file_config(&path).port(0).get("/config", read); + assert_http(prepare(&mut builder, None).await.unwrap(), "occupied").await; } #[tokio::test] @@ -201,18 +201,18 @@ async fn configuration_fallback_is_explicit_and_shared_by_both_apis() { std::fs::write(&path, "label = [").unwrap(); let mut app = TraitApp::new(&path); let builder = || Gotcha::with_config::().with_file_config(&path).get("/config", read); - assert!(matches!(prepare(builder(), None).await, Err(GotchaError::Config(_)))); - assert!(matches!(prepare(&app, None).await, Err(GotchaError::Config(_)))); + assert!(matches!(prepare(&mut builder(), None).await, Err(GotchaError::Config(_)))); + assert!(matches!(prepare(&mut &app, None).await, Err(GotchaError::Config(_)))); assert_eq!(*app.events.lock().unwrap(), ["logger", "config"]); app.fallback = true; assert_http( - prepare(builder().config_error_policy(ConfigErrorPolicy::Fallback(fallback_settings)), None) + prepare(&mut builder().config_error_policy(ConfigErrorPolicy::Fallback(fallback_settings)), None) .await .unwrap(), "fallback", ) .await; - assert_http(prepare(&app, None).await.unwrap(), "fallback").await; + assert_http(prepare(&mut &app, None).await.unwrap(), "fallback").await; } #[tokio::test] @@ -228,7 +228,7 @@ async fn initialization_errors_release_the_listener_and_stop_later_stages() { let mut app = TraitApp::new(&path); app.fail_at = Some(stage); app.fallback = true; // A config fallback must not swallow errors from other stages. - assert!(matches!(prepare(&app, None).await, Err(GotchaError::Message(message)) if message == stage)); + assert!(matches!(prepare(&mut &app, None).await, Err(GotchaError::Message(message)) if message == stage)); let bound = app.bound.lock().unwrap().unwrap(); let _rebound = TcpListener::bind(bound).await.expect("failed startup must release its listener"); assert_eq!(app.events.lock().unwrap().last(), Some(&stage)); diff --git a/gotcha/src/task.rs b/gotcha/src/task.rs index 80c839e..e9e0646 100644 --- a/gotcha/src/task.rs +++ b/gotcha/src/task.rs @@ -7,7 +7,7 @@ //! //! - Cron expression based scheduling //! - Fixed interval scheduling -//! - Async task execution +//! - Async task execution with application-owned shutdown //! - Access to application context in tasks //! //! ## Examples @@ -18,9 +18,9 @@ //! //! # #[derive(Clone)] //! # struct Config {} -//! fn setup(scheduler: &TaskScheduler<(), Config>) { +//! fn setup(scheduler: &mut TaskScheduler<(), Config>) { //! // Schedule a cron task (the expression is a `String`) -//! scheduler.cron("daily-cleanup", "0 0 * * *".to_string(), |_ctx| async move { +//! scheduler.cron("daily-cleanup", "0 0 0 * * *".to_string(), |_ctx| async move { //! // Task implementation //! }); //! @@ -40,18 +40,28 @@ //! use std::future::Future; +use std::panic::AssertUnwindSafe; +use std::pin::Pin; use std::str::FromStr; +use std::time::Duration; use chrono::Utc; use cron::Schedule; -use tracing::info; +use futures_util::FutureExt; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; use crate::GotchaContext; -/// Registers background tasks that run alongside the server, each with access to the application -/// context. +pub(crate) type ScheduledTask = Pin + Send>>; +pub(crate) const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30); + +/// Registers tasks without starting them. Applications start all registrations together after +/// successful initialization. Standalone callers must retain the owner returned by [`Self::start`]. pub struct TaskScheduler { context: GotchaContext, + shutdown: CancellationToken, + tasks: Vec, } impl TaskScheduler @@ -59,16 +69,21 @@ where T1: Clone + Send + Sync + 'static, T2: Clone + Send + Sync + 'static, { - /// Create a scheduler bound to the application context tasks will receive. + /// Create a registration-only scheduler. This does not require a running Tokio runtime. pub fn new(context: GotchaContext) -> Self { - Self { context } + Self::with_shutdown(context, CancellationToken::new()) } - /// Schedule a task on a cron expression. - /// - /// An invalid expression is reported (via `tracing::error!`) and the task is - /// simply not started, instead of panicking a detached task at the first tick. - pub fn cron(&self, name: impl AsRef, expression: String, task: F) + pub(crate) fn with_shutdown(context: GotchaContext, shutdown: CancellationToken) -> Self { + Self { + context, + shutdown, + tasks: Vec::new(), + } + } + + /// Register a cron task. Invalid expressions are logged and not registered. + pub fn cron(&mut self, name: impl AsRef, expression: String, task: F) where F: Fn(GotchaContext) -> FF + Send + 'static, FF: Future + Send + 'static, @@ -77,66 +92,172 @@ where let schedule = match Schedule::from_str(&expression) { Ok(schedule) => schedule, Err(e) => { - tracing::error!("cron task {name:?} has an invalid schedule {expression:?}: {e}; task not started"); + tracing::error!("cron task {name:?} has an invalid schedule {expression:?}: {e}; task not registered"); return; } }; - info!("starting cron task: {name}"); - tokio::spawn(cron_proc_macro_wrapper(self.context.clone(), schedule, name, task)); + self.tasks + .push(Box::pin(run_cron(self.context.clone(), schedule, name, task, self.shutdown.clone()))); } - /// Run `task` every `interval`, starting one interval from now. - pub fn interval(&self, name: impl AsRef, interval: std::time::Duration, task: F) + /// Register a task with an immediate first execution, then repeat at the given interval. + /// A zero interval is logged and not registered. Executions of the same task never overlap. + pub fn interval(&mut self, name: impl AsRef, interval: Duration, task: F) where F: Fn(GotchaContext) -> FF + Send + 'static, FF: Future + Send + 'static, { let name = name.as_ref().to_string(); - info!("starting interval task: {name}"); - tokio::spawn(interval_proc_macro_wrapper(self.context.clone(), interval, name, task)); + if interval.is_zero() { + tracing::error!("interval task {name:?} has a zero interval; task not registered"); + return; + } + self.tasks + .push(Box::pin(run_interval(self.context.clone(), interval, name, task, self.shutdown.clone()))); + } + + /// Start the registered tasks on Tokio. Dropping the returned owner aborts them. + /// + /// ```no_run + /// use gotcha::{ConfigWrapper, GotchaContext, TaskScheduler}; + /// use std::time::Duration; + /// # async fn example() { + /// let mut scheduler = TaskScheduler::new(GotchaContext { state: (), config: ConfigWrapper::<()>::default() }); + /// scheduler.interval("cleanup", Duration::from_secs(60), |_| async {}); + /// let running = scheduler.start(); + /// running.shutdown(Duration::from_secs(10)).await; + /// # } + /// ``` + /// + /// Use [`RunningTasks::shutdown`] to stop scheduling and wait for in-flight executions. + pub fn start(self) -> RunningTasks { + RunningTasks::start(self.tasks, self.shutdown) + } + + pub(crate) fn into_tasks(self) -> Vec { + self.tasks } } -/// Run one task execution under supervision: if it panics, log it and keep the -/// scheduler loop alive instead of silently killing the whole task. -async fn run_supervised(name: &str, fut: FF) +/// Owns all running scheduled tasks, including their current executions. +/// Dropping this value requests immediate abortion; [`Self::shutdown`] also waits for cleanup. +/// Tasks spawned independently by application code are outside this owner. +#[must_use = "dropping the task owner aborts all scheduled tasks"] +pub struct RunningTasks { + tasks: JoinSet<()>, + shutdown: CancellationToken, +} + +impl RunningTasks { + pub(crate) fn start(pending: Vec, shutdown: CancellationToken) -> Self { + let mut tasks = JoinSet::new(); + for task in pending { + tasks.spawn(task); + } + Self { tasks, shutdown } + } + + /// Stop new executions and wait up to `timeout` for current executions to complete. + /// At the deadline, abort unfinished tasks and wait for their futures to be dropped. + /// Cancelling this future also aborts the tasks. Cancellation requires tasks to yield; + /// blocking work and application-created detached tasks cannot be forcibly stopped here. + pub async fn shutdown(mut self, timeout: Duration) { + self.shutdown.cancel(); + let drain = async { + while let Some(result) = self.tasks.join_next().await { + if let Err(error) = result { + tracing::error!("scheduled task loop failed: {error}"); + } + } + }; + if tokio::time::timeout(timeout, drain).await.is_err() { + tracing::warn!("scheduled tasks exceeded shutdown timeout {timeout:?}; aborting remaining tasks"); + self.tasks.shutdown().await; + } + } +} + +// Run the callback and its future in the owned loop. Catch both construction and polling panics +// without spawning an execution that could outlive the loop. A mutable borrow retains F: Send +// without requiring F: Sync. +async fn run_supervised(name: &str, context: GotchaContext, task: &mut F) where - FF: Future + Send + 'static, + T1: Clone + Send + Sync + 'static, + T2: Clone + Send + Sync + 'static, + F: Fn(GotchaContext) -> FF + Send, + FF: Future + Send, { - if let Err(join_error) = tokio::spawn(fut).await { - tracing::error!("scheduled task {name:?} panicked: {join_error}"); + if let Err(panic) = AssertUnwindSafe(async move { task(context).await }).catch_unwind().await { + let message = panic + .downcast_ref::<&str>() + .copied() + .or_else(|| panic.downcast_ref::().map(String::as_str)) + .unwrap_or("non-string panic payload"); + tracing::error!("scheduled task {name:?} panicked: {message}"); } } -/// Drives a cron task; called by the scheduler, not directly. -pub async fn cron_proc_macro_wrapper(context: GotchaContext, schedule: Schedule, name: String, task: F) +async fn run_cron(context: GotchaContext, schedule: Schedule, name: String, mut task: F, shutdown: CancellationToken) where T1: Clone + Send + Sync + 'static, T2: Clone + Send + Sync + 'static, F: Fn(GotchaContext) -> FF + Send + 'static, FF: Future + Send + 'static, { + tracing::info!("starting cron task: {name}"); for next_trigger_time in schedule.upcoming(Utc) { - let now = Utc::now(); - // A trigger computed in the past (clock skew, or a long previous run) would make - // `to_std()` fail — run immediately in that case rather than panicking. - let wait = (next_trigger_time - now).to_std().unwrap_or(std::time::Duration::ZERO); - tokio::time::sleep(wait).await; - run_supervised(&name, task(context.clone())).await; + let wait = (next_trigger_time - Utc::now()).to_std().unwrap_or(Duration::ZERO); + tokio::select! { + biased; + _ = shutdown.cancelled() => break, + _ = tokio::time::sleep(wait) => {} + } + run_supervised(&name, context.clone(), &mut task).await; } } -/// Drives an interval task; called by the scheduler, not directly. -pub async fn interval_proc_macro_wrapper(context: GotchaContext, interval: std::time::Duration, name: String, task: F) +async fn run_interval(context: GotchaContext, interval: Duration, name: String, mut task: F, shutdown: CancellationToken) where T1: Clone + Send + Sync + 'static, T2: Clone + Send + Sync + 'static, F: Fn(GotchaContext) -> FF + Send + 'static, FF: Future + Send + 'static, { + tracing::info!("starting interval task: {name}"); let mut interval = tokio::time::interval(interval); loop { - interval.tick().await; - run_supervised(&name, task(context.clone())).await; + tokio::select! { + biased; + _ = shutdown.cancelled() => break, + _ = interval.tick() => {} + } + run_supervised(&name, context.clone(), &mut task).await; } } + +/// Low-level cron driver retained for compatibility. The caller owns this future; +/// prefer [`TaskScheduler`] for managed shutdown and waiting. +pub async fn cron_proc_macro_wrapper(context: GotchaContext, schedule: Schedule, name: String, task: F) +where + T1: Clone + Send + Sync + 'static, + T2: Clone + Send + Sync + 'static, + F: Fn(GotchaContext) -> FF + Send + 'static, + FF: Future + Send + 'static, +{ + run_cron(context, schedule, name, task, CancellationToken::new()).await; +} + +/// Low-level interval driver retained for compatibility. The caller owns this future; +/// prefer [`TaskScheduler`] for managed shutdown and waiting. +pub async fn interval_proc_macro_wrapper(context: GotchaContext, interval: Duration, name: String, task: F) +where + T1: Clone + Send + Sync + 'static, + T2: Clone + Send + Sync + 'static, + F: Fn(GotchaContext) -> FF + Send + 'static, + FF: Future + Send + 'static, +{ + run_interval(context, interval, name, task, CancellationToken::new()).await; +} + +#[cfg(test)] +mod tests; diff --git a/gotcha/src/task/tests.rs b/gotcha/src/task/tests.rs new file mode 100644 index 0000000..37e85fd --- /dev/null +++ b/gotcha/src/task/tests.rs @@ -0,0 +1,169 @@ +use super::*; +use crate::ConfigWrapper; +use std::cell::Cell; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, +}; +use tokio::sync::{mpsc, Notify}; + +fn scheduler() -> TaskScheduler<(), ()> { + TaskScheduler::new(GotchaContext { + state: (), + config: ConfigWrapper::default(), + }) +} + +struct Dropped(mpsc::UnboundedSender<()>); +impl Drop for Dropped { + fn drop(&mut self) { + let _ = self.0.send(()); + } +} + +#[test] +fn registration_is_inert_and_can_be_discarded_without_a_runtime() { + let mut scheduler = scheduler(); + let (dropped, mut drops) = mpsc::unbounded_channel(); + let resource = Dropped(dropped); + scheduler.interval("not-started", Duration::from_secs(1), move |_| -> std::future::Ready<()> { + let _ = &resource; + panic!("registration must not call the task"); + }); + scheduler.cron("invalid", "not a schedule".into(), |_| async {}); + scheduler.interval("zero", Duration::ZERO, |_| async {}); + drop(scheduler); + assert_eq!(drops.try_recv(), Ok(())); +} + +#[tokio::test(start_paused = true)] +async fn graceful_shutdown_finishes_current_execution_without_starting_another() { + let mut scheduler = scheduler(); + let count = Arc::new(AtomicUsize::new(0)); + let finish = Arc::new(Notify::new()); + let (started, mut starts) = mpsc::unbounded_channel(); + scheduler.interval("slow", Duration::from_secs(1), { + let count = count.clone(); + let finish = finish.clone(); + move |_| { + let finish = finish.clone(); + let count = count.clone(); + let started = started.clone(); + async move { + started.send(()).unwrap(); + finish.notified().await; + count.fetch_add(1, Ordering::SeqCst); + } + } + }); + let running = scheduler.start(); + starts.recv().await.unwrap(); + let shutdown = tokio::spawn(running.shutdown(Duration::from_secs(30))); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(5)).await; + assert!(!shutdown.is_finished()); + finish.notify_one(); + shutdown.await.unwrap(); + assert_eq!(count.load(Ordering::SeqCst), 1); + assert_eq!(starts.recv().await, None); +} + +#[tokio::test(start_paused = true)] +async fn timeout_drop_and_cancelled_shutdown_release_in_flight_executions() { + for cron in [false, true] { + for mode in ["timeout", "drop", "cancel-shutdown"] { + let mut scheduler = scheduler(); + let (started, mut starts) = mpsc::unbounded_channel(); + let (dropped, mut drops) = mpsc::unbounded_channel(); + let task = move |_| { + let started = started.clone(); + let dropped = dropped.clone(); + async move { + let _resource = Dropped(dropped); + started.send(()).unwrap(); + std::future::pending::<()>().await; + } + }; + if cron { + scheduler.cron("blocked-cron", "* * * * * * *".into(), task); + } else { + scheduler.interval("blocked-interval", Duration::from_secs(1), task); + } + let running = scheduler.start(); + starts.recv().await.unwrap(); + match mode { + "timeout" => { + let before = tokio::time::Instant::now(); + running.shutdown(Duration::from_secs(10)).await; + assert_eq!(before.elapsed(), Duration::from_secs(10)); + // Explicit shutdown has joined the cancelled future before returning. + assert_eq!(drops.try_recv(), Ok(())); + } + "drop" => { + drop(running); + drops.recv().await.unwrap(); + } + _ => { + let shutdown = tokio::spawn(running.shutdown(Duration::from_secs(30))); + tokio::task::yield_now().await; + shutdown.abort(); + assert!(shutdown.await.unwrap_err().is_cancelled()); + drops.recv().await.unwrap(); + } + } + assert_eq!(starts.recv().await, None); + } + } +} + +#[derive(Clone)] +struct LogBuffer(Arc>>); +impl std::io::Write for LogBuffer { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(bytes); + Ok(bytes.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +#[tokio::test(start_paused = true)] +async fn construction_and_polling_panics_are_logged_and_the_schedule_continues() { + use tracing::instrument::WithSubscriber; + let mut scheduler = scheduler(); + let (completed, mut completions) = mpsc::unbounded_channel(); + let attempts = Cell::new(0); // The callback is Send but deliberately not Sync. + scheduler.interval("recovering", Duration::from_secs(1), move |_| { + let attempt = attempts.get(); + attempts.set(attempt + 1); + assert_ne!(attempt, 0, "construction panic"); + let completed = completed.clone(); + async move { + assert_ne!(attempt, 1, "polling panic"); + completed.send(attempt).unwrap(); + } + }); + let logs = LogBuffer(Arc::default()); + let writer = logs.clone(); + let subscriber = tracing_subscriber::fmt() + .without_time() + .with_ansi(false) + .with_writer(move || writer.clone()) + .finish(); + let dispatch = tracing::Dispatch::new(subscriber); + // Attach the capture subscriber to the scheduled future so it follows Tokio polling. + scheduler.tasks = scheduler + .tasks + .into_iter() + .map(|task| Box::pin(task.with_subscriber(dispatch.clone())) as ScheduledTask) + .collect(); + let running = scheduler.start(); + assert_eq!(completions.recv().await, Some(2)); + running.shutdown(Duration::from_secs(1)).await; + let logs = String::from_utf8(logs.0.lock().unwrap().clone()).unwrap(); + assert!( + logs.contains("recovering") && logs.contains("construction panic") && logs.contains("polling panic"), + "{logs}" + ); +} diff --git a/gotcha/tests/shutdown_signals.rs b/gotcha/tests/shutdown_signals.rs new file mode 100644 index 0000000..382ff67 --- /dev/null +++ b/gotcha/tests/shutdown_signals.rs @@ -0,0 +1,94 @@ +#![cfg(unix)] + +use gotcha::{ConfigWrapper, GotchaApp, GotchaContext, GotchaResult, GotchaRouter, ServerConfig}; +use std::io::{Read, Write}; +use std::path::PathBuf; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +struct SignalApp(PathBuf); +impl GotchaApp for SignalApp { + type State = (); + type Config = (); + fn logger(&self) -> GotchaResult<()> { + Ok(()) + } + async fn config(&self) -> GotchaResult> { + Ok(ConfigWrapper { + app: (), + server: ServerConfig { + host: "127.0.0.1".into(), + port: 0, + }, + }) + } + async fn state(&self, config: &ConfigWrapper<()>) -> GotchaResult<()> { + std::fs::write(&self.0, format!("{}:{}", config.server.host, config.server.port)).unwrap(); + Ok(()) + } + fn routes(&self, router: GotchaRouter>) -> GotchaRouter> { + router.get("/ready", || async { "ready" }) + } +} + +struct Process(Child); +impl Drop for Process { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +#[test] +fn default_signals_shutdown_cleanly() { + const READY_FILE: &str = "GOTCHA_TEST_SHUTDOWN_READY_FILE"; + if let Some(path) = std::env::var_os(READY_FILE) { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(SignalApp(path.into()).run()) + .unwrap(); + return; + } + for signal in ["-INT", "-TERM"] { + let dir = tempfile::tempdir().unwrap(); + let ready = dir.path().join("address"); + let mut process = Process( + Command::new(std::env::current_exe().unwrap()) + .args(["--exact", "default_signals_shutdown_cleanly", "--nocapture"]) + .env(READY_FILE, &ready) + .stdout(Stdio::null()) + .spawn() + .unwrap(), + ); + let deadline = Instant::now() + Duration::from_secs(5); + let address = loop { + if let Ok(address) = std::fs::read_to_string(&ready) { + if !address.is_empty() { + break address; + } + } + assert!(Instant::now() < deadline, "child did not initialize"); + assert!(process.0.try_wait().unwrap().is_none(), "child exited during initialization"); + std::thread::sleep(Duration::from_millis(10)); + }; + // An HTTP round-trip ensures serving and default signal handling have been polled. + let mut stream = std::net::TcpStream::connect(address).unwrap(); + stream.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + stream.write_all(b"GET /ready HTTP/1.0\r\nHost: localhost\r\n\r\n").unwrap(); + let mut response = String::new(); + stream.read_to_string(&mut response).unwrap(); + assert!(response.ends_with("ready")); + assert!(Command::new("kill").args([signal, &process.0.id().to_string()]).status().unwrap().success()); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Some(status) = process.0.try_wait().unwrap() { + assert!(status.success(), "{signal} must return from run() normally: {status}"); + break; + } + assert!(Instant::now() < deadline, "child did not stop after {signal}"); + std::thread::sleep(Duration::from_millis(10)); + } + } +} diff --git a/gotcha/tests/task_lifecycle.rs b/gotcha/tests/task_lifecycle.rs new file mode 100644 index 0000000..6403179 --- /dev/null +++ b/gotcha/tests/task_lifecycle.rs @@ -0,0 +1,285 @@ +#![cfg(feature = "task")] + +use gotcha::{ConfigWrapper, Gotcha, GotchaApp, GotchaContext, GotchaError, GotchaResult, GotchaRouter, ServerConfig, State, TaskScheduler}; +use std::future::Future; +use std::net::SocketAddr; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; +use std::time::Duration; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::sync::{mpsc, Notify}; + +#[derive(Clone)] +struct Controls { + addresses: mpsc::UnboundedSender, + started: mpsc::UnboundedSender<()>, + dropped: mpsc::UnboundedSender<()>, + stop: Arc, + stopping: Arc, + registered: Arc, + finish_task: Arc, + http_started: Arc, + finish_http: Arc, + executions: Arc, + completions: Arc, +} + +fn controls() -> ( + Controls, + mpsc::UnboundedReceiver, + mpsc::UnboundedReceiver<()>, + mpsc::UnboundedReceiver<()>, +) { + let (addresses, address_rx) = mpsc::unbounded_channel(); + let (started, start_rx) = mpsc::unbounded_channel(); + let (dropped, drop_rx) = mpsc::unbounded_channel(); + ( + Controls { + addresses, + started, + dropped, + stop: Arc::default(), + stopping: Arc::default(), + registered: Arc::default(), + finish_task: Arc::default(), + http_started: Arc::default(), + finish_http: Arc::default(), + executions: Arc::default(), + completions: Arc::default(), + }, + address_rx, + start_rx, + drop_rx, + ) +} + +struct Dropped(mpsc::UnboundedSender<()>); +impl Drop for Dropped { + fn drop(&mut self) { + let _ = self.0.send(()); + } +} + +fn config() -> ConfigWrapper<()> { + ConfigWrapper { + app: (), + server: ServerConfig { + host: "127.0.0.1".into(), + port: 0, + }, + } +} + +fn address(config: &ConfigWrapper<()>) -> SocketAddr { + SocketAddr::new(config.server.host.parse().unwrap(), config.server.port) +} + +fn register(scheduler: &mut TaskScheduler) { + scheduler.interval("work", Duration::from_millis(1), |ctx| async move { + let _resource = Dropped(ctx.state.dropped.clone()); + ctx.state.executions.fetch_add(1, Ordering::SeqCst); + ctx.state.addresses.send(address(&ctx.config)).unwrap(); + ctx.state.started.send(()).unwrap(); + ctx.state.finish_task.notified().await; + ctx.state.completions.fetch_add(1, Ordering::SeqCst); + }); +} + +async fn hold(State(ctx): State>) -> &'static str { + ctx.state.http_started.notify_one(); + ctx.state.finish_http.notified().await; + "drained" +} + +#[derive(Clone, Copy)] +enum Registration { + Normal, + Fail, + Wait, +} + +struct App { + controls: Controls, + timeout: Duration, + registration: Registration, +} +impl GotchaApp for App { + type State = Controls; + type Config = (); + fn logger(&self) -> GotchaResult<()> { + Ok(()) + } + async fn config(&self) -> GotchaResult> { + Ok(config()) + } + async fn state(&self, config: &ConfigWrapper<()>) -> GotchaResult { + self.controls.addresses.send(address(config)).unwrap(); + Ok(self.controls.clone()) + } + fn routes(&self, router: GotchaRouter>) -> GotchaRouter> { + router.get("/hold", hold) + } + async fn shutdown_signal(&self) { + // This future borrows the application; no added 'static or Clone bound is needed. + self.controls.stop.notified().await; + self.controls.stopping.notify_one(); + } + fn task_shutdown_timeout(&self) -> Duration { + self.timeout + } + async fn tasks(&self, scheduler: &mut TaskScheduler) -> GotchaResult<()> { + if matches!(self.registration, Registration::Normal) { + register(scheduler); + return Ok(()); + } + let resource = Dropped(self.controls.dropped.clone()); + scheduler.interval("partial-registration", Duration::from_millis(1), move |ctx| { + let _ = &resource; + ctx.state.executions.fetch_add(1, Ordering::SeqCst); + async {} + }); + self.controls.registered.notify_one(); + tokio::task::yield_now().await; // An eagerly spawned registration would run here. + match self.registration { + Registration::Fail => Err(GotchaError::message("registration failed")), + Registration::Wait => std::future::pending().await, + Registration::Normal => unreachable!(), + } + } +} + +async fn within(future: F) -> F::Output { + tokio::time::timeout(Duration::from_secs(5), future) + .await + .expect("lifecycle operation timed out") +} + +async fn request(address: SocketAddr) -> String { + let mut stream = tokio::net::TcpStream::connect(address).await.unwrap(); + stream.write_all(b"GET /hold HTTP/1.0\r\nHost: localhost\r\n\r\n").await.unwrap(); + let mut response = String::new(); + stream.read_to_string(&mut response).await.unwrap(); + response +} + +#[tokio::test] +async fn both_public_apis_share_http_and_task_shutdown_and_own_current_executions() { + for trait_api in [false, true] { + for mode in ["finish", "timeout", "abort"] { + let (controls, mut addresses, mut starts, mut drops) = controls(); + let timeout = if mode == "timeout" { + Duration::from_millis(10) + } else { + Duration::from_secs(30) + }; + let server = if trait_api { + tokio::spawn( + App { + controls: controls.clone(), + timeout, + registration: Registration::Normal, + } + .run(), + ) + } else { + let stop = controls.stop.clone(); + let stopping = controls.stopping.clone(); + tokio::spawn( + Gotcha::from_context(GotchaContext { + state: controls.clone(), + config: config(), + }) + .get("/hold", hold) + .tasks(register) + .shutdown_signal(async move { + stop.notified().await; + stopping.notify_one(); + }) + .task_shutdown_timeout(timeout) + .run(), + ) + }; + let address = within(addresses.recv()).await.unwrap(); + within(starts.recv()).await.unwrap(); + let client = tokio::spawn(request(address)); + within(controls.http_started.notified()).await; + if mode == "abort" { + server.abort(); + assert!(within(server).await.unwrap_err().is_cancelled()); + within(drops.recv()).await.unwrap(); + controls.finish_http.notify_one(); + assert!(within(client).await.unwrap().ends_with("drained")); + } else { + controls.stop.notify_one(); + within(controls.stopping.notified()).await; + if mode == "timeout" { + // The task deadline must run while the HTTP request is still draining. + within(drops.recv()).await.unwrap(); + assert_eq!(controls.completions.load(Ordering::SeqCst), 0); + assert!(!server.is_finished()); + controls.finish_http.notify_one(); + assert!(within(client).await.unwrap().ends_with("drained")); + } else { + controls.finish_http.notify_one(); + assert!(within(client).await.unwrap().ends_with("drained")); + tokio::task::yield_now().await; + assert!(!server.is_finished(), "server must wait for the current scheduled execution"); + controls.finish_task.notify_one(); + } + within(server).await.unwrap().unwrap(); + if mode == "finish" { + assert_eq!(drops.try_recv(), Ok(()), "shutdown must join before returning"); + assert_eq!(controls.completions.load(Ordering::SeqCst), 1); + } + } + assert_eq!(controls.executions.load(Ordering::SeqCst), 1, "shutdown must stop subsequent executions"); + } + } +} + +#[tokio::test] +async fn failed_or_cancelled_registration_discards_tasks_and_releases_the_listener() { + for registration in [Registration::Fail, Registration::Wait] { + let (controls, mut addresses, _starts, mut drops) = controls(); + let app = App { + controls: controls.clone(), + timeout: Duration::from_secs(30), + registration, + }; + let server = tokio::spawn(app.run()); + let address = within(addresses.recv()).await.unwrap(); + within(controls.registered.notified()).await; + if matches!(registration, Registration::Wait) { + server.abort(); + assert!(within(server).await.unwrap_err().is_cancelled()); + } else { + assert!(matches!(within(server).await.unwrap(), Err(GotchaError::Message(message)) if message == "registration failed")); + } + assert_eq!(drops.try_recv(), Ok(())); + assert_eq!(controls.executions.load(Ordering::SeqCst), 0); + let _rebound = tokio::net::TcpListener::bind(address).await.unwrap(); + } +} + +#[tokio::test] +async fn builder_registration_panic_discards_previous_registrations() { + let (controls, _addresses, _starts, mut drops) = controls(); + let dropped = Dropped(controls.dropped.clone()); + let app = Gotcha::from_context(GotchaContext { + state: controls.clone(), + config: config(), + }) + .tasks(move |scheduler| { + scheduler.interval("before-panic", Duration::from_millis(1), move |ctx| { + let _ = &dropped; + ctx.state.executions.fetch_add(1, Ordering::SeqCst); + async {} + }); + }) + .tasks(|_| panic!("registration panic")); + assert!(within(tokio::spawn(app.run())).await.unwrap_err().is_panic()); + assert_eq!(drops.try_recv(), Ok(())); + assert_eq!(controls.executions.load(Ordering::SeqCst), 0); +}