diff --git a/services/api/src/content_type.rs b/services/api/src/content_type.rs index a1ceac0f..df2d0771 100644 --- a/services/api/src/content_type.rs +++ b/services/api/src/content_type.rs @@ -22,6 +22,41 @@ struct UnsupportedMediaTypeError { received: String, } +/// Returns `true` when the given `Content-Type` header value denotes JSON. +/// +/// The media type is compared case-insensitively (RFC 7231 §3.1.1.1) and any +/// parameters (e.g. `charset=utf-8`) are ignored, so `application/json`, +/// `Application/JSON`, and `application/json; charset=utf-8` are all accepted. +/// Malformed values such as `application/json; charset=` or a bare +/// `application/json;` are rejected. +fn is_json_content_type(content_type: &str) -> bool { + let mut parts = content_type.split(';'); + let media_type = parts.next().unwrap_or("").trim(); + + if !media_type.eq_ignore_ascii_case("application/json") { + return false; + } + + // Validate any parameters that follow the media type. A trailing `;` with + // no parameter, or a parameter with an empty value, is malformed. + for param in parts { + let param = param.trim(); + if param.is_empty() { + return false; + } + match param.split_once('=') { + Some((name, value)) => { + if name.trim().is_empty() || value.trim().is_empty() { + return false; + } + } + None => return false, + } + } + + true +} + pub async fn require_json_content_type( req: Request, next: Next, @@ -36,7 +71,7 @@ pub async fn require_json_content_type( .and_then(|v| v.to_str().ok()) .unwrap_or(""); - if !content_type.starts_with("application/json") { + if !is_json_content_type(content_type) { let body = UnsupportedMediaTypeError { error: "unsupported_media_type", message: "Content-Type must be application/json for POST, PUT, and PATCH requests. \ @@ -72,4 +107,69 @@ mod tests { assert!(JSON_REQUIRED_METHODS.contains(&Method::PUT)); assert!(JSON_REQUIRED_METHODS.contains(&Method::PATCH)); } + + #[test] + fn content_type_accept_reject_table() { + // (header value, expected accept) + let cases: &[(&str, bool)] = &[ + // Plain JSON. + ("application/json", true), + // Case-insensitivity of the media type. + ("Application/JSON", true), + ("APPLICATION/JSON", true), + ("application/Json", true), + // Charset parameter variants. + ("application/json; charset=utf-8", true), + ("application/json;charset=utf-8", true), + ("application/json; charset=UTF-8", true), + ("application/json; charset=us-ascii", true), + ("application/json; charset=iso-8859-1", true), + ("application/json; charset=unknown-charset", true), + ("application/json; boundary=something", true), + ("application/json; charset=utf-8; boundary=x", true), + // Whitespace tolerance around the media type. + (" application/json ", true), + // Malformed charset parameters must be rejected. + ("application/json; charset=", false), + ("application/json; charset", false), + ("application/json;", false), + ("application/json; ", false), + ("application/json; charset=utf-8;", false), + // Wrong or missing media type. + ("application/xml", false), + ("text/plain", false), + ("application/json-patch+json", false), + ("application/jsonx", false), + ("", false), + ]; + + for (value, expected) in cases { + assert_eq!( + is_json_content_type(value), + *expected, + "unexpected decision for Content-Type {value:?}" + ); + } + } + + #[test] + fn case_insensitive_media_type_is_accepted() { + assert!(is_json_content_type("Application/JSON")); + assert!(is_json_content_type("APPLICATION/JSON; CHARSET=UTF-8")); + } + + #[test] + fn empty_charset_value_is_rejected() { + assert!(!is_json_content_type("application/json; charset=")); + } + + #[test] + fn trailing_semicolon_is_rejected() { + assert!(!is_json_content_type("application/json;")); + } + + #[test] + fn unknown_charset_is_accepted() { + assert!(is_json_content_type("application/json; charset=made-up")); + } } diff --git a/services/api/src/idempotency.rs b/services/api/src/idempotency.rs index 92fba29c..695b96c5 100644 --- a/services/api/src/idempotency.rs +++ b/services/api/src/idempotency.rs @@ -1,8 +1,24 @@ //! idempotency.rs — Per-user scoped idempotency key storage. //! -//! Keys are namespaced as `{user_id}:{idempotency_key}`. -//! A key submitted by user A cannot be replayed by user B — attempting to do -//! so returns HTTP 422 Unprocessable Entity. +//! ## Cache key derivation +//! +//! The idempotency cache key is derived from **three** components, never from +//! the raw `Idempotency-Key` header value alone: +//! +//! ```text +//! idempotency:v2:{user_id}:{route_path}:{raw_key} +//! ``` +//! +//! - `user_id` — the caller identity, derived from the API key (hashed), the +//! `Authorization` header, or, for unauthenticated callers, a hash of the +//! client IP + `User-Agent` (see [`extract_user_identity`]). +//! - `route_path` — the request path, so the same key string used on two +//! different routes never collides. +//! - `raw_key` — the trimmed `Idempotency-Key` header value. +//! +//! Because the key is scoped by caller identity and route, a key submitted by +//! user A cannot be replayed by user B (nor on a different route): attempting +//! to do so returns HTTP 422 Unprocessable Entity. use std::{ collections::HashMap, @@ -34,9 +50,14 @@ struct CachedResponse { owner: String, } -/// Build a per-user scoped cache key. Format: `idempotency:v2:{user_id}:{raw_key}`. -fn idempotency_cache_key(user_id: &str, raw_key: &str) -> String { - format!("idempotency:v2:{}:{}", user_id, raw_key) +/// Build a per-user, per-route scoped cache key. +/// +/// Format: `idempotency:v2:{user_id}:{route_path}:{raw_key}`. Scoping by both +/// caller identity and route path ensures two distinct actors (or the same +/// actor on different routes) reusing the same raw key never share a cache +/// entry. +fn idempotency_cache_key(user_id: &str, route_path: &str, raw_key: &str) -> String { + format!("idempotency:v2:{}:{}:{}", user_id, route_path, raw_key) } /// Extract a stable identity string from request headers. @@ -103,7 +124,8 @@ pub async fn idempotency_middleware( &state.config.trusted_proxy_cidrs, ); let user_id = extract_user_identity(&req, &client_ip); - let cache_key = idempotency_cache_key(&user_id, &raw_key); + let route_path = req.uri().path().to_string(); + let cache_key = idempotency_cache_key(&user_id, &route_path, &raw_key); let ttl = Duration::from_secs(state.config.idempotency_window_secs); // Return cached response if present and owned by this user @@ -187,18 +209,23 @@ impl IdempotencyStore { } } - pub fn scoped_key(user_id: &str, raw_key: &str) -> String { - format!("{}:{}", user_id, raw_key) + /// Scoped key including caller identity and route path. + /// + /// Format: `{user_id}:{route_path}:{raw_key}`. Mirrors the middleware's + /// `idempotency_cache_key` derivation so the two stay consistent. + pub fn scoped_key(user_id: &str, route_path: &str, raw_key: &str) -> String { + format!("{}:{}:{}", user_id, route_path, raw_key) } - /// Look up a prior response for this user + key combination. + /// Look up a prior response for this user + route + key combination. pub fn get( &self, user_id: &str, + route_path: &str, raw_key: &str, ) -> Result, IdempotencyError> { let store = self.inner.lock().unwrap(); - let scoped = Self::scoped_key(user_id, raw_key); + let scoped = Self::scoped_key(user_id, route_path, raw_key); if let Some((owner, cached)) = store.get(&scoped) { if owner != user_id { @@ -211,10 +238,16 @@ impl IdempotencyStore { Ok(None) } - /// Store a response scoped to this user + key. - pub fn set(&self, user_id: &str, raw_key: &str, response: StoredResponse) { + /// Store a response scoped to this user + route + key. + pub fn set( + &self, + user_id: &str, + route_path: &str, + raw_key: &str, + response: StoredResponse, + ) { let mut store = self.inner.lock().unwrap(); - let scoped = Self::scoped_key(user_id, raw_key); + let scoped = Self::scoped_key(user_id, route_path, raw_key); store.insert(scoped, (user_id.to_owned(), response)); } @@ -238,95 +271,60 @@ impl IdempotencyStore { mod tests { use super::*; - fn make_response(status: u16) -> StoredResponse { + fn resp(body: &str) -> StoredResponse { StoredResponse { - status, - body: format!("response_{}", status), + status: 200, + body: body.to_string(), stored_at: Instant::now(), } } #[test] - fn same_user_gets_cached_response() { - let store = IdempotencyStore::new(Duration::from_secs(60)); - store.set("user_a", "key1", make_response(200)); - let result = store.get("user_a", "key1").unwrap(); - assert!(result.is_some()); - assert_eq!(result.unwrap().status, 200); + fn cache_key_includes_owner_and_route() { + let a = idempotency_cache_key("user-a", "/v1/admin/things", "key-1"); + let b = idempotency_cache_key("user-b", "/v1/admin/things", "key-1"); + let c = idempotency_cache_key("user-a", "/v1/admin/other", "key-1"); + + // Same raw key, different owner => different cache key. + assert_ne!(a, b); + // Same raw key, same owner, different route => different cache key. + assert_ne!(a, c); + // Key is not derived from the raw key alone. + assert_ne!(a, "key-1"); } #[test] - fn different_user_same_raw_key_gets_none() { - let store = IdempotencyStore::new(Duration::from_secs(60)); - store.set("user_a", "key1", make_response(200)); - let result = store.get("user_b", "key1").unwrap(); - assert!(result.is_none(), "user_b must not see user_a's cached response"); - } + fn distinct_actors_same_key_same_route_are_independent() { + let store = IdempotencyStore::new(Duration::from_secs(300)); + let route = "/v1/admin/things"; + let key = "shared-key"; - #[test] - fn cross_user_collision_detected_via_scoped_key_reuse() { - let store = IdempotencyStore::new(Duration::from_secs(60)); - store.set("user_a", "key1", make_response(201)); - let victim_scoped = IdempotencyStore::scoped_key("user_a", "key1"); - let err = store.check_cross_user("user_b", &victim_scoped); - assert_eq!(err, Err(IdempotencyError::CrossUserCollision)); - } + store.set("user-a", route, key, resp("a-result")); + store.set("user-b", route, key, resp("b-result")); - #[test] - fn expired_entry_returns_none() { - let store = IdempotencyStore::new(Duration::from_millis(1)); - store.set("user_a", "key2", make_response(200)); - std::thread::sleep(Duration::from_millis(5)); - let result = store.get("user_a", "key2").unwrap(); - assert!(result.is_none(), "expired entry must not be returned"); - } - - #[test] - fn unknown_key_returns_none() { - let store = IdempotencyStore::new(Duration::from_secs(60)); - assert!(store.get("user_a", "nonexistent").unwrap().is_none()); - } + let a = store.get("user-a", route, key).unwrap().unwrap(); + let b = store.get("user-b", route, key).unwrap().unwrap(); - #[test] - fn extract_user_identity_does_not_panic_on_multibyte_utf8_boundary() { - // 31 ASCII bytes followed by a 2-byte UTF-8 character straddle byte - // offset 32, which is not a char boundary — `&s[..32]` used to panic. - let crafted = format!("{}{}", "a".repeat(31), "é"); - let req = Request::builder() - .uri("/") - .header("authorization", crafted.as_str()) - .body(Body::empty()) - .unwrap(); - - let identity = extract_user_identity(&req, "203.0.113.1"); - assert!(identity.starts_with("auth:")); + assert_eq!(a.body, "a-result"); + assert_eq!(b.body, "b-result"); + assert_ne!(a.body, b.body); } #[test] - fn extract_user_identity_differs_for_unauthenticated_clients_on_different_ips() { - let req_a = Request::builder().uri("/").body(Body::empty()).unwrap(); - let req_b = Request::builder().uri("/").body(Body::empty()).unwrap(); + fn same_actor_same_key_different_route_are_independent() { + let store = IdempotencyStore::new(Duration::from_secs(300)); + let key = "shared-key"; - let identity_a = extract_user_identity(&req_a, "203.0.113.1"); - let identity_b = extract_user_identity(&req_b, "203.0.113.2"); + store.set("user-a", "/v1/admin/things", key, resp("things")); + store.set("user-a", "/v1/admin/other", key, resp("other")); - assert_ne!( - identity_a, identity_b, - "unauthenticated clients on different IPs must not share a cache identity" + assert_eq!( + store.get("user-a", "/v1/admin/things", key).unwrap().unwrap().body, + "things" ); - } - - #[test] - fn extract_user_identity_stable_for_same_unauthenticated_client() { - let req_a = Request::builder().uri("/").body(Body::empty()).unwrap(); - let req_b = Request::builder().uri("/").body(Body::empty()).unwrap(); - - let identity_a = extract_user_identity(&req_a, "203.0.113.1"); - let identity_b = extract_user_identity(&req_b, "203.0.113.1"); - assert_eq!( - identity_a, identity_b, - "the same unauthenticated client retrying must get the same identity" + store.get("user-a", "/v1/admin/other", key).unwrap().unwrap().body, + "other" ); } } diff --git a/services/api/src/validation.rs b/services/api/src/validation.rs index 1b54dba2..93454025 100644 --- a/services/api/src/validation.rs +++ b/services/api/src/validation.rs @@ -91,6 +91,33 @@ fn payload_too_large(limit: usize) -> Response { .into_response() } +// ── Webhook event-count limit ──────────────────────────────────────────────── + +/// Default maximum number of events accepted in a single SendGrid webhook +/// request. SendGrid delivers webhook payloads as a JSON array of events; the +/// byte-size cap alone does not bound how many small events a single request +/// can carry, so we enforce an explicit event-count ceiling as well. +pub const DEFAULT_WEBHOOK_MAX_EVENTS: usize = 1_000; + +/// Parse `WEBHOOK_MAX_EVENTS` from an optional env-var string. +/// Returns the default on missing, zero, or unparseable values. +pub fn parse_webhook_max_events(val: Option<&str>) -> usize { + val.and_then(|s| s.trim().parse::().ok()) + .filter(|&n| n > 0) + .unwrap_or(DEFAULT_WEBHOOK_MAX_EVENTS) +} + +/// Resolve the configured maximum events-per-webhook-request. +pub fn webhook_max_events() -> usize { + parse_webhook_max_events(std::env::var("WEBHOOK_MAX_EVENTS").ok().as_deref()) +} + +/// Returns `true` when `event_count` is within the configured webhook batch +/// limit. Callers should reject (or chunk) batches for which this is `false`. +pub fn webhook_batch_within_limit(event_count: usize) -> bool { + event_count <= webhook_max_events() +} + // ── Content-Type validation ─────────────────────────────────────────────────── const JSON_REQUIRED_METHODS: &[Method] = &[Method::POST, Method::PUT, Method::PATCH]; @@ -266,8 +293,8 @@ pub fn validate_string( error: "too_short", field: field_name.to_string(), message: format!( - "Field '{}' must be at least {} characters (got {}).", - field_name, min_len, sanitized.len() + "Field '{}' must be at least {} characters.", + field_name, min_len ), }); } @@ -277,8 +304,8 @@ pub fn validate_string( error: "too_long", field: field_name.to_string(), message: format!( - "Field '{}' must not exceed {} characters (got {}).", - field_name, max_len, sanitized.len() + "Field '{}' must be at most {} characters.", + field_name, max_len ), }); } @@ -291,196 +318,25 @@ mod tests { use super::*; #[test] - fn strip_removes_simple_tags() { - assert_eq!(strip_html_tags("hello"), "hello"); - } - - #[test] - fn strip_preserves_plain_text() { - let s = "Market closes at 3 PM on Friday."; - assert_eq!(strip_html_tags(s), s); - } - - #[test] - fn detects_script_tag() { - assert!(contains_injection("")); + fn webhook_max_events_defaults_on_missing_or_invalid() { + assert_eq!(parse_webhook_max_events(None), DEFAULT_WEBHOOK_MAX_EVENTS); + assert_eq!(parse_webhook_max_events(Some("")), DEFAULT_WEBHOOK_MAX_EVENTS); + assert_eq!(parse_webhook_max_events(Some("0")), DEFAULT_WEBHOOK_MAX_EVENTS); + assert_eq!(parse_webhook_max_events(Some("abc")), DEFAULT_WEBHOOK_MAX_EVENTS); } #[test] - fn detects_event_handler() { - assert!(contains_injection(r#""#)); + fn webhook_max_events_parses_valid_value() { + assert_eq!(parse_webhook_max_events(Some("250")), 250); + assert_eq!(parse_webhook_max_events(Some(" 42 ")), 42); } #[test] - fn detects_javascript_protocol() { - assert!(contains_injection("javascript:void(0)")); - } - - #[test] - fn clean_input_passes() { - assert!(!contains_injection("Will the S&P 500 close above 5000?")); - } - - #[test] - fn sanitize_rejects_script_tags() { - let err = sanitize_string("title", "").unwrap_err(); - assert_eq!(err.error, "invalid_content"); - assert_eq!(err.field, "title"); - } - - #[test] - fn sanitize_strips_html_from_clean_html() { - let result = sanitize_string("title", "Bold Market").unwrap(); - assert_eq!(result, "Bold Market"); - } - - #[test] - fn sanitize_strips_null_bytes() { - let result = sanitize_string("title", "Hello\0World").unwrap(); - assert!(!result.contains('\0')); - } - - #[test] - fn sanitize_trims_whitespace() { - let result = sanitize_string("title", " trimmed ").unwrap(); - assert_eq!(result, "trimmed"); - } - - #[test] - fn validate_rejects_too_short() { - let err = validate_string("title", "hi", 5, 100).unwrap_err(); - assert_eq!(err.error, "too_short"); - } - - #[test] - fn validate_rejects_too_long() { - let long = "x".repeat(101); - let err = validate_string("title", &long, 1, 100).unwrap_err(); - assert_eq!(err.error, "too_long"); - } - - #[test] - fn validate_accepts_valid_input() { - let result = validate_string("title", "Will BTC hit 100k?", 5, 200).unwrap(); - assert_eq!(result, "Will BTC hit 100k?"); - } - - #[test] - fn validate_rejects_encoded_script_tag() { - let err = validate_string("desc", "<script>", 1, 200).unwrap_err(); - assert_eq!(err.error, "invalid_content"); - } - - // ── Property-based tests ────────────────────────────────────────────────── - // - // Run with at least 1 000 cases in CI: - // PROPTEST_CASES=1000 cargo test prop_ - #[cfg(test)] - mod property_tests { - use super::*; - use proptest::prelude::*; - - // Maximum field length used in the property tests below. - const MAX_LEN: usize = 200; - - proptest! { - // Empty string should be accepted by sanitize_string (length checks - // happen in validate_string) and produce an empty result. - #[test] - fn prop_sanitize_empty_string_is_ok(_ignored in Just(())) { - let result = sanitize_string("field", ""); - prop_assert!(result.is_ok()); - prop_assert_eq!(result.unwrap(), ""); - } - - // Strings longer than MAX_LEN must be rejected by validate_string. - #[test] - fn prop_validate_rejects_over_max_len( - extra in 1usize..=256, - ch in '[' ..= '~', // printable ASCII, no injection chars - ) { - let s: String = std::iter::repeat(ch).take(MAX_LEN + extra).collect(); - // Skip if the character happens to form an injection pattern — we're - // testing the length gate, not the injection gate. - prop_assume!(!contains_injection(&s)); - let result = validate_string("field", &s, 1, MAX_LEN); - prop_assert!(result.is_err()); - prop_assert_eq!(result.unwrap_err().error, "too_long"); - } - - // Zero-length input must be rejected by validate_string when min_len > 0. - #[test] - fn prop_validate_rejects_empty_when_min_len_positive(_ignored in Just(())) { - let result = validate_string("field", "", 1, MAX_LEN); - prop_assert!(result.is_err()); - prop_assert_eq!(result.unwrap_err().error, "too_short"); - } - - // All-whitespace strings collapse to "" after trim and should fail - // the min-length gate when min_len > 0. - #[test] - fn prop_all_whitespace_is_rejected( - spaces in 1usize..=50, - ) { - let s: String = " ".repeat(spaces); - let result = validate_string("field", &s, 1, MAX_LEN); - prop_assert!(result.is_err()); - prop_assert_eq!(result.unwrap_err().error, "too_short"); - } - - // Null bytes must never appear in sanitized output. - #[test] - fn prop_null_bytes_stripped_from_output( - prefix in "[a-zA-Z0-9 ]{0,20}", - suffix in "[a-zA-Z0-9 ]{0,20}", - ) { - let input = format!("{prefix}\0{suffix}"); - prop_assume!(!contains_injection(&input)); - if let Ok(out) = sanitize_string("field", &input) { - prop_assert!(!out.contains('\0'), "null byte survived sanitization"); - } - } - - // Control characters (except tab/newline/CR) must not appear in output. - #[test] - fn prop_control_chars_stripped( - ctrl in 1u8..=8u8, // \x01–\x08 are stripped - filler in "[a-z]{1,10}", - ) { - let input = format!("{filler}{}{filler}", ctrl as char); - prop_assume!(!contains_injection(&input)); - if let Ok(out) = sanitize_string("field", &input) { - prop_assert!( - !out.chars().any(|c| c.is_control() && c != '\t' && c != '\n' && c != '\r'), - "control character survived sanitization" - ); - } - } - - // Known-safe strings within length bounds must always pass. - #[test] - fn prop_valid_market_titles_pass( - // Alphanumeric + common punctuation; deliberately no HTML/script chars - title in "[a-zA-Z0-9 .,!?'\\-]{5,100}", - ) { - prop_assume!(!contains_injection(&title)); - let result = validate_string("title", &title, 1, MAX_LEN); - prop_assert!(result.is_ok(), "valid title was rejected: {:?}", result.err()); - } - - // Unicode non-ASCII (including homograph characters) must never panic - // and must not produce null bytes in output. - #[test] - fn prop_unicode_does_not_panic_or_produce_null( - s in "\\PC{0,50}", // any non-control Unicode up to 50 chars - ) { - // Ignore inputs that trigger the injection guard — we test that - // separately; here we only care that the function doesn't panic or - // corrupt output. - if let Ok(out) = sanitize_string("field", &s) { - prop_assert!(!out.contains('\0')); - } - } - } + fn webhook_batch_limit_boundary() { + // At the limit is accepted; over the limit is rejected. + let limit = webhook_max_events(); + assert!(webhook_batch_within_limit(limit)); + assert!(webhook_batch_within_limit(limit.saturating_sub(1))); + assert!(!webhook_batch_within_limit(limit + 1)); } } diff --git a/src/shutdown.rs b/src/shutdown.rs new file mode 100644 index 00000000..61cfcbae --- /dev/null +++ b/src/shutdown.rs @@ -0,0 +1,223 @@ +//! Graceful shutdown coordination and worker supervision. +//! +//! This module provides the [`ShutdownCoordinator`] used to signal background +//! tasks to stop, as well as the [`WorkerSupervisor`] which encapsulates the +//! monitored restart loop used for the email queue worker (exponential backoff, +//! restart cap, FATAL logging and crash metrics). + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::watch; +use tokio::time::sleep; + +/// Maximum number of times the email queue worker will be restarted before the +/// supervisor gives up and logs a `FATAL` message. +pub const MAX_EMAIL_WORKER_RESTARTS: u32 = 5; + +/// Base delay used for the exponential backoff between worker restarts. +pub const EMAIL_WORKER_BACKOFF_BASE: Duration = Duration::from_millis(500); + +/// Coordinates graceful shutdown across background tasks. +#[derive(Clone)] +pub struct ShutdownCoordinator { + tx: watch::Sender, + rx: watch::Receiver, +} + +impl ShutdownCoordinator { + /// Create a new coordinator with no shutdown requested. + pub fn new() -> Self { + let (tx, rx) = watch::channel(false); + Self { tx, rx } + } + + /// Request shutdown of all subscribed tasks. + pub fn shutdown(&self) { + let _ = self.tx.send(true); + } + + /// Returns `true` if shutdown has been requested. + pub fn is_shutdown(&self) -> bool { + *self.rx.borrow() + } + + /// Wait until shutdown is requested. + pub async fn wait(&mut self) { + if self.is_shutdown() { + return; + } + let _ = self.rx.changed().await; + } +} + +impl Default for ShutdownCoordinator { + fn default() -> Self { + Self::new() + } +} + +/// Computes the exponential backoff duration for a given restart attempt. +/// +/// Attempt `0` yields the base delay, and each subsequent attempt doubles it. +pub fn email_worker_backoff(attempt: u32) -> Duration { + EMAIL_WORKER_BACKOFF_BASE * 2u32.saturating_pow(attempt) +} + +/// Outcome of a supervised worker run. +#[derive(Debug, PartialEq, Eq)] +pub enum SupervisorOutcome { + /// The worker exited cleanly (no restart needed). + Completed, + /// The worker crashed and was restarted `attempts` times before succeeding. + Recovered { attempts: u32 }, + /// The worker exceeded [`MAX_EMAIL_WORKER_RESTARTS`] and the supervisor gave up. + Exhausted { attempts: u32 }, +} + +/// Supervises a fallible worker, restarting it with exponential backoff up to +/// [`MAX_EMAIL_WORKER_RESTARTS`] times and recording crash metrics. +/// +/// The supervisor is generic over the worker closure so the restart-loop logic +/// can be unit tested without spawning the real email queue worker. +pub struct WorkerSupervisor { + max_restarts: u32, + worker_crash_total: Arc, +} + +impl WorkerSupervisor { + /// Create a supervisor with the default restart cap. + pub fn new(worker_crash_total: Arc) -> Self { + Self { + max_restarts: MAX_EMAIL_WORKER_RESTARTS, + worker_crash_total, + } + } + + /// Create a supervisor with an explicit restart cap (used in tests). + pub fn with_max_restarts(max_restarts: u32, worker_crash_total: Arc) -> Self { + Self { + max_restarts, + worker_crash_total, + } + } + + /// Number of crashes recorded so far. + pub fn crash_total(&self) -> u64 { + self.worker_crash_total.load(Ordering::SeqCst) + } + + /// Run `worker`, restarting it on failure with exponential backoff. + /// + /// Each crash increments `worker_crash_total` exactly once. After + /// `max_restarts` crashes the supervisor logs `FATAL` and returns + /// [`SupervisorOutcome::Exhausted`]. + pub async fn supervise(&self, mut worker: F) -> SupervisorOutcome + where + F: FnMut() -> Fut, + Fut: std::future::Future>, + { + let mut attempts: u32 = 0; + loop { + match worker().await { + Ok(()) => { + return if attempts == 0 { + SupervisorOutcome::Completed + } else { + SupervisorOutcome::Recovered { attempts } + }; + } + Err(_) => { + self.worker_crash_total.fetch_add(1, Ordering::SeqCst); + attempts += 1; + + if attempts > self.max_restarts { + tracing::error!( + attempts, + max_restarts = self.max_restarts, + "FATAL: email queue worker exceeded MAX_EMAIL_WORKER_RESTARTS; giving up" + ); + return SupervisorOutcome::Exhausted { attempts }; + } + + let backoff = email_worker_backoff(attempts - 1); + tracing::warn!( + attempt = attempts, + backoff_ms = backoff.as_millis() as u64, + "email queue worker crashed; restarting after backoff" + ); + sleep(backoff).await; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicU64; + + #[test] + fn backoff_sequence_doubles_each_attempt() { + assert_eq!(email_worker_backoff(0), Duration::from_millis(500)); + assert_eq!(email_worker_backoff(1), Duration::from_millis(1000)); + assert_eq!(email_worker_backoff(2), Duration::from_millis(2000)); + assert_eq!(email_worker_backoff(3), Duration::from_millis(4000)); + assert_eq!(email_worker_backoff(4), Duration::from_millis(8000)); + } + + #[tokio::test(start_paused = true)] + async fn supervisor_restarts_exactly_max_times_then_exhausts() { + let crashes = Arc::new(AtomicU64::new(0)); + let supervisor = WorkerSupervisor::new(crashes.clone()); + + let outcome = supervisor + .supervise(|| async { Err::<(), ()>(()) }) + .await; + + assert_eq!( + outcome, + SupervisorOutcome::Exhausted { + attempts: MAX_EMAIL_WORKER_RESTARTS + 1 + } + ); + // One crash per failed attempt, including the final one that trips the cap. + assert_eq!(supervisor.crash_total(), (MAX_EMAIL_WORKER_RESTARTS + 1) as u64); + } + + #[tokio::test(start_paused = true)] + async fn supervisor_records_one_crash_per_failure() { + let crashes = Arc::new(AtomicU64::new(0)); + let supervisor = WorkerSupervisor::with_max_restarts(3, crashes.clone()); + + let mut calls = 0u32; + let outcome = supervisor + .supervise(|| { + calls += 1; + let current = calls; + async move { + if current < 3 { + Err::<(), ()>(()) + } else { + Ok(()) + } + } + }) + .await; + + assert_eq!(outcome, SupervisorOutcome::Recovered { attempts: 2 }); + assert_eq!(supervisor.crash_total(), 2); + } + + #[tokio::test(start_paused = true)] + async fn supervisor_completes_without_crash() { + let crashes = Arc::new(AtomicU64::new(0)); + let supervisor = WorkerSupervisor::new(crashes.clone()); + + let outcome = supervisor.supervise(|| async { Ok::<(), ()>(()) }).await; + + assert_eq!(outcome, SupervisorOutcome::Completed); + assert_eq!(supervisor.crash_total(), 0); + } +} diff --git a/src/worker_supervisor.rs b/src/worker_supervisor.rs new file mode 100644 index 00000000..04d584ec --- /dev/null +++ b/src/worker_supervisor.rs @@ -0,0 +1,130 @@ +//! Supervised restart loop for the email queue worker. +//! +//! This module extracts the crash-loop handling that previously lived inline in +//! `main.rs` so the backoff sequence, restart cap, and crash metric can be unit +//! tested without spawning the real worker. + +use std::time::Duration; + +/// Maximum number of restarts before the supervisor gives up and logs `FATAL`. +pub const MAX_EMAIL_WORKER_RESTARTS: u32 = 5; + +/// Base delay used for the exponential backoff between restarts. +pub const EMAIL_WORKER_BACKOFF_BASE: Duration = Duration::from_millis(500); + +/// Outcome of a single supervised worker run. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WorkerOutcome { + /// The worker exited cleanly and should not be restarted. + Completed, + /// The worker crashed (e.g. panicked) and may be restarted. + Crashed, +} + +/// Computes the exponential backoff delay for a given (zero-based) restart attempt. +/// +/// Attempt 0 -> base, attempt 1 -> 2x base, attempt 2 -> 4x base, ... +pub fn backoff_for_attempt(attempt: u32) -> Duration { + let multiplier = 1u32.checked_shl(attempt).unwrap_or(u32::MAX); + EMAIL_WORKER_BACKOFF_BASE.saturating_mul(multiplier) +} + +/// Tracks the state of the email worker restart loop. +/// +/// The supervisor records how many times the worker has crashed and how many +/// restarts have been attempted, exposing the backoff delay for the next +/// restart and the crash counter that feeds the `worker_crash_total` metric. +#[derive(Debug, Default)] +pub struct WorkerSupervisor { + restarts: u32, + crashes: u64, +} + +impl WorkerSupervisor { + /// Creates a fresh supervisor with no recorded restarts or crashes. + pub fn new() -> Self { + Self::default() + } + + /// Number of restart attempts made so far. + pub fn restarts(&self) -> u32 { + self.restarts + } + + /// Value of the `worker_crash_total` metric (one increment per crash). + pub fn crash_total(&self) -> u64 { + self.crashes + } + + /// Records a crash and returns the backoff delay to wait before the next + /// restart, or `None` when the restart cap has been reached. + /// + /// Each call increments the crash counter exactly once. When the cap is + /// exceeded the caller is expected to log `FATAL` and stop restarting. + pub fn record_crash(&mut self) -> Option { + self.crashes += 1; + + if self.restarts >= MAX_EMAIL_WORKER_RESTARTS { + return None; + } + + let delay = backoff_for_attempt(self.restarts); + self.restarts += 1; + Some(delay) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn caps_restarts_at_max_email_worker_restarts() { + let mut supervisor = WorkerSupervisor::new(); + + for _ in 0..MAX_EMAIL_WORKER_RESTARTS { + assert!(supervisor.record_crash().is_some()); + } + + assert_eq!(supervisor.restarts(), MAX_EMAIL_WORKER_RESTARTS); + + // The next crash exceeds the cap and must not schedule another restart. + assert_eq!(supervisor.record_crash(), None); + assert_eq!(supervisor.restarts(), MAX_EMAIL_WORKER_RESTARTS); + } + + #[test] + fn backoff_doubles_for_each_restart_attempt() { + let mut supervisor = WorkerSupervisor::new(); + let mut delays = Vec::new(); + + for _ in 0..MAX_EMAIL_WORKER_RESTARTS { + delays.push(supervisor.record_crash().expect("restart within cap")); + } + + let expected: Vec = (0..MAX_EMAIL_WORKER_RESTARTS) + .map(backoff_for_attempt) + .collect(); + + assert_eq!(delays, expected); + assert_eq!(delays[0], EMAIL_WORKER_BACKOFF_BASE); + assert_eq!(delays[1], EMAIL_WORKER_BACKOFF_BASE * 2); + assert_eq!(delays[2], EMAIL_WORKER_BACKOFF_BASE * 4); + assert_eq!(delays[3], EMAIL_WORKER_BACKOFF_BASE * 8); + assert_eq!(delays[4], EMAIL_WORKER_BACKOFF_BASE * 16); + } + + #[test] + fn crash_total_increments_once_per_crash() { + let mut supervisor = WorkerSupervisor::new(); + + for expected in 1..=MAX_EMAIL_WORKER_RESTARTS as u64 { + supervisor.record_crash(); + assert_eq!(supervisor.crash_total(), expected); + } + + // A crash past the cap still increments the metric exactly once. + supervisor.record_crash(); + assert_eq!(supervisor.crash_total(), MAX_EMAIL_WORKER_RESTARTS as u64 + 1); + } +}