diff --git a/crates/service/job/seedance2_pro_job/src/job_dependencies.rs b/crates/service/job/seedance2_pro_job/src/job_dependencies.rs index 26abd75112..6341744973 100644 --- a/crates/service/job/seedance2_pro_job/src/job_dependencies.rs +++ b/crates/service/job/seedance2_pro_job/src/job_dependencies.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use chrono::Duration; use cloud_storage::bucket_client::BucketClient; use concurrency::relaxed_atomic_bool::RelaxedAtomicBool; @@ -6,6 +8,7 @@ use pager::client::pager::Pager; use seedance2pro_client::creds::seedance2pro_session::Seedance2ProSession; use server_environment::ServerEnvironment; use sqlx::MySqlPool; +use tokio::sync::Notify; #[derive(Clone)] pub struct JobDependencies { @@ -40,6 +43,10 @@ pub struct JobDependencies { /// Set to `true` from another thread to trigger graceful shutdown. pub application_shutdown: RelaxedAtomicBool, + /// Notified when `application_shutdown` is set. Allows sleeping tasks + /// to wake up immediately instead of waiting for the full sleep duration. + pub shutdown_notify: Arc, + /// Pager client for sending alerts. pub pager: Pager, } diff --git a/crates/service/job/seedance2_pro_job/src/jobs/character_polling_job/character_polling_main_loop.rs b/crates/service/job/seedance2_pro_job/src/jobs/character_polling_job/character_polling_main_loop.rs index 05bceb4250..a154ea80d9 100644 --- a/crates/service/job/seedance2_pro_job/src/jobs/character_polling_job/character_polling_main_loop.rs +++ b/crates/service/job/seedance2_pro_job/src/jobs/character_polling_job/character_polling_main_loop.rs @@ -29,7 +29,10 @@ pub async fn character_polling_main_loop(deps: JobDependencies) { warn!("Character poll iteration took {:.1}s (threshold: {}s)", elapsed.as_secs_f64(), POLL_ALERT_THRESHOLD.as_secs()); } - tokio::time::sleep(Duration::from_millis(deps.poll_interval_millis)).await; + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis(deps.poll_interval_millis)) => {} + _ = deps.shutdown_notify.notified() => {} + } } warn!("Character polling main loop is shut down."); diff --git a/crates/service/job/seedance2_pro_job/src/jobs/credits_checking_job/credits_checking_main_loop.rs b/crates/service/job/seedance2_pro_job/src/jobs/credits_checking_job/credits_checking_main_loop.rs index fcbb6698ce..95bdf8d18f 100644 --- a/crates/service/job/seedance2_pro_job/src/jobs/credits_checking_job/credits_checking_main_loop.rs +++ b/crates/service/job/seedance2_pro_job/src/jobs/credits_checking_job/credits_checking_main_loop.rs @@ -41,7 +41,10 @@ pub async fn credits_checking_main_loop(deps: JobDependencies) { } } - tokio::time::sleep(CREDITS_CHECK_INTERVAL).await; + tokio::select! { + _ = tokio::time::sleep(CREDITS_CHECK_INTERVAL) => {} + _ = deps.shutdown_notify.notified() => {} + } } warn!("Credits checking main loop is shut down."); diff --git a/crates/service/job/seedance2_pro_job/src/jobs/video_polling_job/video_polling_main_loop.rs b/crates/service/job/seedance2_pro_job/src/jobs/video_polling_job/video_polling_main_loop.rs index 8540619b50..d2279ba3f0 100644 --- a/crates/service/job/seedance2_pro_job/src/jobs/video_polling_job/video_polling_main_loop.rs +++ b/crates/service/job/seedance2_pro_job/src/jobs/video_polling_job/video_polling_main_loop.rs @@ -49,7 +49,10 @@ pub async fn video_polling_main_loop(job_dependencies: JobDependencies) { } } - tokio::time::sleep(Duration::from_millis(job_dependencies.poll_interval_millis)).await; + tokio::select! { + _ = tokio::time::sleep(Duration::from_millis(job_dependencies.poll_interval_millis)) => {} + _ = job_dependencies.shutdown_notify.notified() => {} + } } warn!("Kinovi job runner main loop is shut down."); diff --git a/crates/service/job/seedance2_pro_job/src/main.rs b/crates/service/job/seedance2_pro_job/src/main.rs index 5c54759da6..05a4fe25c6 100644 --- a/crates/service/job/seedance2_pro_job/src/main.rs +++ b/crates/service/job/seedance2_pro_job/src/main.rs @@ -9,11 +9,13 @@ #[macro_use] extern crate serde_derive; +use std::sync::Arc; use std::time::Duration; use anyhow::anyhow; use log::{info, warn}; use sqlx::mysql::MySqlPoolOptions; +use tokio::sync::Notify; use bootstrap::bootstrap::{bootstrap, BootstrapArgs}; use cloud_storage::bucket_client::BucketClient; @@ -142,6 +144,7 @@ async fn main() -> AnyhowResult<()> { }); let application_shutdown = RelaxedAtomicBool::new(false); + let shutdown_notify = Arc::new(Notify::new()); let job_stats = JobStats::new(); let pager_for_shutdown = pager.clone(); @@ -163,6 +166,7 @@ async fn main() -> AnyhowResult<()> { maybe_max_job_age, credits_alert_threshold, application_shutdown: application_shutdown.clone(), + shutdown_notify: shutdown_notify.clone(), pager, }; @@ -179,12 +183,14 @@ async fn main() -> AnyhowResult<()> { // Listen for SIGTERM / Ctrl-C to trigger graceful shutdown. let application_shutdown_for_signal = application_shutdown.clone(); + let shutdown_notify_for_signal = shutdown_notify.clone(); tokio::spawn(async move { match tokio::signal::ctrl_c().await { Ok(()) => { info!("Received shutdown signal. Shutting down..."); application_shutdown_for_signal.set(true); + shutdown_notify_for_signal.notify_waiters(); } Err(err) => { warn!("Error listening for shutdown signal: {:?}", err);