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

@@ -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<ahash::AHashSet<u64>>,
+    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<u64> {
+        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),
+        }
+    }
+}

crates/common/src/lib.rs

@@ -279,6 +279,8 @@ pub struct HttpAuthCache {
 pub struct Ipc {
     pub push_tx: mpsc::Sender<PushEvent>,
     pub task_tx: Arc<Notify>,
+    // inbuxa: task locks held by this node, released on a graceful stop
+    pub task_locks: Arc<crate::ipc::TaskLocks>,
     pub queue_tx: mpsc::Sender<QueueEvent>,
     pub report_tx: mpsc::Sender<ReportingEvent>,
     pub broadcast_tx: Option<mpsc::Sender<BroadcastEvent>>,

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 {

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();
 

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;
                         }
                     };

crates/services/src/task_manager/lock.rs

@@ -2,6 +2,8 @@
  * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
  *
  * 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()
+}

crates/services/src/task_manager/manager.rs

@@ -2,6 +2,8 @@
  * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <[email protected]>
  *
  * 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<Inner>) {
                     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::<Task>(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::<Task>(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::<Vec<_>>();
+                    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<Inner>) {
                     let server = inner.build_server();
                     let mut refresh_queue = false;
 
-                    match server
-                        .store()
-                        .get_value::<Task>(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::<ValueClass> {
             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<ServerInstance>,
+) -> 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<TaskDetails> {
+    match server
+        .store()
+        .get_value::<Task>(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<u64> {
     matches!(
         typ,

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<TaskJob>; TaskType::COUNT],