From 6198ab5a466990ebdbfa386cb44096b94b4101d0 Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Fri, 3 Jul 2026 16:40:17 -0400 Subject: [PATCH 1/5] Fix deadlock --- components/spider-storage/src/db/mariadb.rs | 96 ++++++++++++++++----- 1 file changed, 74 insertions(+), 22 deletions(-) diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index cac6d9ea2..a70af2731 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -343,10 +343,10 @@ impl InternalJobOrchestration for MariaDbStorageConnector { ) -> Result, DbError> { const DELETE_BATCH_SIZE: usize = 1000; - const SELECT_QUERY: &str = formatcp!( + const SELECT_CANDIDATES_QUERY: &str = formatcp!( "SELECT `id` FROM `{table}` WHERE `state` IN \ ('{succeeded_state}','{failed_state}','{cancelled_state}') AND `ended_at` < NOW() - \ - INTERVAL ? SECOND LIMIT {DELETE_BATCH_SIZE} FOR UPDATE;", + INTERVAL ? SECOND LIMIT {DELETE_BATCH_SIZE};", table = JOBS_TABLE_NAME, succeeded_state = JobState::Succeeded.as_str(), failed_state = JobState::Failed.as_str(), @@ -357,28 +357,53 @@ impl InternalJobOrchestration for MariaDbStorageConnector { let mut tx = self.pool.begin().await?; loop { - let job_id_batch: Vec = sqlx::query_scalar(SELECT_QUERY) + let candidate_ids: Vec = sqlx::query_scalar(SELECT_CANDIDATES_QUERY) .bind(expire_after_sec) .fetch_all(&mut *tx) .await?; - if job_id_batch.is_empty() { + if candidate_ids.is_empty() { break; } - let placeholders = std::iter::repeat_n("?", job_id_batch.len()) + let placeholders = std::iter::repeat_n("?", candidate_ids.len()) .collect::>() .join(","); - let delete_query = - format!("DELETE FROM `{JOBS_TABLE_NAME}` WHERE `id` IN ({placeholders})"); + let select_for_update_query = format!( + "SELECT `id` FROM `{table}` FORCE INDEX (PRIMARY) WHERE `id` IN ({placeholders}) \ + AND `state` IN ('{succeeded_state}','{failed_state}','{cancelled_state}') AND \ + `ended_at` < NOW() - INTERVAL ? SECOND ORDER BY `id` FOR UPDATE;", + table = JOBS_TABLE_NAME, + succeeded_state = JobState::Succeeded.as_str(), + failed_state = JobState::Failed.as_str(), + cancelled_state = JobState::Cancelled.as_str(), + ); + let mut select_query = sqlx::query_scalar::<_, JobId>(&select_for_update_query); + for job_id in &candidate_ids { + select_query = select_query.bind(job_id); + } + let confirmed_ids: Vec = select_query + .bind(expire_after_sec) + .fetch_all(&mut *tx) + .await?; - let mut query = sqlx::query(&delete_query); - for job_id in &job_id_batch { - query = query.bind(job_id); + if !confirmed_ids.is_empty() { + let placeholders = std::iter::repeat_n("?", confirmed_ids.len()) + .collect::>() + .join(","); + let delete_query = + format!("DELETE FROM `{JOBS_TABLE_NAME}` WHERE `id` IN ({placeholders})"); + let mut delete_query = sqlx::query(&delete_query); + for job_id in &confirmed_ids { + delete_query = delete_query.bind(job_id); + } + delete_query.execute(&mut *tx).await?; + deleted_job_ids.extend(confirmed_ids); } - query.execute(&mut *tx).await?; - deleted_job_ids.extend(job_id_batch); + if candidate_ids.len() < DELETE_BATCH_SIZE { + break; + } } tx.commit().await?; @@ -558,21 +583,46 @@ impl ExecutionManagerLivenessManagement for MariaDbStorageConnector { ) -> Result, DbError> { const UPDATE_BATCH_SIZE: usize = 1000; - const SELECT_QUERY: &str = formatcp!( + const SELECT_CANDIDATES_QUERY: &str = formatcp!( "SELECT `id` FROM `{table}` WHERE `state` = '{alive_state}' AND `last_heartbeat_at` < \ - CURRENT_TIMESTAMP - INTERVAL ? SECOND FOR UPDATE;", + CURRENT_TIMESTAMP - INTERVAL ? SECOND ORDER BY `id`;", table = EXECUTION_MANAGERS_TABLE_NAME, alive_state = ExecutionManagerState::Alive.as_str(), ); let mut tx = self.pool.begin().await?; - let execution_manager_ids: Vec = sqlx::query_scalar(SELECT_QUERY) + let candidate_ids: Vec = sqlx::query_scalar(SELECT_CANDIDATES_QUERY) .bind(stale_after_sec) .fetch_all(&mut *tx) .await?; - for execution_manager_id_batch in execution_manager_ids.chunks(UPDATE_BATCH_SIZE) { - let placeholders = std::iter::repeat_n("?", execution_manager_id_batch.len()) + let mut dead_ids: Vec = Vec::with_capacity(candidate_ids.len()); + for candidate_batch in candidate_ids.chunks(UPDATE_BATCH_SIZE) { + let placeholders = std::iter::repeat_n("?", candidate_batch.len()) + .collect::>() + .join(","); + let select_for_update_query = format!( + "SELECT `id` FROM `{table}` FORCE INDEX (PRIMARY) WHERE `id` IN ({placeholders}) \ + AND `state` = '{alive_state}' AND `last_heartbeat_at` < CURRENT_TIMESTAMP - \ + INTERVAL ? SECOND ORDER BY `id` FOR UPDATE;", + table = EXECUTION_MANAGERS_TABLE_NAME, + alive_state = ExecutionManagerState::Alive.as_str(), + ); + let mut select_query = + sqlx::query_scalar::<_, ExecutionManagerId>(&select_for_update_query); + for execution_manager_id in candidate_batch { + select_query = select_query.bind(execution_manager_id); + } + let confirmed_ids: Vec = select_query + .bind(stale_after_sec) + .fetch_all(&mut *tx) + .await?; + + if confirmed_ids.is_empty() { + continue; + } + + let placeholders = std::iter::repeat_n("?", confirmed_ids.len()) .collect::>() .join(","); let update_query = format!( @@ -580,15 +630,17 @@ impl ExecutionManagerLivenessManagement for MariaDbStorageConnector { `death_confirmed_at` = CURRENT_TIMESTAMP WHERE `id` IN ({placeholders})", dead_state = ExecutionManagerState::Dead.as_str(), ); - let mut query = sqlx::query(&update_query); - for execution_manager_id in execution_manager_id_batch { - query = query.bind(execution_manager_id); + let mut update_query = sqlx::query(&update_query); + for execution_manager_id in &confirmed_ids { + update_query = update_query.bind(execution_manager_id); } - query.execute(&mut *tx).await?; + update_query.execute(&mut *tx).await?; + + dead_ids.extend(confirmed_ids); } tx.commit().await?; - Ok(execution_manager_ids) + Ok(dead_ids) } } From 4e0d7aa4c7865df97c13fa8cef92a7d94feb7442 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 7 Jul 2026 17:32:23 -0400 Subject: [PATCH 2/5] Variable renaming. --- components/spider-storage/src/db/mariadb.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index a70af2731..4a21143f1 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -369,7 +369,7 @@ impl InternalJobOrchestration for MariaDbStorageConnector { let placeholders = std::iter::repeat_n("?", candidate_ids.len()) .collect::>() .join(","); - let select_for_update_query = format!( + let select_for_update_stmt = format!( "SELECT `id` FROM `{table}` FORCE INDEX (PRIMARY) WHERE `id` IN ({placeholders}) \ AND `state` IN ('{succeeded_state}','{failed_state}','{cancelled_state}') AND \ `ended_at` < NOW() - INTERVAL ? SECOND ORDER BY `id` FOR UPDATE;", @@ -378,7 +378,7 @@ impl InternalJobOrchestration for MariaDbStorageConnector { failed_state = JobState::Failed.as_str(), cancelled_state = JobState::Cancelled.as_str(), ); - let mut select_query = sqlx::query_scalar::<_, JobId>(&select_for_update_query); + let mut select_query = sqlx::query_scalar::<_, JobId>(&select_for_update_stmt); for job_id in &candidate_ids { select_query = select_query.bind(job_id); } @@ -391,9 +391,9 @@ impl InternalJobOrchestration for MariaDbStorageConnector { let placeholders = std::iter::repeat_n("?", confirmed_ids.len()) .collect::>() .join(","); - let delete_query = - format!("DELETE FROM `{JOBS_TABLE_NAME}` WHERE `id` IN ({placeholders})"); - let mut delete_query = sqlx::query(&delete_query); + let delete_stmt = + format!("DELETE FROM `{JOBS_TABLE_NAME}` WHERE `id` IN ({placeholders});"); + let mut delete_query = sqlx::query(&delete_stmt); for job_id in &confirmed_ids { delete_query = delete_query.bind(job_id); } @@ -601,7 +601,7 @@ impl ExecutionManagerLivenessManagement for MariaDbStorageConnector { let placeholders = std::iter::repeat_n("?", candidate_batch.len()) .collect::>() .join(","); - let select_for_update_query = format!( + let select_for_update_stmt = format!( "SELECT `id` FROM `{table}` FORCE INDEX (PRIMARY) WHERE `id` IN ({placeholders}) \ AND `state` = '{alive_state}' AND `last_heartbeat_at` < CURRENT_TIMESTAMP - \ INTERVAL ? SECOND ORDER BY `id` FOR UPDATE;", @@ -609,7 +609,7 @@ impl ExecutionManagerLivenessManagement for MariaDbStorageConnector { alive_state = ExecutionManagerState::Alive.as_str(), ); let mut select_query = - sqlx::query_scalar::<_, ExecutionManagerId>(&select_for_update_query); + sqlx::query_scalar::<_, ExecutionManagerId>(&select_for_update_stmt); for execution_manager_id in candidate_batch { select_query = select_query.bind(execution_manager_id); } @@ -625,12 +625,12 @@ impl ExecutionManagerLivenessManagement for MariaDbStorageConnector { let placeholders = std::iter::repeat_n("?", confirmed_ids.len()) .collect::>() .join(","); - let update_query = format!( + let update_stmt = format!( "UPDATE `{EXECUTION_MANAGERS_TABLE_NAME}` SET `state` = '{dead_state}', \ - `death_confirmed_at` = CURRENT_TIMESTAMP WHERE `id` IN ({placeholders})", + `death_confirmed_at` = CURRENT_TIMESTAMP WHERE `id` IN ({placeholders});", dead_state = ExecutionManagerState::Dead.as_str(), ); - let mut update_query = sqlx::query(&update_query); + let mut update_query = sqlx::query(&update_stmt); for execution_manager_id in &confirmed_ids { update_query = update_query.bind(execution_manager_id); } From 391db346abe07074edc86aaa574c60ba521f270e Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Tue, 7 Jul 2026 18:45:21 -0400 Subject: [PATCH 3/5] Do not check for job state change --- components/spider-storage/src/db/mariadb.rs | 44 +++++++-------------- 1 file changed, 15 insertions(+), 29 deletions(-) diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index 4a21143f1..fcc9287b7 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -366,42 +366,28 @@ impl InternalJobOrchestration for MariaDbStorageConnector { break; } - let placeholders = std::iter::repeat_n("?", candidate_ids.len()) + let candidate_count = candidate_ids.len(); + let placeholders = std::iter::repeat_n("?", candidate_count) .collect::>() .join(","); - let select_for_update_stmt = format!( - "SELECT `id` FROM `{table}` FORCE INDEX (PRIMARY) WHERE `id` IN ({placeholders}) \ - AND `state` IN ('{succeeded_state}','{failed_state}','{cancelled_state}') AND \ - `ended_at` < NOW() - INTERVAL ? SECOND ORDER BY `id` FOR UPDATE;", - table = JOBS_TABLE_NAME, - succeeded_state = JobState::Succeeded.as_str(), - failed_state = JobState::Failed.as_str(), - cancelled_state = JobState::Cancelled.as_str(), - ); - let mut select_query = sqlx::query_scalar::<_, JobId>(&select_for_update_stmt); + let delete_stmt = + format!("DELETE FROM `{JOBS_TABLE_NAME}` WHERE `id` IN ({placeholders});"); + let mut delete_query = sqlx::query(&delete_stmt); for job_id in &candidate_ids { - select_query = select_query.bind(job_id); + delete_query = delete_query.bind(job_id); } - let confirmed_ids: Vec = select_query - .bind(expire_after_sec) - .fetch_all(&mut *tx) - .await?; + let rows_affected = delete_query.execute(&mut *tx).await?.rows_affected(); - if !confirmed_ids.is_empty() { - let placeholders = std::iter::repeat_n("?", confirmed_ids.len()) - .collect::>() - .join(","); - let delete_stmt = - format!("DELETE FROM `{JOBS_TABLE_NAME}` WHERE `id` IN ({placeholders});"); - let mut delete_query = sqlx::query(&delete_stmt); - for job_id in &confirmed_ids { - delete_query = delete_query.bind(job_id); - } - delete_query.execute(&mut *tx).await?; - deleted_job_ids.extend(confirmed_ids); + if rows_affected != candidate_count as u64 { + return Err(DbError::CorruptedDbState(format!( + "expected to delete {candidate_count} rows but only {rows_affected} rows were \ + deleted" + ))); } - if candidate_ids.len() < DELETE_BATCH_SIZE { + deleted_job_ids.extend(candidate_ids); + + if candidate_count < DELETE_BATCH_SIZE { break; } } From 184d13d44f4e4cef2be4698bed2f32fabfe809ef Mon Sep 17 00:00:00 2001 From: Sitao Wang Date: Tue, 7 Jul 2026 18:50:14 -0400 Subject: [PATCH 4/5] Check for state change --- components/spider-storage/src/db/mariadb.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index fcc9287b7..1787ff375 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -370,8 +370,14 @@ impl InternalJobOrchestration for MariaDbStorageConnector { let placeholders = std::iter::repeat_n("?", candidate_count) .collect::>() .join(","); - let delete_stmt = - format!("DELETE FROM `{JOBS_TABLE_NAME}` WHERE `id` IN ({placeholders});"); + let delete_stmt = format!( + "DELETE FROM `{table}` WHERE `id` IN ({placeholders}) AND `state` IN \ + ('{succeeded_state}','{failed_state}','{cancelled_state}');", + table = JOBS_TABLE_NAME, + succeeded_state = JobState::Succeeded.as_str(), + failed_state = JobState::Failed.as_str(), + cancelled_state = JobState::Cancelled.as_str(), + ); let mut delete_query = sqlx::query(&delete_stmt); for job_id in &candidate_ids { delete_query = delete_query.bind(job_id); @@ -380,7 +386,7 @@ impl InternalJobOrchestration for MariaDbStorageConnector { if rows_affected != candidate_count as u64 { return Err(DbError::CorruptedDbState(format!( - "expected to delete {candidate_count} rows but only {rows_affected} rows were \ + "expected to delete {candidate_count} rows but only {rows_affected} rows \ deleted" ))); } From c0c0911d1aaeadfaf6b1f2e16721b040d6e3ef24 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Wed, 8 Jul 2026 01:27:24 -0400 Subject: [PATCH 5/5] Add unit tests; Add . --- components/spider-storage/src/db/mariadb.rs | 201 ++++++++++++------ .../spider-storage/tests/mariadb_test.rs | 62 +++++- 2 files changed, 200 insertions(+), 63 deletions(-) diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index 1787ff375..513516681 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -12,7 +12,7 @@ use spider_core::{ }, }; use spider_derive::MySqlEnum; -use sqlx::{MySqlPool, mysql::MySqlDatabaseError}; +use sqlx::{Connection, MySqlPool, mysql::MySqlDatabaseError}; use crate::{ config::DatabaseConfig, @@ -573,66 +573,10 @@ impl ExecutionManagerLivenessManagement for MariaDbStorageConnector { &self, stale_after_sec: u64, ) -> Result, DbError> { - const UPDATE_BATCH_SIZE: usize = 1000; - - const SELECT_CANDIDATES_QUERY: &str = formatcp!( - "SELECT `id` FROM `{table}` WHERE `state` = '{alive_state}' AND `last_heartbeat_at` < \ - CURRENT_TIMESTAMP - INTERVAL ? SECOND ORDER BY `id`;", - table = EXECUTION_MANAGERS_TABLE_NAME, - alive_state = ExecutionManagerState::Alive.as_str(), - ); - - let mut tx = self.pool.begin().await?; - let candidate_ids: Vec = sqlx::query_scalar(SELECT_CANDIDATES_QUERY) - .bind(stale_after_sec) - .fetch_all(&mut *tx) - .await?; - - let mut dead_ids: Vec = Vec::with_capacity(candidate_ids.len()); - for candidate_batch in candidate_ids.chunks(UPDATE_BATCH_SIZE) { - let placeholders = std::iter::repeat_n("?", candidate_batch.len()) - .collect::>() - .join(","); - let select_for_update_stmt = format!( - "SELECT `id` FROM `{table}` FORCE INDEX (PRIMARY) WHERE `id` IN ({placeholders}) \ - AND `state` = '{alive_state}' AND `last_heartbeat_at` < CURRENT_TIMESTAMP - \ - INTERVAL ? SECOND ORDER BY `id` FOR UPDATE;", - table = EXECUTION_MANAGERS_TABLE_NAME, - alive_state = ExecutionManagerState::Alive.as_str(), - ); - let mut select_query = - sqlx::query_scalar::<_, ExecutionManagerId>(&select_for_update_stmt); - for execution_manager_id in candidate_batch { - select_query = select_query.bind(execution_manager_id); - } - let confirmed_ids: Vec = select_query - .bind(stale_after_sec) - .fetch_all(&mut *tx) - .await?; - - if confirmed_ids.is_empty() { - continue; - } - - let placeholders = std::iter::repeat_n("?", confirmed_ids.len()) - .collect::>() - .join(","); - let update_stmt = format!( - "UPDATE `{EXECUTION_MANAGERS_TABLE_NAME}` SET `state` = '{dead_state}', \ - `death_confirmed_at` = CURRENT_TIMESTAMP WHERE `id` IN ({placeholders});", - dead_state = ExecutionManagerState::Dead.as_str(), - ); - let mut update_query = sqlx::query(&update_stmt); - for execution_manager_id in &confirmed_ids { - update_query = update_query.bind(execution_manager_id); - } - update_query.execute(&mut *tx).await?; - - dead_ids.extend(confirmed_ids); - } - - tx.commit().await?; - Ok(dead_ids) + run_read_committed_tx(self.pool.clone(), async move |connection| { + get_dead_execution_managers(connection, stale_after_sec).await + }) + .await } } @@ -956,3 +900,138 @@ async fn transition_job_state( Ok(()) } + +/// Runs `tx` on a freshly acquired pooled connection whose next transaction uses the `READ +/// COMMITTED` isolation level. +/// +/// A `SET TRANSACTION ISOLATION LEVEL READ COMMITTED` statement is issued on the connection before +/// `tx` runs. Because it omits `SESSION`/`GLOBAL`, it applies only to the next transaction started +/// on that connection. `tx` is expected to begin exactly one transaction to consume the setting, +/// and to commit or roll it back itself. +/// +/// If `tx` returns an error, the connection is detached from the pool and closed rather than being +/// released back into it. This ensures a failed attempt can never hand a later, unrelated borrower +/// a connection still carrying the pending isolation change (or a half-open transaction). +/// +/// # Type Parameters +/// +/// * `ReturnType` - The return type of `tx`. +/// * `TransactionType` - The type of `tx`, which is an async function that takes a mutable +/// reference to the connection. +/// +/// # Returns +/// +/// The value returned by `tx` on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards [`sqlx::Pool::acquire`]'s return values on failure. +/// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. +/// * Forwards `tx`'s return values on failure. +async fn run_read_committed_tx( + pool: MySqlPool, + tx: TransactionType, +) -> Result +where + for<'connection_lifetime> TransactionType: + AsyncFnOnce(&'connection_lifetime mut sqlx::MySqlConnection) -> Result, +{ + const SET_READ_COMMITTED: &str = "SET TRANSACTION ISOLATION LEVEL READ COMMITTED"; + let mut conn = pool.acquire().await?; + sqlx::query(SET_READ_COMMITTED).execute(&mut *conn).await?; + let result = tx(&mut *conn).await; + if result.is_err() { + let _ = conn.detach().close().await; + } + result +} + +/// Marks stale execution managers dead and returns their IDs, in a transaction run on `conn`. +/// +/// # Note +/// +/// It is assumed that `conn` must be set to run its next transaction at the `READ COMMITTED` +/// isolation level. Under the default `REPEATABLE REA`, the confirming `SELECT ... FOR UPDATE` can +/// observe a row changed by a concurrent +/// [`ExecutionManagerLivenessManagement::update_execution_manager_heartbeat`] since the +/// transaction's read view was established and failed with a "record has changed" error (error code +/// 1020). +/// +/// # Returns +/// +/// A vector of dead execution manager IDs (removed from the table) on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Forwards [`sqlx::Connection::begin`]'s return values on failure. +/// * Forwards [`sqlx::query::QueryScalar::fetch_all`]'s return values on failure. +/// * Forwards [`sqlx::query::Query::execute`]'s return values on failure. +/// * Forwards [`sqlx::Transaction::commit`]'s return values on failure. +async fn get_dead_execution_managers( + conn: &mut sqlx::MySqlConnection, + stale_after_sec: u64, +) -> Result, DbError> { + const UPDATE_BATCH_SIZE: usize = 1000; + + const SELECT_CANDIDATES_QUERY: &str = formatcp!( + "SELECT `id` FROM `{table}` WHERE `state` = '{alive_state}' AND `last_heartbeat_at` < \ + CURRENT_TIMESTAMP - INTERVAL ? SECOND ORDER BY `id`;", + table = EXECUTION_MANAGERS_TABLE_NAME, + alive_state = ExecutionManagerState::Alive.as_str(), + ); + + let mut tx = Connection::begin(conn).await?; + let candidate_ids: Vec = sqlx::query_scalar(SELECT_CANDIDATES_QUERY) + .bind(stale_after_sec) + .fetch_all(&mut *tx) + .await?; + + let mut dead_ids: Vec = Vec::with_capacity(candidate_ids.len()); + for candidate_batch in candidate_ids.chunks(UPDATE_BATCH_SIZE) { + let placeholders = std::iter::repeat_n("?", candidate_batch.len()) + .collect::>() + .join(","); + let select_for_update_stmt = format!( + "SELECT `id` FROM `{table}` FORCE INDEX (PRIMARY) WHERE `id` IN ({placeholders}) AND \ + `state` = '{alive_state}' AND `last_heartbeat_at` < CURRENT_TIMESTAMP - INTERVAL ? \ + SECOND ORDER BY `id` FOR UPDATE;", + table = EXECUTION_MANAGERS_TABLE_NAME, + alive_state = ExecutionManagerState::Alive.as_str(), + ); + let mut select_query = sqlx::query_scalar::<_, ExecutionManagerId>(&select_for_update_stmt); + for execution_manager_id in candidate_batch { + select_query = select_query.bind(execution_manager_id); + } + let confirmed_ids: Vec = select_query + .bind(stale_after_sec) + .fetch_all(&mut *tx) + .await?; + + if confirmed_ids.is_empty() { + continue; + } + + let placeholders = std::iter::repeat_n("?", confirmed_ids.len()) + .collect::>() + .join(","); + let update_stmt = format!( + "UPDATE `{EXECUTION_MANAGERS_TABLE_NAME}` SET `state` = '{dead_state}', \ + `death_confirmed_at` = CURRENT_TIMESTAMP WHERE `id` IN ({placeholders});", + dead_state = ExecutionManagerState::Dead.as_str(), + ); + let mut update_query = sqlx::query(&update_stmt); + for execution_manager_id in &confirmed_ids { + update_query = update_query.bind(execution_manager_id); + } + update_query.execute(&mut *tx).await?; + + dead_ids.extend(confirmed_ids); + } + + tx.commit().await?; + Ok(dead_ids) +} diff --git a/components/spider-storage/tests/mariadb_test.rs b/components/spider-storage/tests/mariadb_test.rs index 91e7bfe1e..ff7eff582 100644 --- a/components/spider-storage/tests/mariadb_test.rs +++ b/components/spider-storage/tests/mariadb_test.rs @@ -20,9 +20,10 @@ use spider_storage::db::{ SchedulerRegistrationManagement, SessionManagement, }; +use tokio::task::JoinSet; use super::{ - mariadb_infra::{create_mariadb_connector, create_test_resource_group}, + mariadb_infra::{create_mariadb_config, create_mariadb_connector, create_test_resource_group}, task_graph_builder::{SubmittedTaskGraph, build_flat_task_graph, create_validated_submission}, }; @@ -953,7 +954,7 @@ async fn test_get_dead_execution_managers_atomic() { .expect("first get_dead_execution_managers should succeed"); assert!( dead_first.contains(&em_id), - "expected em_id in first dead list, got {dead_first:?}" + "expected {em_id} in first dead list, got {dead_first:?}" ); // Second call should not return the same EM again. @@ -997,6 +998,63 @@ async fn test_get_dead_execution_managers_multiple() { } } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires MariaDB"] +#[serial_test::file_serial] +async fn test_liveness_operations_no_deadlock_under_concurrency() { + // Regression test for . + const TEST_NUM_CONCURRENT_EMS: usize = 16; + const TEST_NUM_DEADLOCK_ITERATIONS: usize = 10; + const TEST_STALE_AFTER_SEC: u64 = 1; + + let mut config = create_mariadb_config(); + config.max_connections = + u32::try_from(TEST_NUM_CONCURRENT_EMS).expect("EM count should fit in u32") + 4; + let storage = MariaDbStorageConnector::connect(&config) + .await + .expect("connect should succeed"); + let cutoff_with_margin = Duration::from_secs(TEST_STALE_AFTER_SEC) + Duration::from_millis(100); + + for iteration in 0..TEST_NUM_DEADLOCK_ITERATIONS { + let mut em_ids = Vec::with_capacity(TEST_NUM_CONCURRENT_EMS); + for _ in 0..TEST_NUM_CONCURRENT_EMS { + em_ids.push(register_test_em(&storage).await); + } + + tokio::time::sleep(cutoff_with_margin).await; + + // Race every heartbeat against the single get_dead on one join set. The get_dead result is + // mapped to `()` so all tasks share the `Result<(), DbError>` type and a single outcome + // check: a deadlock or any other unexpected failure surfaces as an `Err` other than + // `ExecutionManagerAlreadyDead`. + let mut join_set: JoinSet> = JoinSet::new(); + for em_id in em_ids { + let storage = storage.clone(); + join_set.spawn(async move { storage.update_execution_manager_heartbeat(em_id).await }); + } + + let storage = storage.clone(); + join_set.spawn(async move { + storage + .get_dead_execution_managers(TEST_STALE_AFTER_SEC) + .await + .map(|_| ()) + }); + + while let Some(joined) = join_set.join_next().await { + let result = joined.expect("task should not panic"); + assert!( + matches!( + result, + Ok(()) | Err(DbError::ExecutionManagerAlreadyDead(_)) + ), + "iteration {iteration}: liveness operation should succeed or find the EM already \ + dead, got {result:?}" + ); + } + } +} + #[tokio::test] #[ignore = "requires MariaDB"] #[serial_test::file_serial]