From 5d2e5d5b48bbfac4c965e979ebb289dc9d878ceb Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Wed, 12 Mar 2025 11:51:53 +0100 Subject: [PATCH] fix: prevent monitoring task to die without sending killpill (#5472) --- backend/ee-repo-ref.txt | 2 +- backend/src/main.rs | 20 ++---- backend/src/monitor.rs | 34 +++------- backend/tests/worker.rs | 10 ++- backend/windmill-api/src/lib.rs | 6 +- backend/windmill-common/src/lib.rs | 64 ++++++++++++++++++- backend/windmill-common/src/worker.rs | 7 +- .../windmill-indexer/src/service_logs_ee.rs | 4 +- .../windmill-worker/src/dedicated_worker.rs | 17 +++-- .../windmill-worker/src/result_processor.rs | 13 ++-- backend/windmill-worker/src/worker.rs | 7 +- 11 files changed, 109 insertions(+), 75 deletions(-) diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index 4c2e4d7211..5d5699c639 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -d3e20c73b06a7c6820868769a49b5fa621591653 \ No newline at end of file +9a3f58425136b0f9e9c8106151537b1e271ec838 \ No newline at end of file diff --git a/backend/src/main.rs b/backend/src/main.rs index 33d562b1dd..656aca22a9 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -46,7 +46,7 @@ use windmill_common::{ stats_ee::schedule_stats, utils::{hostname, rd_string, Mode, GIT_VERSION, MODE_AND_ADDONS}, worker::{reload_custom_tags_setting, HUB_CACHE_DIR, TMP_DIR, TMP_LOGS_DIR, WORKER_GROUP}, - DB, METRICS_ENABLED, + KillpillSender, DB, METRICS_ENABLED, }; #[cfg(all(not(target_env = "msvc"), feature = "jemalloc"))] @@ -361,7 +361,7 @@ async fn windmill_main() -> anyhow::Result<()> { let db = windmill_common::connect_db(server_mode, indexer_mode, worker_mode).await?; - let (killpill_tx, mut killpill_rx) = tokio::sync::broadcast::channel::<()>(2); + let (killpill_tx, mut killpill_rx) = KillpillSender::new(2); let mut monitor_killpill_rx = killpill_tx.subscribe(); let (killpill_phase2_tx, _killpill_phase2_rx) = tokio::sync::broadcast::channel::<()>(2); let server_killpill_rx = killpill_phase2_tx.subscribe(); @@ -619,7 +619,7 @@ Windmill Community Edition {GIT_VERSION} ) .await?; tracing::info!("All workers exited."); - killpill_tx.send(())?; + killpill_tx.send(); } else { rx.recv().await?; } @@ -643,7 +643,6 @@ Windmill Community Edition {GIT_VERSION} let base_internal_url = base_internal_url.to_string(); let h = tokio::spawn(async move { let mut listener = retry_listen_pg(&db).await; - loop { tokio::select! { biased; @@ -868,6 +867,7 @@ Windmill Community Edition {GIT_VERSION} tracing::error!("Error waiting for monitor handle: {e:#}") } tracing::info!("Monitor exited"); + killpill_tx.send(); Ok(()) as anyhow::Result<()> }; @@ -982,7 +982,7 @@ fn display_config(envs: &[&str]) { pub async fn run_workers( db: Pool, mut rx: tokio::sync::broadcast::Receiver<()>, - tx: tokio::sync::broadcast::Sender<()>, + tx: KillpillSender, num_workers: i32, base_internal_url: String, agent_mode: bool, @@ -1095,11 +1095,7 @@ pub async fn run_workers( Ok(()) } -async fn send_delayed_killpill( - tx: &tokio::sync::broadcast::Sender<()>, - mut max_delay_secs: u64, - context: &str, -) { +async fn send_delayed_killpill(tx: &KillpillSender, mut max_delay_secs: u64, context: &str) { if max_delay_secs == 0 { max_delay_secs = 1; } @@ -1108,7 +1104,5 @@ async fn send_delayed_killpill( tracing::info!("Scheduling {context} shutdown in {rd_delay}s"); tokio::time::sleep(Duration::from_secs(rd_delay)).await; - if let Err(e) = tx.send(()) { - tracing::error!(error = %e, "Could not send killpill for {context}"); - } + tx.send(); } diff --git a/backend/src/monitor.rs b/backend/src/monitor.rs index 6fe8070014..f348bb9fab 100644 --- a/backend/src/monitor.rs +++ b/backend/src/monitor.rs @@ -34,10 +34,7 @@ use windmill_common::ee::{jobs_waiting_alerts, worker_groups_alerts}; #[cfg(feature = "oauth2")] use windmill_common::global_settings::OAUTH_SETTING; use windmill_common::{ - ee::CriticalErrorChannel, - error, - flow_status::{FlowStatus, FlowStatusModule}, - global_settings::{ + ee::CriticalErrorChannel, error, flow_status::{FlowStatus, FlowStatusModule}, global_settings::{ BASE_URL_SETTING, BUNFIG_INSTALL_SCOPES_SETTING, CRITICAL_ALERT_MUTE_UI_SETTING, CRITICAL_ERROR_CHANNELS_SETTING, DEFAULT_TAGS_PER_WORKSPACE_SETTING, DEFAULT_TAGS_WORKSPACES_SETTING, EXPOSE_DEBUG_METRICS_SETTING, EXPOSE_METRICS_SETTING, @@ -47,24 +44,11 @@ use windmill_common::{ NUGET_CONFIG_SETTING, OTEL_SETTING, PIP_INDEX_URL_SETTING, REQUEST_SIZE_LIMIT_SETTING, REQUIRE_PREEXISTING_USER_FOR_OAUTH_SETTING, RETENTION_PERIOD_SECS_SETTING, SAML_METADATA_SETTING, SCIM_TOKEN_SETTING, TIMEOUT_WAIT_RESULT_SETTING, - }, - indexer::load_indexer_config, - jobs::QueuedJob, - jwt::JWT_SECRET, - oauth2::REQUIRE_PREEXISTING_USER_FOR_OAUTH, - server::load_smtp_config, - tracing_init::JSON_FMT, - users::truncate_token, - utils::{now_from_db, rd_string, report_critical_error, Mode}, - worker::{ + }, indexer::load_indexer_config, jobs::QueuedJob, jwt::JWT_SECRET, oauth2::REQUIRE_PREEXISTING_USER_FOR_OAUTH, server::load_smtp_config, tracing_init::JSON_FMT, users::truncate_token, utils::{now_from_db, rd_string, report_critical_error, Mode}, worker::{ load_worker_config, make_pull_query, make_suspended_pull_query, reload_custom_tags_setting, update_min_version, DEFAULT_TAGS_PER_WORKSPACE, DEFAULT_TAGS_WORKSPACES, INDEXER_CONFIG, SMTP_CONFIG, TMP_DIR, WORKER_CONFIG, WORKER_GROUP, - }, - BASE_URL, CRITICAL_ALERT_MUTE_UI_ENABLED, CRITICAL_ERROR_CHANNELS, DB, DEFAULT_HUB_BASE_URL, - HUB_BASE_URL, JOB_RETENTION_SECS, METRICS_DEBUG_ENABLED, METRICS_ENABLED, - MONITOR_LOGS_ON_OBJECT_STORE, OTEL_LOGS_ENABLED, OTEL_METRICS_ENABLED, OTEL_TRACING_ENABLED, - SERVICE_LOG_RETENTION_SECS, + }, KillpillSender, BASE_URL, CRITICAL_ALERT_MUTE_UI_ENABLED, CRITICAL_ERROR_CHANNELS, DB, DEFAULT_HUB_BASE_URL, HUB_BASE_URL, JOB_RETENTION_SECS, METRICS_DEBUG_ENABLED, METRICS_ENABLED, MONITOR_LOGS_ON_OBJECT_STORE, OTEL_LOGS_ENABLED, OTEL_METRICS_ENABLED, OTEL_TRACING_ENABLED, SERVICE_LOG_RETENTION_SECS }; use windmill_queue::cancel_job; use windmill_worker::{ @@ -133,7 +117,7 @@ lazy_static::lazy_static! { pub async fn initial_load( db: &Pool, - tx: tokio::sync::broadcast::Sender<()>, + tx: KillpillSender, worker_mode: bool, server_mode: bool, #[cfg(feature = "parquet")] disable_s3_store: bool, @@ -1265,7 +1249,7 @@ pub async fn monitor_db( server_mode: bool, _worker_mode: bool, initial_load: bool, - _killpill_tx: tokio::sync::broadcast::Sender<()>, + _killpill_tx: KillpillSender, ) { let zombie_jobs_f = async { if server_mode && !initial_load { @@ -1441,7 +1425,7 @@ pub async fn reload_indexer_config(db: &Pool) { pub async fn reload_worker_config( db: &DB, - tx: tokio::sync::broadcast::Sender<()>, + tx: KillpillSender, kill_if_change: bool, ) { let config = load_worker_config(&db, tx.clone()).await; @@ -1456,17 +1440,17 @@ pub async fn reload_worker_config( || (*wc).dedicated_worker != config.dedicated_worker { tracing::info!("Dedicated worker config changed, sending killpill. Expecting to be restarted by supervisor."); - let _ = tx.send(()); + let _ = tx.send(); } if (*wc).init_bash != config.init_bash { tracing::info!("Init bash config changed, sending killpill. Expecting to be restarted by supervisor."); - let _ = tx.send(()); + let _ = tx.send(); } if (*wc).cache_clear != config.cache_clear { tracing::info!("Cache clear changed, sending killpill. Expecting to be restarted by supervisor."); - let _ = tx.send(()); + let _ = tx.send(); tracing::info!("Waiting 5 seconds to allow others workers to start potential jobs that depend on a potential shared cache volume"); tokio::time::sleep(Duration::from_secs(5)).await; if let Err(e) = windmill_worker::common::clean_cache().await { diff --git a/backend/tests/worker.rs b/backend/tests/worker.rs index 2f7abde33d..0e0e209a5b 100644 --- a/backend/tests/worker.rs +++ b/backend/tests/worker.rs @@ -2,6 +2,7 @@ use serde::de::DeserializeOwned; use std::future::Future; use std::{str::FromStr, sync::Arc}; use windmill_api_client::types::{NewScript, ScriptLang as NewScriptLanguage}; +use windmill_common::KillpillSender; #[cfg(feature = "enterprise")] use chrono::Timelike; @@ -999,7 +1000,7 @@ async fn in_test_worker( }; /* ensure the worker quits before we return */ - quit.send(()).expect("send"); + quit.send(); let _: () = worker .await @@ -1011,16 +1012,13 @@ async fn in_test_worker( fn spawn_test_worker( db: &Pool, port: u16, -) -> ( - tokio::sync::broadcast::Sender<()>, - tokio::task::JoinHandle<()>, -) { +) -> (KillpillSender, tokio::task::JoinHandle<()>) { std::fs::DirBuilder::new() .recursive(true) .create(windmill_worker::GO_BIN_CACHE_DIR) .expect("could not create initial worker dir"); - let (tx, rx) = tokio::sync::broadcast::channel(1); + let (tx, rx) = KillpillSender::new(1); let db = db.to_owned(); let worker_instance: &str = "test worker instance"; let worker_name: String = next_worker_name(); diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 843a409906..11113775e0 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -401,7 +401,7 @@ pub async fn run_server( let db_killpill_rx = rx.resubscribe(); postgres_triggers::start_database(db.clone(), db_killpill_rx); } - + #[cfg(feature = "mqtt_trigger")] { let mqtt_killpill_rx = rx.resubscribe(); @@ -600,7 +600,9 @@ pub async fn run_server( .on_failure(MyOnFailure {}), ) }; - let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); + let listener = tokio::net::TcpListener::bind(addr) + .await + .context("binding main windmill server")?; let port = listener.local_addr().map(|x| x.port()).unwrap_or(8000); let ip = listener .local_addr() diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 9eeeec3cdf..c4387cb8c3 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -9,9 +9,14 @@ use std::{ net::SocketAddr, str::FromStr, - sync::{atomic::AtomicBool, Arc}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, }; +use tokio::sync::broadcast; + use ee::CriticalErrorChannel; use error::Error; use scripts::ScriptLang; @@ -115,7 +120,7 @@ lazy_static::lazy_static! { } pub async fn shutdown_signal( - tx: tokio::sync::broadcast::Sender<()>, + tx: KillpillSender, mut rx: tokio::sync::broadcast::Receiver<()>, ) -> anyhow::Result<()> { use std::io; @@ -151,7 +156,7 @@ pub async fn shutdown_signal( } tracing::info!("signal received, starting graceful shutdown"); - let _ = tx.send(()); + let _ = tx.send(); Ok(()) } @@ -420,3 +425,56 @@ pub async fn get_latest_hash_for_path<'c>( script.created_by, )) } + +pub struct KillpillSender { + tx: broadcast::Sender<()>, + already_sent: Arc, +} + +impl Clone for KillpillSender { + fn clone(&self) -> Self { + KillpillSender { tx: self.tx.clone(), already_sent: self.already_sent.clone() } + } +} + +impl KillpillSender { + pub fn new(capacity: usize) -> (Self, broadcast::Receiver<()>) { + let (tx, rx) = broadcast::channel(capacity); + let sender = KillpillSender { tx, already_sent: Arc::new(AtomicBool::new(false)) }; + (sender, rx) + } + + pub fn clone(&self) -> Self { + KillpillSender { tx: self.tx.clone(), already_sent: self.already_sent.clone() } + } + + pub fn subscribe(&self) -> broadcast::Receiver<()> { + self.tx.subscribe() + } + + // Try to send the killpill if it hasn't been sent already + pub fn send(&self) -> bool { + // Check if it's already been sent, and if not, set the flag to true + if !self.already_sent.swap(true, Ordering::SeqCst) { + // We're the first to set it to true, so send the signal + if let Err(e) = self.tx.send(()) { + tracing::error!("failed to send killpill: {:?}", e); + } + true + } else { + // Signal was already sent + false + } + } + + // // Force send a signal regardless of previous sends + // fn force_send(&self) -> Result> { + // self.already_sent.store(true, Ordering::SeqCst); + // self.tx.send(()) + // } + + // // Check if the killpill has been sent + // fn is_sent(&self) -> bool { + // self.already_sent.load(Ordering::SeqCst) + // } +} diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index b2429f5ae0..7a1f53774b 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -19,7 +19,8 @@ use tokio::sync::RwLock; use windmill_macros::annotations; use crate::{ - error, global_settings::CUSTOM_TAGS_SETTING, indexer::TantivyIndexerSettings, server::Smtp, DB, + error, global_settings::CUSTOM_TAGS_SETTING, indexer::TantivyIndexerSettings, server::Smtp, + KillpillSender, DB, }; lazy_static::lazy_static! { @@ -716,7 +717,7 @@ pub async fn update_ping(worker_instance: &str, worker_name: &str, ip: &str, db: pub async fn load_worker_config( db: &DB, - killpill_tx: tokio::sync::broadcast::Sender<()>, + killpill_tx: KillpillSender, ) -> error::Result { tracing::info!("Loading config from WORKER_GROUP: {}", *WORKER_GROUP); let mut config: WorkerConfigOpt = sqlx::query_scalar!( @@ -750,7 +751,7 @@ pub async fn load_worker_config( .map(|x| { let splitted = x.split(':').to_owned().collect_vec(); if splitted.len() != 2 { - killpill_tx.send(()).expect("send"); + killpill_tx.send(); return Err(anyhow::anyhow!( "Invalid dedicated_worker format. Got {x}, expects :" )); diff --git a/backend/windmill-indexer/src/service_logs_ee.rs b/backend/windmill-indexer/src/service_logs_ee.rs index 49a4f12447..e04575d3f5 100644 --- a/backend/windmill-indexer/src/service_logs_ee.rs +++ b/backend/windmill-indexer/src/service_logs_ee.rs @@ -1,7 +1,7 @@ use anyhow::anyhow; use sqlx::{Pool, Postgres}; use windmill_common::error::Error; - +use windmill_common::KillpillSender; #[derive(Clone)] pub struct ServiceLogIndexReader; @@ -10,7 +10,7 @@ pub struct ServiceLogIndexWriter; pub async fn init_index( _db: &Pool, - mut _killpill_rx: tokio::sync::broadcast::Sender<()>, + mut _killpill_tx: KillpillSender, ) -> Result<(ServiceLogIndexReader, ServiceLogIndexWriter), Error> { Err(anyhow!("Cannot initialize index: not in EE").into()) } diff --git a/backend/windmill-worker/src/dedicated_worker.rs b/backend/windmill-worker/src/dedicated_worker.rs index fb3a30176c..b67483dd9c 100644 --- a/backend/windmill-worker/src/dedicated_worker.rs +++ b/backend/windmill-worker/src/dedicated_worker.rs @@ -15,6 +15,7 @@ use tokio::{ use windmill_common::error::Error; use windmill_common::flows::FlowValue; use windmill_common::worker::WORKER_CONFIG; +use windmill_common::KillpillSender; use windmill_common::{ cache, error, flows::{FlowModule, FlowModuleValue}, @@ -252,7 +253,7 @@ async fn spawn_dedicated_workers_for_flow( modules: &Vec, w_id: &str, path: &str, - killpill_tx: tokio::sync::broadcast::Sender<()>, + killpill_tx: KillpillSender, killpill_rx: &tokio::sync::broadcast::Receiver<()>, db: &DB, worker_dir: &str, @@ -438,7 +439,7 @@ async fn spawn_dedicated_workers_for_flow( } pub async fn create_dedicated_worker_map( - killpill_tx: &tokio::sync::broadcast::Sender<()>, + killpill_tx: &KillpillSender, killpill_rx: &tokio::sync::broadcast::Receiver<()>, db: &DB, worker_dir: &str, @@ -551,7 +552,7 @@ pub enum SpawnWorker { async fn spawn_dedicated_worker( sw: SpawnWorker, w_id: &str, - killpill_tx: tokio::sync::broadcast::Sender<()>, + killpill_tx: KillpillSender, killpill_rx: &tokio::sync::broadcast::Receiver<()>, db: &DB, worker_dir: &str, @@ -638,7 +639,7 @@ async fn spawn_dedicated_worker( } } else { tracing::error!("Failed to fetch script for dedicated worker"); - killpill_tx.send(()).expect("send"); + killpill_tx.send(); return None; } } @@ -670,7 +671,7 @@ async fn spawn_dedicated_worker( .await { tracing::error!("failed to create token for dedicated worker: {:?}", e); - killpill_tx.clone().send(()).expect("send"); + killpill_tx.clone().send(); }; token }; @@ -682,7 +683,7 @@ async fn spawn_dedicated_worker( #[cfg(not(feature = "python"))] { tracing::error!("Python requires the python feature to be enabled"); - killpill_tx.send(()).expect("send"); + killpill_tx.send(); return; } @@ -744,9 +745,7 @@ async fn spawn_dedicated_worker( } { tracing::error!("error in dedicated worker for {sw:#?}: {:?}", e); }; - if let Err(e) = killpill_tx.clone().send(()) { - tracing::error!("failed to send final killpill to dedicated worker: {:?}", e); - } + killpill_tx.clone().send(); }); return Some((node_id.unwrap_or(path2), dedicated_worker_tx, Some(handle))); // (Some(dedi_path), Some(dedicated_worker_tx), Some(handle)) diff --git a/backend/windmill-worker/src/result_processor.rs b/backend/windmill-worker/src/result_processor.rs index bb4c1937ae..153d6f32fc 100644 --- a/backend/windmill-worker/src/result_processor.rs +++ b/backend/windmill-worker/src/result_processor.rs @@ -22,7 +22,7 @@ use windmill_common::{ jobs::{JobKind, QueuedJob}, utils::WarnAfterExt, worker::{to_raw_value, WORKER_GROUP}, - DB, + KillpillSender, DB, }; #[cfg(feature = "benchmark")] @@ -33,10 +33,7 @@ use windmill_queue::{append_logs, get_queued_job, CanceledBy, WrappedError}; use serde_json::{json, value::RawValue}; use tokio::{ - sync::{ - self, - mpsc::{Receiver, Sender}, - }, + sync::mpsc::{Receiver, Sender}, task::JoinHandle, }; @@ -59,7 +56,7 @@ pub fn start_background_processor( worker_dir: String, same_worker_tx: SameWorkerSender, worker_name: String, - killpill_tx: sync::broadcast::Sender<()>, + killpill_tx: KillpillSender, is_dedicated_worker: bool, ) -> JoinHandle<()> { tokio::spawn(async move { @@ -150,7 +147,7 @@ pub fn start_background_processor( if is_init_script_and_failure { tracing::error!("init script errored, exiting"); - killpill_tx.send(()).unwrap_or_default(); + killpill_tx.send(); break; } if is_dependency_job && is_dedicated_worker { @@ -162,7 +159,7 @@ pub fn start_background_processor( .execute(&db) .await .expect("update config to trigger restart of all dedicated workers at that config"); - killpill_tx.send(()).unwrap_or_default(); + killpill_tx.send(); } add_time!(bench, "job completed processed"); diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index db6ea56bb5..7c474efd40 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -20,6 +20,7 @@ use windmill_common::{ get_memory, get_vcpus, get_windmill_memory_usage, get_worker_memory_usage, write_file, ROOT_CACHE_DIR, ROOT_CACHE_NOMOUNT_DIR, TMP_DIR, }, + KillpillSender, }; #[cfg(feature = "enterprise")] @@ -746,7 +747,7 @@ pub async fn run_worker( _num_workers: u32, ip: &str, mut killpill_rx: tokio::sync::broadcast::Receiver<()>, - killpill_tx: tokio::sync::broadcast::Sender<()>, + killpill_tx: KillpillSender, base_internal_url: &str, agent_mode: bool, ) { @@ -1105,7 +1106,7 @@ pub async fn run_worker( if i_worker == 1 { if let Err(e) = queue_init_bash_maybe(db, same_worker_tx.clone(), &worker_name).await { - killpill_tx.send(()).unwrap_or_default(); + killpill_tx.send(); tracing::error!(worker = %worker_name, hostname = %hostname, "Error queuing init bash script for worker {worker_name}: {e:#}"); return; } @@ -1231,7 +1232,7 @@ pub async fn run_worker( tracing::error!( worker = %worker_name, hostname = %hostname, "failed to update worker ping, exiting: {}", e); - killpill_tx.send(()).unwrap_or_default(); + killpill_tx.send(); } tracing::info!( worker = %worker_name, hostname = %hostname,