fix(spurctld): reclaim agent allocations the controller no longer tracks on that node

crates/spurctld/src/cluster.rs

@@ -897,6 +897,20 @@ impl ClusterManager {
         self.jobs.read().get(&job_id).map(|j| j.state)
     }
 
+    /// Next id this controller would assign. Ids at or above it were never
+    /// issued here; ids below it were not necessarily issued either.
+    pub fn peek_next_job_id(&self) -> JobId {
+        self.next_job_id.load(Ordering::Relaxed)
+    }
+
+    /// Whether `node` is part of the job's current allocation.
+    pub fn job_holds_node(&self, job_id: JobId, node: &str) -> bool {
+        self.jobs
+            .read()
+            .get(&job_id)
+            .is_some_and(|j| j.allocated_nodes.iter().any(|n| n == node))
+    }
+
     /// Get a job by ID, synthesizing an aggregate record for an array *parent*
     /// id (which has no stored job — Spur stores only per-task jobs) so
     /// `scontrol show job <array_parent>` matches Slurm instead of returning
@@ -5528,6 +5542,8 @@ impl StateMachineApply for ClusterManager {
         // (NOT config-derived like license_pool/burst_buffer) — restore them.
         *self.k0s.write() = snap.k0s;
 
+        // Must stay inside the `jobs` write guard: a reader seeing the cleared
+        // map with this watermark would treat every live job as reclaimable.
         self.next_job_id.store(next_id, Ordering::Relaxed);
 
         // Re-evaluate partition membership and NodeConfig policy

crates/spurctld/src/server.rs

@@ -228,10 +228,10 @@ impl ControllerService {
         Err(self.not_leader_status())
     }
 
-    /// Re-send a terminal-job cancel to a node still reporting it on heartbeat,
-    /// freeing an allocation whose terminal cancel never landed. Spawned, best-effort.
+    /// Re-send a cancel to a node still reporting a job the controller considers
+    /// finished, freeing an allocation whose cancel never landed. Best-effort.
     fn reclaim_stale_agent_jobs(&self, node: &str, reported: &[RunningJobStatus]) {
-        let stale = stale_reported_jobs(&self.cluster, reported);
+        let stale = stale_reported_jobs(&self.cluster, node, reported);
         if stale.is_empty() {
             return;
         }
@@ -241,13 +241,13 @@ impl ControllerService {
             for job_id in stale {
                 // Re-check: a requeue since the snapshot above would otherwise
                 // send an unguarded cancel into the job's new run.
-                if !is_still_terminal(&cluster, job_id) {
+                if !is_reclaimable(&cluster, &node, job_id) {
                     continue;
                 }
                 warn!(
                     job_id,
                     node = %node,
-                    "agent still holds a terminal job — re-sending cancel to reclaim its allocation"
+                    "agent still holds a job the controller no longer tracks — re-sending cancel to reclaim its allocation"
                 );
                 // Signal 0 = graceful release, no-op on an unknown id. Not
                 // epoch-gated — a requeue racing this send is still possible.
@@ -394,18 +394,28 @@ fn is_k0s_admin(cache: &crate::association_cache::AssociationCache, caller: &str
     caller.is_empty() || caller == "root" || cache.is_admin(caller)
 }
 
-/// Whether the controller still considers `job_id` terminal right now — guards
-/// the spawned reclaim loop against a requeue landing after its stale snapshot.
-fn is_still_terminal(cluster: &ClusterManager, job_id: u32) -> bool {
-    cluster.job_state(job_id).is_some_and(|s| s.is_terminal())
+/// Whether `node` may release what it holds for `job_id`: the run is over, the
+/// job is active elsewhere, or the id is untracked but was issued by us.
+fn is_reclaimable(cluster: &ClusterManager, node: &str, job_id: u32) -> bool {
+    match cluster.job_state(job_id) {
+        Some(state) if state.is_terminal() => true,
+        // Only an active job has an authoritative nodelist; anything earlier may
+        // be mid-dispatch to this very node, so spare it.
+        Some(state) => state.is_active() && !cluster.job_holds_node(job_id, node),
+        None => job_id < cluster.peek_next_job_id(),
+    }
 }
 
-/// Reported ids the controller's own record marks terminal. Non-terminal (incl.
-/// Pending mid-dispatch) and unknown ids are spared, so no live job is reclaimed.
-fn stale_reported_jobs(cluster: &ClusterManager, reported: &[RunningJobStatus]) -> Vec<u32> {
+/// Reported ids `node` may release; ids still allocated here, not yet started,
+/// and never issued by this controller are spared.
+fn stale_reported_jobs(
+    cluster: &ClusterManager,
+    node: &str,
+    reported: &[RunningJobStatus],
+) -> Vec<u32> {
     reported
         .iter()
-        .filter_map(|r| is_still_terminal(cluster, r.job_id).then_some(r.job_id))
+        .filter_map(|r| is_reclaimable(cluster, node, r.job_id).then_some(r.job_id))
         .collect()
 }
 
@@ -3669,18 +3679,216 @@ mod tests {
             })
             .collect();
 
-        let stale = stale_reported_jobs(&cluster, &reported);
+        let stale = stale_reported_jobs(&cluster, "n1", &reported);
         assert_eq!(
             stale,
             vec![12],
-            "only the terminal job is reclaimed; Pending/Running/Completing/Suspended/Preempted/unknown are spared"
+            "terminal is reclaimed; Pending/Running/Completing/Suspended/Preempted spared, and 999 was never issued"
+        );
+    }
+
+    /// GATE: a terminal job aged out of the job map must stay reclaimable, or an
+    /// agent still holding it keeps that allocation forever.
+    #[tokio::test]
+    async fn stale_reported_jobs_reclaims_job_evicted_from_memory() {
+        use crate::raft::StateMachineApply;
+        use spur_core::job::{JobSpec, JobState};
+        use spur_core::wal::WalOperation;
+
+        let dir = tempfile::TempDir::new().unwrap();
+        let cluster =
+            Arc::new(crate::cluster::ClusterManager::new(test_slurm_config(), dir.path()).unwrap());
+        let apply = |op: &WalOperation| {
+            <crate::cluster::ClusterManager as StateMachineApply>::apply_operation(
+                cluster.as_ref(),
+                op,
+            );
+        };
+
+        apply(&WalOperation::JobSubmit {
+            job_id: 20,
+            spec: Box::new(JobSpec {
+                name: "evicted".into(),
+                user: "alice".into(),
+                num_nodes: 1,
+                num_tasks: 1,
+                cpus_per_task: 1,
+                work_dir: "/tmp".into(),
+                ..Default::default()
+            }),
+        });
+        apply(&WalOperation::job_state_change(
+            20,
+            JobState::Pending,
+            JobState::Cancelled,
+        ));
+        apply(&WalOperation::EvictTerminalJobs { job_ids: vec![20] });
+
+        assert_eq!(cluster.job_state(20), None, "eviction drops the record");
+        let issued = cluster.peek_next_job_id();
+        assert!(20 < issued, "id 20 was issued by this controller");
+
+        let reported: Vec<RunningJobStatus> = [20, issued, issued + 5]
+            .into_iter()
+            .map(|job_id| RunningJobStatus {
+                job_id,
+                ..Default::default()
+            })
+            .collect();
+
+        assert_eq!(
+            stale_reported_jobs(&cluster, "n1", &reported),
+            vec![20],
+            "the evicted id is reclaimed; ids at or above next_job_id were never issued here"
+        );
+    }
+
+    /// Pins a known false positive: an array parent id is consumed but never
+    /// stored, so it reads as reclaimable. Agents only report dispatched tasks.
+    #[tokio::test]
+    async fn reclaimable_reports_true_for_a_consumed_but_unstored_id() {
+        use crate::raft::StateMachineApply;
+        use spur_core::job::JobSpec;
+        use spur_core::wal::WalOperation;
+
+        let dir = tempfile::TempDir::new().unwrap();
+        let cluster =
+            Arc::new(crate::cluster::ClusterManager::new(test_slurm_config(), dir.path()).unwrap());
+
+        for task_id in [31, 32] {
+            <crate::cluster::ClusterManager as StateMachineApply>::apply_operation(
+                cluster.as_ref(),
+                &WalOperation::JobSubmit {
+                    job_id: task_id,
+                    spec: Box::new(JobSpec {
+                        name: "array-task".into(),
+                        user: "alice".into(),
+                        num_nodes: 1,
+                        num_tasks: 1,
+                        cpus_per_task: 1,
+                        work_dir: "/tmp".into(),
+                        array_job_id: Some(30),
+                        ..Default::default()
+                    }),
+                },
+            );
+        }
+
+        assert_eq!(cluster.job_state(30), None, "parent id is never stored");
+        assert!(30 < cluster.peek_next_job_id());
+        assert!(
+            is_reclaimable(&cluster, "n1", 30),
+            "an unstored id below the watermark is reclaimable — reachable only if an agent reports it"
+        );
+    }
+
+    /// A node evicted mid-run keeps the job and its devices; the controller
+    /// restarts it elsewhere and must let the old node release.
+    #[tokio::test]
+    async fn stale_reported_jobs_reclaims_a_job_restarted_on_another_node() {
+        use crate::raft::StateMachineApply;
+        use spur_core::job::{JobSpec, JobState};
+        use spur_core::wal::WalOperation;
+
+        let dir = tempfile::TempDir::new().unwrap();
+        let cluster =
+            Arc::new(crate::cluster::ClusterManager::new(test_slurm_config(), dir.path()).unwrap());
+        let apply = |op: &WalOperation| {
+            <crate::cluster::ClusterManager as StateMachineApply>::apply_operation(
+                cluster.as_ref(),
+                op,
+            );
+        };
+
+        apply(&WalOperation::JobSubmit {
+            job_id: 40,
+            spec: Box::new(JobSpec {
+                name: "moved".into(),
+                user: "alice".into(),
+                num_nodes: 1,
+                num_tasks: 1,
+                cpus_per_task: 1,
+                work_dir: "/tmp".into(),
+                ..Default::default()
+            }),
+        });
+
+        let res = spur_core::resource::ResourceAllocations {
+            cpus: 1,
+            memory_mb: 0,
+            devices: std::collections::HashMap::new(),
+        };
+        let mut per_node = std::collections::HashMap::new();
+        per_node.insert("n2".to_string(), res.clone());
+        apply(&WalOperation::job_start(
+            40,
+            vec!["n2".into()],
+            res,
+            per_node,
+        ));
+        apply(&WalOperation::job_state_change(
+            40,
+            JobState::Pending,
+            JobState::Running,
+        ));
+
+        let reported = vec![RunningJobStatus {
+            job_id: 40,
+            ..Default::default()
+        }];
+        assert_eq!(
+            stale_reported_jobs(&cluster, "n1", &reported),
+            vec![40],
+            "the node it no longer runs on may release it"
+        );
+        assert!(
+            stale_reported_jobs(&cluster, "n2", &reported).is_empty(),
+            "the node actually running it must never be told to release"
+        );
+    }
+
+    /// A job dispatched but not yet started has no nodelist, so the node it is
+    /// being launched on must not be told to release it.
+    #[tokio::test]
+    async fn stale_reported_jobs_spares_a_job_mid_dispatch() {
+        use crate::raft::StateMachineApply;
+        use spur_core::job::JobSpec;
+        use spur_core::wal::WalOperation;
+
+        let dir = tempfile::TempDir::new().unwrap();
+        let cluster =
+            Arc::new(crate::cluster::ClusterManager::new(test_slurm_config(), dir.path()).unwrap());
+
+        <crate::cluster::ClusterManager as StateMachineApply>::apply_operation(
+            cluster.as_ref(),
+            &WalOperation::JobSubmit {
+                job_id: 50,
+                spec: Box::new(JobSpec {
+                    name: "launching".into(),
+                    user: "alice".into(),
+                    num_nodes: 1,
+                    num_tasks: 1,
+                    cpus_per_task: 1,
+                    work_dir: "/tmp".into(),
+                    ..Default::default()
+                }),
+            },
+        );
+
+        let reported = vec![RunningJobStatus {
+            job_id: 50,
+            ..Default::default()
+        }];
+        assert!(
+            stale_reported_jobs(&cluster, "n1", &reported).is_empty(),
+            "a Pending job the agent already holds is mid-launch, not stale"
         );
     }
 
     /// GATE: a job requeued (Timeout -> Pending) between the reclaim snapshot
     /// and the spawned loop's send must fail the re-check, not just the snapshot.
     #[tokio::test]
-    async fn is_still_terminal_false_after_requeue_race() {
+    async fn is_reclaimable_false_after_requeue_race() {
         use crate::raft::StateMachineApply;
         use spur_core::job::{JobSpec, JobState};
         use spur_core::wal::WalOperation;
@@ -3731,7 +3939,7 @@ mod tests {
             JobState::Running,
             JobState::Timeout,
         ));
-        assert!(is_still_terminal(&cluster, 77), "snapshot sees Timeout");
+        assert!(is_reclaimable(&cluster, "n1", 77), "snapshot sees Timeout");
 
         // Concurrent requeue lands before the reclaim loop's re-check.
         apply(&WalOperation::job_state_change(
@@ -3741,7 +3949,7 @@ mod tests {
         ));
 
         assert!(
-            !is_still_terminal(&cluster, 77),
+            !is_reclaimable(&cluster, "n1", 77),
             "re-check must skip a job requeued since the snapshot"
         );
     }