Coordinator: join the cluster when NATS comes up, report the connection

A node that started while NATS was down never got a coordinator. The
connect failed at boot, bootstrap recorded a build error and the node ran
with Coordinator::None until restarted. It had no broadcast subscriber
or publisher, so cross-node push and cache invalidation to it stayed
broken, and its healthcheck said nothing about it. Losing NATS after
startup was silent too.

- The NATS client now connects in the background
  (retry_on_initial_connect): startup never waits on NATS or fails over
  it, the node gets its coordinator, subscriber and publisher at once,
  and the client keeps trying (async-nats's backoff, at most 4 s apart)
  until NATS answers. Subscriptions made meanwhile start delivering when
  it does. A configured maxReconnects still ends the attempts.
- Three new events report the connection: cluster.coordinator-connected
  (info), cluster.coordinator-disconnected (warn: lost, closed, gave up,
  or not connected within the connection timeout at startup) and
  cluster.coordinator-error (warn: a failed attempt, reported once per
  outage rather than every retry, and server errors, slow consumers and
  lame duck mode). They are in the packaged schema, ids 644 to 646.
- GET /healthz/cluster reports the coordinator: 200
  {"coordinator":"connected"}, 503 {"coordinator":"disconnected"}, or
  200 with "none" (no coordinator) or "unknown" (a backend that doesn't
  track its connection). /healthz/live and /healthz/ready are unchanged
  on purpose: a node without its coordinator still serves mail, and
  failing those would have orchestrators restart, or pull out of
  service, every node at once whenever NATS is down.

Only NATS connects lazily; the other coordinator backends still fail at
boot as before.

cluster::coordinator::coordinator_reconnect_tests starts a node against a
NATS port with nothing behind it, checks it boots with a coordinator and
reports it disconnected, subscribes, then starts NATS on that port: the
node connects on its own and the subscription receives a message from a
second client. Stopping and restarting NATS shows disconnected, then
connected, and the same subscription keeps working.

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],

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;

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<u64> {
+    (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::<Task>(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;
+    }
+}