diff --git a/Cargo.toml b/Cargo.toml index d1de917..2d28279 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ log = "0.4.29" toml = "0.9.10" reqwest = { version = "0.13.1", features = ["json", "blocking", "multipart", "stream", "query"] } anyhow = "1.0.100" -tokio = { version = "1.49.0", features = ["rt", "rt-multi-thread", "macros", "fs"] } +tokio = { version = "1.49.0", features = ["rt", "rt-multi-thread", "macros", "fs", "sync", "time"] } async-trait = "0.1.89" tempfile = "3.24.0" openssl = "0.10.75" diff --git a/src/services/backup/dispatcher.rs b/src/services/backup/dispatcher.rs index bf9145a..5401e87 100644 --- a/src/services/backup/dispatcher.rs +++ b/src/services/backup/dispatcher.rs @@ -1,9 +1,37 @@ use super::service::BackupService; use crate::services::api::models::agent::status::DatabaseStorage; use crate::services::config::DatabasesConfig; +use crate::settings::CONFIG; use crate::utils::common::BackupMethod; +use once_cell::sync::Lazy; +use std::future::Future; +use tokio::sync::Semaphore; use tracing::error; +pub(crate) static BACKUP_SEMAPHORE: Lazy> = + Lazy::new(|| CONFIG.max_concurrent_backups.map(Semaphore::new)); + +/// Runs `job` after acquiring a permit from `semaphore` (or unthrottled if `None`), +/// holding it for the duration of `job`. Shared by `dispatch()` and its tests so the +/// concurrency-capping behavior under test is the same code path production uses. +pub(crate) async fn run_with_permit(semaphore: Option<&Semaphore>, job: F) +where + F: Future, +{ + let _permit = match semaphore { + Some(semaphore) => match semaphore.acquire().await { + Ok(permit) => Some(permit), + Err(e) => { + error!("Failed to acquire backup concurrency permit: {}", e); + return; + } + }, + None => None, + }; + + job.await; +} + impl BackupService { pub async fn dispatch( &self, @@ -31,12 +59,15 @@ impl BackupService { let generated_id = generated_id.clone(); tokio::spawn(async move { - if let Err(e) = service - .execute_backup(generated_id, db_cfg, method, storages, encrypt) - .await - { - error!("Backup execution failed: {}", e); - } + run_with_permit(BACKUP_SEMAPHORE.as_ref(), async move { + if let Err(e) = service + .execute_backup(generated_id, db_cfg, method, storages, encrypt) + .await + { + error!("Backup execution failed: {}", e); + } + }) + .await; }); } } diff --git a/src/settings.rs b/src/settings.rs index b4076d1..0d27b28 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -1,6 +1,7 @@ use dotenvy::dotenv; use once_cell::sync::Lazy; use std::env; +use tokio::sync::Semaphore; #[derive(Debug)] #[allow(dead_code)] @@ -15,7 +16,8 @@ pub struct Settings { pub pooling: usize, pub timezone: String, pub log: String, - pub chunk_size: usize, // bytes + pub chunk_size: usize, // bytes + pub max_concurrent_backups: Option, // None = unlimited } impl Settings { @@ -49,6 +51,9 @@ impl Settings { let chunk_size = chunk_size_mb * 1024 * 1024; + let max_concurrent_backups = + parse_max_concurrent_backups(env::var("MAX_CONCURRENT_BACKUPS")); + let tz = env::var("TZ").unwrap_or_else(|_| "UTC".to_string()); Self { @@ -64,9 +69,86 @@ impl Settings { pooling: pooling_seconds, timezone: tz, log: env::var("LOG").unwrap_or_else(|_| "info".into()), - chunk_size + chunk_size, + max_concurrent_backups, } } } pub static CONFIG: Lazy = Lazy::new(Settings::from_env); + +/// None = unset/blank (unlimited). Panics on a malformed or zero value rather than +/// silently falling back to unlimited, matching this file's other env-parsed fields. +pub(crate) fn parse_max_concurrent_backups(value: Result) -> Option { + match value { + Ok(val) if val.trim().is_empty() => None, + Ok(val) => { + let parsed = val + .trim() + .parse::() + .expect("MAX_CONCURRENT_BACKUPS must be a valid positive integer"); + + if parsed == 0 { + panic!("MAX_CONCURRENT_BACKUPS must be at least 1"); + } + if parsed > Semaphore::MAX_PERMITS { + panic!( + "MAX_CONCURRENT_BACKUPS must not exceed {}", + Semaphore::MAX_PERMITS + ); + } + + Some(parsed) + } + Err(_) => None, + } +} + +#[cfg(test)] +mod tests { + use super::parse_max_concurrent_backups; + use std::env::VarError; + + #[test] + fn unset_means_unlimited() { + assert_eq!( + parse_max_concurrent_backups(Err(VarError::NotPresent)), + None + ); + } + + #[test] + fn blank_means_unlimited() { + assert_eq!(parse_max_concurrent_backups(Ok("".to_string())), None); + assert_eq!(parse_max_concurrent_backups(Ok(" ".to_string())), None); + } + + #[test] + fn valid_value_is_parsed() { + assert_eq!(parse_max_concurrent_backups(Ok("2".to_string())), Some(2)); + } + + #[test] + fn surrounding_whitespace_is_trimmed() { + assert_eq!(parse_max_concurrent_backups(Ok(" 2 ".to_string())), Some(2)); + } + + #[test] + #[should_panic(expected = "must be at least 1")] + fn zero_panics() { + parse_max_concurrent_backups(Ok("0".to_string())); + } + + #[test] + #[should_panic(expected = "must be a valid positive integer")] + fn malformed_value_panics() { + parse_max_concurrent_backups(Ok("not-a-number".to_string())); + } + + #[test] + #[should_panic(expected = "must not exceed")] + fn value_exceeding_semaphore_max_permits_panics() { + let over_limit = (tokio::sync::Semaphore::MAX_PERMITS as u128 + 1).to_string(); + parse_max_concurrent_backups(Ok(over_limit)); + } +} diff --git a/src/tests/services/backup_dispatcher_tests.rs b/src/tests/services/backup_dispatcher_tests.rs new file mode 100644 index 0000000..21d8ccb --- /dev/null +++ b/src/tests/services/backup_dispatcher_tests.rs @@ -0,0 +1,71 @@ +//! Regression coverage for the concurrency cap added alongside the #94 panic fixes. +//! `BackupService::dispatch` gates `execute_backup` behind `BACKUP_SEMAPHORE` via the +//! shared `run_with_permit` helper (`src/services/backup/dispatcher.rs`), configurable +//! via `MAX_CONCURRENT_BACKUPS`. + +use crate::services::backup::dispatcher::{BACKUP_SEMAPHORE, run_with_permit}; +use crate::tests::init_tracing_for_test; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use tokio::sync::Semaphore; +use tokio::time::{Duration, sleep}; + +#[tokio::test] +async fn backup_semaphore_defaults_to_unlimited_when_max_concurrent_backups_unset() { + init_tracing_for_test(); + + // Neither this test suite nor docker-compose.test.yml sets MAX_CONCURRENT_BACKUPS, + // so the real, process-wide BACKUP_SEMAPHORE must be None: dispatch() must not + // throttle backups unless an operator explicitly opts in. This depends on ambient + // environment (CONFIG is a process-wide Lazy singleton) — see settings::tests for + // hermetic coverage of the underlying parsing logic that doesn't have this caveat. + assert!( + BACKUP_SEMAPHORE.is_none(), + "expected no concurrency cap by default (MAX_CONCURRENT_BACKUPS unset)" + ); +} + +#[tokio::test] +async fn semaphore_gated_execution_never_exceeds_configured_limit() { + init_tracing_for_test(); + + // Calls the same run_with_permit() that dispatch() uses, so this exercises the + // real gating code path rather than a parallel reimplementation. A local Semaphore + // is used instead of the global BACKUP_SEMAPHORE because CONFIG.max_concurrent_backups + // is a process-wide singleton read once from the environment and can't be + // reconfigured per-test. + let limit = 2; + let semaphore = Arc::new(Semaphore::new(limit)); + let concurrent = Arc::new(AtomicUsize::new(0)); + let max_observed = Arc::new(AtomicUsize::new(0)); + + let mut handles = Vec::new(); + for _ in 0..6 { + let semaphore = semaphore.clone(); + let concurrent = concurrent.clone(); + let max_observed = max_observed.clone(); + + handles.push(tokio::spawn(async move { + run_with_permit(Some(semaphore.as_ref()), async move { + let now = concurrent.fetch_add(1, Ordering::SeqCst) + 1; + max_observed.fetch_max(now, Ordering::SeqCst); + + sleep(Duration::from_millis(50)).await; + + concurrent.fetch_sub(1, Ordering::SeqCst); + }) + .await; + })); + } + + for handle in handles { + handle.await.unwrap(); + } + + assert!( + max_observed.load(Ordering::SeqCst) <= limit, + "observed {} concurrent jobs, expected at most {}", + max_observed.load(Ordering::SeqCst), + limit + ); +} diff --git a/src/tests/services/mod.rs b/src/tests/services/mod.rs index af4f3eb..d8e5c74 100644 --- a/src/tests/services/mod.rs +++ b/src/tests/services/mod.rs @@ -1,3 +1,4 @@ mod api_models_tests; +mod backup_dispatcher_tests; mod backup_uploader_tests; mod config_tests; diff --git a/src/tests/utils/mod.rs b/src/tests/utils/mod.rs index 858d6c1..af011da 100644 --- a/src/tests/utils/mod.rs +++ b/src/tests/utils/mod.rs @@ -5,3 +5,4 @@ mod edge_key_tests; mod file_tests; mod normalize_cron_tests; mod stream_tests; +mod task_manager_tests; diff --git a/src/tests/utils/task_manager_tests.rs b/src/tests/utils/task_manager_tests.rs new file mode 100644 index 0000000..2d1a998 --- /dev/null +++ b/src/tests/utils/task_manager_tests.rs @@ -0,0 +1,111 @@ +use crate::tests::init_tracing_for_test; +use crate::utils::task_manager::cron::check_and_update_cron; +use crate::utils::task_manager::scheduler::{execute_task, scheduler_loop}; +use crate::utils::task_manager::tasks::SCHEDULE_KEY; +use redis::AsyncCommands; +use redis::aio::MultiplexedConnection; +use std::time::Duration; +use testcontainers::ContainerAsync; +use testcontainers::runners::AsyncRunner; +use testcontainers_modules::redis::Redis; +use url::Host; + +#[tokio::test] +async fn execute_task_errors_instead_of_panicking_on_missing_generated_id() { + let result = execute_task("tasks.database.periodic_backup", vec![], None).await; + + let err = result.expect_err("expected an error, not a panic, for empty args"); + assert!( + err.to_string().contains("generated_id"), + "expected error to mention missing generated_id, got: {err}" + ); +} + +#[tokio::test] +async fn execute_task_errors_instead_of_panicking_on_missing_dbms() { + let args = vec!["some-generated-id".to_string()]; + let result = execute_task("tasks.database.periodic_backup", args, None).await; + + let err = result.expect_err("expected an error, not a panic, for a single arg"); + assert!( + err.to_string().contains("dbms"), + "expected error to mention missing dbms, got: {err}" + ); +} + +async fn start_redis() -> (ContainerAsync, MultiplexedConnection) { + let container = Redis::default().start().await.unwrap(); + + let host = container + .get_host() + .await + .unwrap_or(Host::parse("127.0.0.1").unwrap()); + let port = container.get_host_port_ipv4(6379).await.unwrap_or(6379); + + let client = redis::Client::open(format!("redis://{host}:{port}")).unwrap(); + let conn = client.get_multiplexed_async_connection().await.unwrap(); + + (container, conn) +} + +#[tokio::test] +async fn check_and_update_cron_does_not_panic_on_malformed_stored_data() { + init_tracing_for_test(); + let (_container, mut conn) = start_redis().await; + + let task_name = "malformed-task"; + let redis_key = format!("redbeat:{task_name}"); + + let _: () = conn + .hset(&redis_key, "data", "not-valid-json") + .await + .unwrap(); + + // Before the fix this unwrapped serde_json::from_str and panicked. It should + // now log the parse failure and return gracefully instead. + check_and_update_cron( + &mut conn, + Some("*/5 * * * *".to_string()), + vec![], + "tasks.database.periodic_backup", + task_name.to_string(), + None, + ) + .await; +} + +#[tokio::test] +async fn scheduler_loop_skips_malformed_due_task_without_panicking() { + init_tracing_for_test(); + let (_container, mut conn) = start_redis().await; + + let task_name = "malformed-due-task"; + let redis_key = format!("redbeat:{task_name}"); + let now = chrono::Local::now().timestamp(); + + let _: () = conn + .hset(&redis_key, "data", "not-valid-json") + .await + .unwrap(); + let _: () = conn.zadd(SCHEDULE_KEY, &redis_key, now).await.unwrap(); + + let handle = tokio::spawn(scheduler_loop(conn)); + let abort_handle = handle.abort_handle(); + + // Before the fix this unwrapped serde_json::from_str inside the loop body + // (not a spawned subtask), so the panic would unwind scheduler_loop itself + // and the JoinHandle would resolve within the timeout. Post-fix it logs the + // parse failure, continues, and keeps looping forever, so the timeout should + // elapse instead. + let outcome = tokio::time::timeout(Duration::from_millis(1500), handle).await; + abort_handle.abort(); + + match outcome { + Err(_elapsed) => { + // Still running after processing the due tick: it survived the bad payload. + } + Ok(join_result) => { + panic!("scheduler_loop exited instead of continuing to loop: {join_result:?}"); + } + } +} diff --git a/src/utils/task_manager/cron.rs b/src/utils/task_manager/cron.rs index 002662f..318ba91 100644 --- a/src/utils/task_manager/cron.rs +++ b/src/utils/task_manager/cron.rs @@ -43,8 +43,20 @@ pub async fn check_and_update_cron( debug!("Task cron normalized: unix \"{}\" -> crate \"{}\"", raw_cron, cron); if exists { - let raw: String = conn.hget(&redis_key, "data").await.unwrap(); - let stored: models::PeriodicTask = serde_json::from_str(&raw).unwrap(); + let raw: String = match conn.hget(&redis_key, "data").await { + Ok(raw) => raw, + Err(e) => { + tracing::error!("Failed to load task data for {}: {:?}", task_name, e); + return; + } + }; + let stored: models::PeriodicTask = match serde_json::from_str(&raw) { + Ok(stored) => stored, + Err(e) => { + tracing::error!("Failed to parse task data for {}: {:?}", task_name, e); + return; + } + }; let cron_changed = stored.cron != cron; let args_changed = stored.args != args; diff --git a/src/utils/task_manager/scheduler.rs b/src/utils/task_manager/scheduler.rs index 2972f60..adfb842 100644 --- a/src/utils/task_manager/scheduler.rs +++ b/src/utils/task_manager/scheduler.rs @@ -17,13 +17,28 @@ pub async fn scheduler_loop(mut conn: MultiplexedConnection) { loop { let now = chrono::Local::now().timestamp(); - let due: Vec = conn - .zrangebyscore(SCHEDULE_KEY, 0, now) - .await - .unwrap_or_default(); + let due: Vec = match conn.zrangebyscore(SCHEDULE_KEY, 0, now).await { + Ok(due) => due, + Err(e) => { + error!("Failed to fetch due tasks from {}: {:?}", SCHEDULE_KEY, e); + Vec::new() + } + }; for key in due { - let raw: String = conn.hget(&key, "data").await.unwrap(); - let task: PeriodicTask = serde_json::from_str(&raw).unwrap(); + let raw: String = match conn.hget(&key, "data").await { + Ok(raw) => raw, + Err(e) => { + error!("Failed to load task data for key={}: {:?}", key, e); + continue; + } + }; + let task: PeriodicTask = match serde_json::from_str(&raw) { + Ok(task) => task, + Err(e) => { + error!("Failed to parse task data for key={}: {:?}", key, e); + continue; + } + }; if !task.enabled { continue; @@ -50,7 +65,14 @@ pub async fn scheduler_loop(mut conn: MultiplexedConnection) { } match next_run_timestamp(&task_clone.cron) { Some(next_ts) => { - let _: () = conn_clone.zadd(SCHEDULE_KEY, &key, next_ts).await.unwrap(); + let result: redis::RedisResult<()> = + conn_clone.zadd(SCHEDULE_KEY, &key, next_ts).await; + if let Err(e) = result { + error!( + "Failed to reschedule task={} key={}: {:?}", + task_clone.task, key, e + ); + } } None => { error!( @@ -72,14 +94,18 @@ pub async fn execute_task( ) -> Result<(), anyhow::Error> { match task { "tasks.database.periodic_backup" => { - let generated_id = &args[0]; - let dbms = &args[1]; + let generated_id = args + .first() + .ok_or_else(|| anyhow::anyhow!("Missing generated_id argument"))?; + let dbms = args + .get(1) + .ok_or_else(|| anyhow::anyhow!("Missing dbms argument"))?; info!("{} | {}", generated_id, dbms); let ctx = Arc::new(Context::new()); let config_service = ConfigService::new(ctx.clone()); let backup_service = BackupService::new(ctx.clone()); - let config = config_service.load(None).unwrap(); + let config = config_service.load(None).map_err(|e| anyhow::anyhow!(e))?; let metadata_obj = metadata .into_iter()