From 7c80a12d751483bccbcbad51df86ae63823e674d Mon Sep 17 00:00:00 2001 From: John Coffey Date: Thu, 24 Sep 2026 08:19:05 -0700 Subject: [PATCH] Task manager: release task locks on stop, recheck claims held elsewhere A cluster rehearsal (PostgreSQL + NATS) left index tasks pending well past the one-hour task lock after the node that claimed them was stopped or killed. The exact cause there isn't confirmed; this closes every path found in the task manager that stretches a takeover past the lock, or keeps a task claimed without running it: - A graceful stop never released the locks it held, so every task the node had claimed stayed blocked for an hour. The server now tracks the locks it holds (common::ipc::TaskLocks) and, once the shutdown signal arrives, stops claiming and releases them before exiting. - A node that failed to claim a task (another node held it) set its own local hold for a full lock lifetime from that scan. If the holder claimed it just after the scan began, or ran on a clock ahead, that hold ran out a moment before the lock did and was set for another hour: two hours in all. Such claims are now tried again every five minutes (a twelfth of the lock lifetime), and the task manager wakes up for them: before, a node without a coordinator could sleep up to five minutes past the recheck, or until something else woke it. - A worker that panicked took its task type down on that node for good, while the scan kept claiming that type's tasks and failing to hand them over, re-taking each lock as it expired and so starving every other node of them. Each batch now runs on a task of its own; a panic is logged, the batch's locks are released and the worker carries on. A failed hand-over releases the lock too. - A claimed task the worker couldn't read, or found gone, kept its lock for the hour. It is released. - An IndexDocument task for a file (not indexed) returned no result, which shifted every later result in the batch onto the wrong task in update_tasks. It returns Ignored. Nothing queues such a task today. The lock lifetime stays one hour; it now lives per server so the tests can shorten it. store::task_locks::task_lock_tests plays a second node by writing its locks straight into the in-memory store: tasks it claimed and abandoned run here once its locks expire, including locks that outlive this node's view of them, and a graceful stop hands this node's locks back at once and claims nothing more. It passes on RocksDB, SQLite and PostgreSQL. With the old recheck it fails. --- crates/common/src/ipc.rs | 54 ++++ crates/common/src/lib.rs | 2 + crates/common/src/manager/boot.rs | 1 + crates/main/src/main.rs | 4 + crates/services/src/task_manager/index.rs | 10 +- crates/services/src/task_manager/lock.rs | 36 ++- crates/services/src/task_manager/manager.rs | 337 ++++++++++++-------- crates/services/src/task_manager/mod.rs | 4 +- tests/src/store/mod.rs | 1 + tests/src/store/task_locks.rs | 184 +++++++++++ 10 files changed, 489 insertions(+), 144 deletions(-) create mode 100644 tests/src/store/task_locks.rs diff --git a/crates/common/src/ipc.rs b/crates/common/src/ipc.rs index 7e0c87b..a7fab1c 100644 --- a/crates/common/src/ipc.rs +++ b/crates/common/src/ipc.rs @@ -335,3 +335,57 @@ impl EmailPush { } } } + +/// inbuxa: the task locks this node holds, so a graceful stop can hand them +/// back instead of leaving the tasks blocked until the locks expire. +pub struct TaskLocks { + held: parking_lot::Mutex>, + stopping: AtomicBool, + expiry: std::sync::atomic::AtomicU64, +} + +impl TaskLocks { + /// How long a task lock lasts, in seconds, unless it is released first. + pub const DEFAULT_EXPIRY: u64 = 60 * 60; + + pub fn is_stopping(&self) -> bool { + self.stopping.load(Ordering::Acquire) + } + + /// Stops new claims and returns the ids of every lock still held. + pub fn stop(&self) -> Vec { + self.stopping.store(true, Ordering::Release); + self.held.lock().drain().collect() + } + + pub fn insert(&self, id: u64) { + self.held.lock().insert(id); + } + + pub fn remove(&self, id: u64) { + self.held.lock().remove(&id); + } + + pub fn held(&self) -> usize { + self.held.lock().len() + } + + pub fn expiry(&self) -> u64 { + self.expiry.load(Ordering::Relaxed) + } + + /// Changes the lock lifetime; the tests shorten it. + pub fn set_expiry(&self, seconds: u64) { + self.expiry.store(seconds.max(1), Ordering::Relaxed); + } +} + +impl Default for TaskLocks { + fn default() -> Self { + Self { + held: Default::default(), + stopping: AtomicBool::new(false), + expiry: std::sync::atomic::AtomicU64::new(Self::DEFAULT_EXPIRY), + } + } +} diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index c0b4bdd..d797b14 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -279,6 +279,8 @@ pub struct HttpAuthCache { pub struct Ipc { pub push_tx: mpsc::Sender, pub task_tx: Arc, + // inbuxa: task locks held by this node, released on a graceful stop + pub task_locks: Arc, pub queue_tx: mpsc::Sender, pub report_tx: mpsc::Sender, pub broadcast_tx: Option>, diff --git a/crates/common/src/manager/boot.rs b/crates/common/src/manager/boot.rs index 9f2f85a..b4e7b13 100644 --- a/crates/common/src/manager/boot.rs +++ b/crates/common/src/manager/boot.rs @@ -297,6 +297,7 @@ pub fn build_ipc(has_pubsub: bool) -> (Ipc, IpcReceivers) { report_tx, broadcast_tx: has_pubsub.then_some(broadcast_tx), task_tx: Arc::new(Notify::new()), + task_locks: Arc::new(crate::ipc::TaskLocks::default()), train_task_controller: Arc::new(TrainTaskController::default()), }, IpcReceivers { diff --git a/crates/main/src/main.rs b/crates/main/src/main.rs index 646f810..c61d5c3 100644 --- a/crates/main/src/main.rs +++ b/crates/main/src/main.rs @@ -109,6 +109,10 @@ async fn main() -> std::io::Result<()> { // Wait for shutdown signal wait_for_shutdown().await; + // inbuxa: hand back the task locks this node holds, so other nodes can + // run those tasks now rather than when the locks expire + services::task_manager::lock::release_task_locks(&inner.build_server()).await; + // Shutdown collector Collector::shutdown(); diff --git a/crates/services/src/task_manager/index.rs b/crates/services/src/task_manager/index.rs index 95296d1..07ce970 100644 --- a/crates/services/src/task_manager/index.rs +++ b/crates/services/src/task_manager/index.rs @@ -91,7 +91,15 @@ impl SearchIndexTask for Server { build_contact_document(self, account_id, document_id).await } IndexDocumentType::File => { - // File indexing not implemented yet + // File indexing not implemented yet. inbuxa: still + // one result per task: update_tasks pairs them by + // position, and a missing one shifts every result + // after it onto the wrong task + results.push(IndexTaskResult { + task_type: TaskType::Insert, + index: task.document_type, + result: TaskResult::Ignored, + }); continue; } }; diff --git a/crates/services/src/task_manager/lock.rs b/crates/services/src/task_manager/lock.rs index cb11a2e..861d5b9 100644 --- a/crates/services/src/task_manager/lock.rs +++ b/crates/services/src/task_manager/lock.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use crate::task_manager::*; @@ -13,13 +15,21 @@ pub trait TaskLockManager: Sync + Send { impl TaskLockManager for Server { async fn try_lock_task(&self, id: u64) -> bool { + // inbuxa: a node that is stopping claims nothing new + let locks = &self.inner.ipc.task_locks; + if locks.is_stopping() { + return false; + } + match self .in_memory_store() - .try_lock(KV_LOCK_TASK, &id.to_be_bytes(), DEFAULT_LOCK_EXPIRY) + .try_lock(KV_LOCK_TASK, &id.to_be_bytes(), locks.expiry()) .await { Ok(result) => { - if !result { + if result { + locks.insert(id); + } else { trc::event!( TaskManager(TaskManagerEvent::TaskLocked), Id = id, @@ -48,5 +58,27 @@ impl TaskLockManager for Server { .caused_by(trc::location!()) ); } + self.inner.ipc.task_locks.remove(id); } } + +/// inbuxa: on a graceful stop, stops claiming tasks and releases every task +/// lock this node holds, so the rest of the cluster can pick the tasks up at +/// once instead of after the lock expires. Returns how many were released. +pub async fn release_task_locks(server: &Server) -> usize { + let ids = server.inner.ipc.task_locks.stop(); + for id in &ids { + if let Err(err) = server + .in_memory_store() + .remove_lock(KV_LOCK_TASK, &id.to_be_bytes()) + .await + { + trc::error!( + err.details("Failed to release task lock on shutdown") + .ctx(trc::Key::Id, *id) + .caused_by(trc::location!()) + ); + } + } + ids.len() +} diff --git a/crates/services/src/task_manager/manager.rs b/crates/services/src/task_manager/manager.rs index 180aeac..4d32b9a 100644 --- a/crates/services/src/task_manager/manager.rs +++ b/crates/services/src/task_manager/manager.rs @@ -2,6 +2,8 @@ * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC * * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL + * + * Modified by Coffey Labs in 2026 for INBUXA. */ use crate::task_manager::acme::AcmeTask; @@ -18,7 +20,7 @@ use crate::task_manager::report::{self, SubmitReportTask}; use crate::task_manager::restore_item::RestoreItemTask; use crate::task_manager::spam_classifier::SpamFilterMaintenanceTask; use crate::task_manager::{ - DEFAULT_LOCK_EXPIRY, Locked, QUEUE_REFRESH_INTERVAL, TaskDetails, TaskFailureType, TaskInfo, + CLAIM_RECHECK_INTERVAL, Locked, QUEUE_REFRESH_INTERVAL, TaskDetails, TaskFailureType, TaskInfo, TaskJob, TaskManagerIpc, TaskResult, }; use common::BuildServer; @@ -124,72 +126,47 @@ pub fn spawn_task_manager(inner: Arc) { let server = inner.build_server(); let batch_size = server.core.email.index_batch_size; let mut batch = Vec::with_capacity(batch_size); - match server - .store() - .get_value::(ValueKey::from(ValueClass::TaskQueue( - TaskQueueClass::Task { id: job.id }, - ))) - .await - { - Ok(Some(task)) => { - batch.push(TaskDetails { task, info: job }); - } - Ok(None) => { - trc::event!( - TaskManager(TaskManagerEvent::TaskIgnored), - Id = job.id, - Reason = "Task not found in store, likely already processed.", - ); - } - Err(err) => { - trc::error!( - err.id(job.id) - .details("Failed to retrieve task details.") - .caused_by(trc::location!()) - ); - } + if let Some(task) = fetch_task(&server, job).await { + batch.push(task); } while batch.len() < batch_size { match rx.try_recv() { Ok(job) => { - match server - .store() - .get_value::(ValueKey::from(ValueClass::TaskQueue( - TaskQueueClass::Task { id: job.id }, - ))) - .await - { - Ok(Some(task)) => { - batch.push(TaskDetails { task, info: job }); - } - Ok(None) => { - trc::event!( - TaskManager(TaskManagerEvent::TaskIgnored), - Id = job.id, - Reason = "Task not found in store, likely already processed.", - ); - } - Err(err) => { - trc::error!( - err.id(job.id) - .details("Failed to retrieve task details.") - .caused_by(trc::location!()) - ); - } + if let Some(task) = fetch_task(&server, job).await { + batch.push(task); } } Err(_) => break, } } - // Dispatch + // Dispatch. inbuxa: on a task of its own, so a panic + // releases the batch's locks and leaves this worker + // running; a dead worker would keep claiming tasks it + // can never run let mut refresh_queue = false; - let results = server.index(&batch).await.into_iter().map(|r| { - refresh_queue |= r.result.is_retry(); - r.result - }); - update_tasks(&server, &mut batch, results).await; + let ids = batch.iter().map(|task| task.info.id).collect::>(); + let run = { + let server = server.clone(); + tokio::spawn(async move { + let results = server.index(&batch).await; + (batch, results) + }) + }; + match run.await { + Ok((mut batch, results)) => { + let results = results.into_iter().map(|r| { + refresh_queue |= r.result.is_retry(); + r.result + }); + update_tasks(&server, &mut batch, results).await; + } + Err(err) => { + worker_failed(&server, &ids, err).await; + refresh_queue = true; + } + } if refresh_queue || rx.is_empty() { server.notify_task_queue(); @@ -203,83 +180,31 @@ pub fn spawn_task_manager(inner: Arc) { let server = inner.build_server(); let mut refresh_queue = false; - match server - .store() - .get_value::(ValueKey::from(ValueClass::TaskQueue( - TaskQueueClass::Task { id: job.id }, - ))) - .await - { - Ok(Some(task)) => { - let result = match &task { - Task::CalendarAlarmEmail(task) => { - server.send_email_alarm(task, server_instance.clone()).await - } - Task::CalendarAlarmNotification(task) => { - server.send_display_alarm(task).await - } - Task::CalendarItipMessage(task) => { - server.send_imip(task, server_instance.clone()).await - } - Task::MergeThreads(task) => server.merge_threads(task).await, - Task::DmarcReport(task) => { - server - .submit_report(report::ReportId::Dmarc(task.report_id.id())) - .await - } - Task::TlsReport(task) => { - server - .submit_report(report::ReportId::Tls(task.report_id.id())) - .await - } - Task::RestoreArchivedItem(task) => server.restore_item(task).await, - Task::DestroyAccount(task) => server.destroy_account(task).await, - Task::AccountMaintenance(task) => { - server.account_maintenance(task).await - } - Task::TenantMaintenance(task) => { - server.tenant_maintenance(task).await - } - Task::StoreMaintenance(task) => { - server.store_maintenance(task).await - } - Task::SpamFilterMaintenance(task) => { - Box::pin(server.spam_filter_maintenance(task)).await - } - Task::AcmeRenewal(task) => server.acme_management(task).await, - Task::DkimManagement(task_dkim_rotation) => { - server.dkim_management(task_dkim_rotation).await - } - Task::DnsManagement(task_dns_management) => { - server.dns_management(task_dns_management).await - } - Task::IndexDocument(_) - | Task::UnindexDocument(_) - | Task::IndexTrace(_) => unreachable!(), - }; + if let Some(TaskDetails { task, info }) = fetch_task(&server, job).await { + // inbuxa: on a task of its own, as above + let run = { + let server = server.clone(); + let server_instance = server_instance.clone(); + tokio::spawn(async move { + let result = run_task(&server, &task, server_instance).await; + (task, result) + }) + }; + match run.await { + Ok((task, result)) => { + refresh_queue = result.is_retry(); - refresh_queue = result.is_retry(); - - update_tasks( - &server, - &mut [TaskDetails { task, info: job }], - vec![result], - ) - .await; - } - Ok(None) => { - trc::event!( - TaskManager(TaskManagerEvent::TaskIgnored), - Id = job.id, - Reason = "Task not found in store, likely already processed.", - ); - } - Err(err) => { - trc::error!( - err.id(job.id) - .details("Failed to retrieve task details.") - .caused_by(trc::location!()) - ); + update_tasks( + &server, + &mut [TaskDetails { task, info }], + vec![result], + ) + .await; + } + Err(err) => { + worker_failed(&server, &[info.id], err).await; + refresh_queue = true; + } } } @@ -318,6 +243,13 @@ pub(crate) trait TaskQueueManager: Sync + Send { impl TaskQueueManager for Server { async fn process_tasks(&self, ipc: &mut TaskManagerIpc) -> Duration { + // inbuxa: a node that is stopping has released its locks and claims + // nothing new + let task_locks = &self.inner.ipc.task_locks; + if task_locks.is_stopping() { + return Duration::from_secs(QUEUE_REFRESH_INTERVAL); + } + let lock_expiry = task_locks.expiry(); let now_timestamp = now(); let from_key = ValueKey:: { account_id: 0, @@ -393,9 +325,7 @@ impl TaskQueueManager for Server { let locked = entry.get_mut(); if locked.expires <= now || locked.due < task_due { locked.expires = Instant::now() - + std::time::Duration::from_secs( - DEFAULT_LOCK_EXPIRY + 1, - ); + + std::time::Duration::from_secs(lock_expiry + 1); locked.due = task_due; tasks.push(( TaskJob { @@ -411,9 +341,7 @@ impl TaskQueueManager for Server { Entry::Vacant(entry) => { entry.insert(Locked { expires: Instant::now() - + std::time::Duration::from_secs( - DEFAULT_LOCK_EXPIRY + 1, - ), + + std::time::Duration::from_secs(lock_expiry + 1), due: task_due, revision: ipc.revision, }); @@ -464,12 +392,26 @@ impl TaskQueueManager for Server { let tx = &ipc.txs[task_type_idx as usize]; if tx.capacity() > 0 { - if self.try_lock_task(task_job.id).await && tx.send(task_job).await.is_err() { + let id = task_job.id; + if !self.try_lock_task(id).await { + // inbuxa: another node holds the task. Look again after a + // short while rather than a full lock lifetime from now: + // the holder may have claimed it after this scan began, + // or run on a clock ahead of this one, and waiting the + // whole lifetime again would leave the task stuck for + // another hour past its lock if that holder died + if let Some(locked) = ipc.locked.get_mut(&id) { + locked.expires = + Instant::now() + Duration::from_secs(claim_recheck_interval(lock_expiry)); + } + } else if tx.send(task_job).await.is_err() { trc::event!( Server(trc::ServerEvent::ThreadError), Details = "Error sending task.", CausedBy = trc::location!() ); + // inbuxa: nothing will run it here, so don't hold it + self.remove_index_lock(id).await; } } else { // If the channel is full, release the lock so it can be picked up in the next iteration @@ -481,9 +423,117 @@ impl TaskQueueManager for Server { let now = Instant::now(); ipc.locked .retain(|_, locked| locked.expires > now && locked.revision == ipc.revision); - Duration::from_secs(next_event.map_or(QUEUE_REFRESH_INTERVAL, |timestamp| { + let sleep_for = Duration::from_secs(next_event.map_or(QUEUE_REFRESH_INTERVAL, |timestamp| { timestamp.saturating_sub(store::write::now()) - })) + })); + + // inbuxa: wake up when a claim held elsewhere is due to be tried + // again, rather than only on the next task or refresh + ipc.locked + .values() + .map(|locked| locked.expires.saturating_duration_since(now)) + .min() + .map_or(sleep_for, |recheck| sleep_for.min(recheck.max(Duration::from_secs(1)))) + } +} + +async fn run_task( + server: &Server, + task: &Task, + server_instance: Arc, +) -> TaskResult { + match task { + Task::CalendarAlarmEmail(task) => { + server.send_email_alarm(task, server_instance.clone()).await + } + Task::CalendarAlarmNotification(task) => { + server.send_display_alarm(task).await + } + Task::CalendarItipMessage(task) => { + server.send_imip(task, server_instance.clone()).await + } + Task::MergeThreads(task) => server.merge_threads(task).await, + Task::DmarcReport(task) => { + server + .submit_report(report::ReportId::Dmarc(task.report_id.id())) + .await + } + Task::TlsReport(task) => { + server + .submit_report(report::ReportId::Tls(task.report_id.id())) + .await + } + Task::RestoreArchivedItem(task) => server.restore_item(task).await, + Task::DestroyAccount(task) => server.destroy_account(task).await, + Task::AccountMaintenance(task) => { + server.account_maintenance(task).await + } + Task::TenantMaintenance(task) => { + server.tenant_maintenance(task).await + } + Task::StoreMaintenance(task) => { + server.store_maintenance(task).await + } + Task::SpamFilterMaintenance(task) => { + Box::pin(server.spam_filter_maintenance(task)).await + } + Task::AcmeRenewal(task) => server.acme_management(task).await, + Task::DkimManagement(task_dkim_rotation) => { + server.dkim_management(task_dkim_rotation).await + } + Task::DnsManagement(task_dns_management) => { + server.dns_management(task_dns_management).await + } + Task::IndexDocument(_) + | Task::UnindexDocument(_) + | Task::IndexTrace(_) => unreachable!(), + } +} + +/// Reads a claimed task. When it is gone or can't be read, the claim is +/// released: inbuxa: holding it would block the task, everywhere, until +/// the lock expired. +async fn fetch_task(server: &Server, job: TaskJob) -> Option { + match server + .store() + .get_value::(ValueKey::from(ValueClass::TaskQueue(TaskQueueClass::Task { + id: job.id, + }))) + .await + { + Ok(Some(task)) => Some(TaskDetails { task, info: job }), + Ok(None) => { + trc::event!( + TaskManager(TaskManagerEvent::TaskIgnored), + Id = job.id, + Reason = "Task not found in store, likely already processed.", + ); + server.remove_index_lock(job.id).await; + None + } + Err(err) => { + trc::error!( + err.id(job.id) + .details("Failed to retrieve task details.") + .caused_by(trc::location!()) + ); + server.remove_index_lock(job.id).await; + None + } + } +} + +/// inbuxa: a task panicked: its locks are released so it runs again, here or +/// on another node, and the worker carries on. +async fn worker_failed(server: &Server, ids: &[u64], err: tokio::task::JoinError) { + trc::event!( + Server(trc::ServerEvent::ThreadError), + Details = "Task worker failed", + Reason = err.to_string(), + CausedBy = trc::location!() + ); + for id in ids { + server.remove_index_lock(*id).await; } } @@ -614,6 +664,13 @@ async fn update_tasks( } } +/// inbuxa: how long to wait before trying again to claim a task another node +/// holds: a twelfth of the lock lifetime, so five minutes for the one-hour +/// lock, never more than that and never under a second. +pub(crate) fn claim_recheck_interval(lock_expiry: u64) -> u64 { + (lock_expiry / 12).clamp(1, CLAIM_RECHECK_INTERVAL) +} + pub fn perpetual_retry_time(typ: TaskType, attempt: u64) -> Option { matches!( typ, diff --git a/crates/services/src/task_manager/mod.rs b/crates/services/src/task_manager/mod.rs index d2631f5..d8c1976 100644 --- a/crates/services/src/task_manager/mod.rs +++ b/crates/services/src/task_manager/mod.rs @@ -35,7 +35,9 @@ pub mod scheduler; pub mod spam_classifier; const QUEUE_REFRESH_INTERVAL: u64 = 60 * 5; // 5 minutes -const DEFAULT_LOCK_EXPIRY: u64 = 60 * 60; // 1 hour +// inbuxa: the lock lifetime (one hour) lives in common::ipc::TaskLocks, per +// server, so a graceful stop can release the locks and the tests can shorten it +const CLAIM_RECHECK_INTERVAL: u64 = 60 * 5; // 5 minutes pub(crate) struct TaskManagerIpc { txs: [mpsc::Sender; TaskType::COUNT], diff --git a/tests/src/store/mod.rs b/tests/src/store/mod.rs index ff4b253..e863401 100644 --- a/tests/src/store/mod.rs +++ b/tests/src/store/mod.rs @@ -21,6 +21,7 @@ pub mod replica_cluster; // inbuxa: read replicas across nodes pub mod scaleout; // inbuxa: scale-out storage #[cfg(any(feature = "postgres", feature = "mysql"))] pub mod sql_timeout; +pub mod task_locks; // inbuxa: task locks across nodes use crate::utils::server::TestServerBuilder; use std::io::Read; diff --git a/tests/src/store/task_locks.rs b/tests/src/store/task_locks.rs new file mode 100644 index 0000000..bcfacb6 --- /dev/null +++ b/tests/src/store/task_locks.rs @@ -0,0 +1,184 @@ +/* + * SPDX-FileCopyrightText: 2026 Coffey Labs + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +//! Task locks across nodes: tasks claimed by a node that then disappears +//! run elsewhere once its locks expire, and a node that stops gracefully +//! hands its locks back at once. The other node is played by writing its +//! locks straight into the shared in-memory store, as a node that claimed +//! the tasks and died leaves them. + +use crate::utils::server::TestServerBuilder; +use common::{KV_LOCK_TASK, Server}; +use registry::schema::{ + enums::IndexDocumentType, + structs::{Task, TaskIndexDocument, TaskStatus}, +}; +use services::task_manager::lock::{TaskLockManager, release_task_locks}; +use std::time::{Duration, Instant}; +use store::{ + ValueKey, + write::{BatchBuilder, TaskQueueClass, ValueClass}, +}; +use utils::snowflake::SnowflakeIdGenerator; + +// Short enough for a test, long enough that the recheck interval (a twelfth +// of it) is well below it +const LOCK_EXPIRY: u64 = 12; + +#[tokio::test(flavor = "multi_thread")] +pub async fn task_lock_tests() { + let test = TestServerBuilder::new("task_lock_tests") + .await + .build() + .await; + let server = test.server.clone(); + println!( + "Running task lock tests on {}...", + std::env::var("STORE").unwrap_or_default() + ); + server.inner.ipc.task_locks.set_expiry(LOCK_EXPIRY); + + // 1. Another node claimed the tasks and died. Its locks block them until + // they expire; then this node runs them, without waiting for anything + // else to wake it + let ids = new_task_ids(4); + for id in &ids { + assert!(foreign_lock(&server, *id, LOCK_EXPIRY).await); + } + schedule(&server, &ids).await; + let started = Instant::now(); + server.notify_task_queue(); + tokio::time::sleep(Duration::from_secs(3)).await; + assert_eq!(pending(&server, &ids).await, ids.len(), "held by the other node"); + wait_until_done(&server, &ids, Duration::from_secs(LOCK_EXPIRY + 10)).await; + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_secs(LOCK_EXPIRY - 2), + "ran before the other node's locks expired: {elapsed:?}" + ); + + // 2. The other node's locks outlive what this node expects: claimed just + // after this node looked, or by a node whose clock runs ahead. This node + // keeps checking at the recheck interval, so the tasks run soon after + // those locks expire, not a whole lock lifetime later + let held_for = LOCK_EXPIRY + LOCK_EXPIRY / 2; + let ids = new_task_ids(4); + for id in &ids { + assert!(foreign_lock(&server, *id, held_for).await); + } + schedule(&server, &ids).await; + let started = Instant::now(); + server.notify_task_queue(); + wait_until_done(&server, &ids, Duration::from_secs(held_for + 8)).await; + let elapsed = started.elapsed(); + assert!( + elapsed >= Duration::from_secs(held_for - 2), + "ran before the other node's locks expired: {elapsed:?}" + ); + + // 3. A graceful stop releases the locks this node holds: another node + // can claim those tasks at once, and this one claims nothing more + let ids = new_task_ids(3); + for id in &ids { + assert!(server.try_lock_task(*id).await, "claim {id}"); + } + assert_eq!(server.inner.ipc.task_locks.held(), ids.len()); + for id in &ids { + assert!( + !foreign_lock(&server, *id, LOCK_EXPIRY).await, + "held while this node runs" + ); + } + assert_eq!(release_task_locks(&server).await, ids.len()); + assert_eq!(server.inner.ipc.task_locks.held(), 0); + for id in &ids { + assert!( + foreign_lock(&server, *id, LOCK_EXPIRY).await, + "released on stop: {id}" + ); + } + let [id] = new_task_ids(1)[..] else { + unreachable!() + }; + assert!(!server.try_lock_task(id).await, "a stopping node claims nothing"); + + for id in ids { + let _ = server + .in_memory_store() + .remove_lock(KV_LOCK_TASK, &id.to_be_bytes()) + .await; + } + if test.is_reset() { + test.temp_dir.delete(); + } +} + +fn new_task_ids(count: usize) -> Vec { + (0..count) + .map(|_| SnowflakeIdGenerator::global_id().unwrap()) + .collect() +} + +/// The other node's claim on a task, as its task manager takes it. +async fn foreign_lock(server: &Server, id: u64, seconds: u64) -> bool { + server + .in_memory_store() + .try_lock(KV_LOCK_TASK, &id.to_be_bytes(), seconds) + .await + .unwrap() +} + +/// Unindex tasks for files that don't exist: files aren't search-indexed and +/// there is no undelete note, so running one only drops it from the queue. +async fn schedule(server: &Server, ids: &[u64]) { + let mut batch = BatchBuilder::new(); + for (n, id) in ids.iter().enumerate() { + batch.schedule_task_with_id( + *id, + Task::UnindexDocument(TaskIndexDocument { + account_id: 0u32.into(), + document_id: (u32::MAX - n as u32).into(), + document_type: IndexDocumentType::File, + status: TaskStatus::now(), + }), + ); + } + server.store().write(batch.build_all()).await.unwrap(); +} + +async fn pending(server: &Server, ids: &[u64]) -> usize { + let mut count = 0; + for id in ids { + if server + .store() + .get_value::(ValueKey::from(ValueClass::TaskQueue( + TaskQueueClass::Task { id: *id }, + ))) + .await + .unwrap() + .is_some() + { + count += 1; + } + } + count +} + +async fn wait_until_done(server: &Server, ids: &[u64], within: Duration) { + let started = Instant::now(); + loop { + let left = pending(server, ids).await; + if left == 0 { + return; + } + assert!( + started.elapsed() < within, + "{left} task(s) still pending after {:?}", + started.elapsed() + ); + tokio::time::sleep(Duration::from_millis(250)).await; + } +} -- 2.54.0