diff --git a/crates/workspace-server/src/backend_job_recovery_store.rs b/crates/workspace-server/src/backend_job_recovery_store.rs index 373cfc57a..cf33c8967 100644 --- a/crates/workspace-server/src/backend_job_recovery_store.rs +++ b/crates/workspace-server/src/backend_job_recovery_store.rs @@ -852,7 +852,7 @@ mod tests { assert!(acquired); assert_eq!( store - .recover_backend_job_worker_cleanups("workspace", "later", 100) + .list_backend_job_worker_cleanups("workspace", None, 100) .unwrap(), vec![executing.clone()] ); @@ -953,7 +953,7 @@ mod tests { }).unwrap(); assert!( store - .list_backend_job_worker_cleanups("workspace", 100) + .list_backend_job_worker_cleanups("workspace", None, 100) .unwrap() .is_empty() ); @@ -977,7 +977,7 @@ mod tests { Ok(()) }).unwrap(); let selected = store - .recover_backend_job_worker_cleanups("workspace", NOW, 100) + .list_backend_job_worker_cleanups("workspace", None, 100) .unwrap(); assert_eq!(selected.len(), 1); assert_eq!(selected[0].attempt_id, journal.attempt.attempt_id); @@ -1000,7 +1000,7 @@ mod tests { .unwrap(); assert!( store - .recover_backend_job_worker_cleanups("workspace", NOW, 100) + .list_backend_job_worker_cleanups("workspace", None, 100) .unwrap() .is_empty() ); diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index fa3aee993..bead2e4c8 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -5386,25 +5386,28 @@ impl WorkspaceApi { // Available input. Its owner takes the retirement-only step before claim; // neither a live Worker nor a new removal operation is required. let page_limit = usize::from(ABSOLUTE_MAX_CONCURRENT_JOBS); - let mut attempts = self.store.recover_backend_job_worker_cleanups( - &self.config.workspace_id, - &Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true), - page_limit, - )?; + let mut after: Option<(String, String)> = None; loop { - if attempts.is_empty() { + let attempts = self.store.list_backend_job_worker_cleanups( + &self.config.workspace_id, + after + .as_ref() + .map(|(job, attempt)| (job.as_str(), attempt.as_str())), + page_limit, + )?; + let Some(last) = attempts.last() else { + return Ok(()); + }; + let next = (last.job_id.clone(), last.attempt_id.clone()); + for attempt in &attempts { + self.schedule_backend_job_worker_cleanup(attempt); + } + if attempts.len() < page_limit { return Ok(()); } - let scheduled = attempts - .iter() - .filter(|attempt| self.schedule_backend_job_worker_cleanup(attempt)) - .count(); - if attempts.len() < page_limit || scheduled == 0 { - return Ok(()); - } - attempts = self - .store - .list_backend_job_worker_cleanups(&self.config.workspace_id, page_limit)?; + // Deduplicated/Executing owners must not hide a later page. Their + // progress and timestamps are independent of this read-only cursor. + after = Some(next); } } diff --git a/crates/workspace-server/src/server/backend_job_recovery.rs b/crates/workspace-server/src/server/backend_job_recovery.rs index 50bd68365..b6cc9e42e 100644 --- a/crates/workspace-server/src/server/backend_job_recovery.rs +++ b/crates/workspace-server/src/server/backend_job_recovery.rs @@ -23,9 +23,21 @@ pub(super) struct RecoveryOwners { pub(super) storage_probe_started: tokio::sync::Notify, #[cfg(test)] pub(super) nonjournal_lookup_pending: tokio::sync::Notify, + #[cfg(test)] + pub(super) cleanup_before_claim: Mutex>>, } impl RecoveryOwners { + #[cfg(test)] + pub(super) fn resume_after_test_pause(&self) { + self.stopping.store(false, Ordering::Release); + } + + #[cfg(test)] + pub(super) fn cleanup_owner_ids(&self) -> HashSet { + self.cleanup_running.lock().unwrap().clone() + } + async fn pause(&self, duration: Duration) { let notified = self.changed.notified(); tokio::pin!(notified); @@ -389,6 +401,21 @@ impl WorkspaceApi { // inside the result's Job lease would make close_and_drain wait on itself. runtime.spawn(async move { let _lease = lease; + #[cfg(test)] + { + let gate = api + .backend_job_recovery + .cleanup_before_claim + .lock() + .unwrap() + .clone(); + if let Some(gate) = gate { + let Ok(permit) = gate.acquire().await else { + return; + }; + drop(permit); + } + } if let Err(error) = api.run_backend_job_cleanup(expected.clone()).await { tracing::warn!(job_id=%expected.job_id, attempt_id=%expected.attempt_id, %error, "Job cleanup stopped; execution/input/resource fences remain retained"); diff --git a/crates/workspace-server/src/server/tests/backend_job_tests.rs b/crates/workspace-server/src/server/tests/backend_job_tests.rs index ef623d7d2..2d11d5c70 100644 --- a/crates/workspace-server/src/server/tests/backend_job_tests.rs +++ b/crates/workspace-server/src/server/tests/backend_job_tests.rs @@ -1421,6 +1421,120 @@ async fn backend_job_input_response_loss_is_owned_until_deadline_without_resubmi ); } +#[tokio::test] +async fn cleanup_recovery_owns_later_pages_while_first_page_is_held_before_claim() { + let temp = tempfile::tempdir().unwrap(); + let (api, execution) = test_api_with_recording_backend(temp.path()).await; + // Seed accepted/terminal Jobs without automatic owner tasks. Keep the actual + // original Host owners: this fixture pause is not a reconstructed restart fence. + api.stop_backend_job_recovery().await; + let count = usize::from(ABSOLUTE_MAX_CONCURRENT_JOBS) + 2; + let mut attempts = Vec::new(); + for index in 0..count { + let request = recovery_request(&format!("cleanup-page-{index:03}")); + let dispatched = api.dispatch_backend_job(&request).unwrap(); + seed_worker_session_for_cleanup(temp.path(), dispatched.attempt.worker.as_ref().unwrap()); + let terminal = api + .store + .finish_backend_job_attempt( + TEST_WORKSPACE_ID, + &request.job_id, + &dispatched.attempt.attempt_id, + BackendJobAttemptState::Failed, + "fixture_failure", + "Original result retained", + &Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true), + ) + .unwrap(); + attempts.push(terminal); + } + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + *api.backend_job_recovery + .cleanup_before_claim + .lock() + .unwrap() = Some(gate.clone()); + api.backend_job_recovery.resume_after_test_pause(); + api.recover_backend_job_worker_cleanups().unwrap(); + let owners = api.backend_job_recovery.cleanup_owner_ids(); + assert_eq!( + owners.len(), + count, + "later pages must be owned before any prior cleanup claim progresses" + ); + assert!( + attempts + .iter() + .all(|attempt| owners.contains(&attempt.attempt_id)) + ); + for attempt in &attempts { + assert_eq!( + api.store + .get_backend_job_attempt(TEST_WORKSPACE_ID, &attempt.job_id, &attempt.attempt_id) + .unwrap() + .unwrap() + .worker_cleanup_state, + Some(BackendJobWorkerCleanupState::Pending) + ); + } + // A concurrent sweep sees the same admitted owners and must advance past + // them too, without creating another owner or needing a later user request. + api.recover_backend_job_worker_cleanups().unwrap(); + assert_eq!(api.backend_job_recovery.cleanup_owner_ids(), owners); + gate.add_permits(count); + // The two finite cleanup workers perform real canonical removal/retention + // writes for every fixture Worker; allow bounded per-Worker drain time. + tokio::time::timeout(std::time::Duration::from_secs(3 * count as u64), async { + loop { + let mut complete = true; + for attempt in &attempts { + let saved = api + .store + .get_backend_job_attempt( + TEST_WORKSPACE_ID, + &attempt.job_id, + &attempt.attempt_id, + ) + .unwrap() + .unwrap(); + assert_eq!(saved.worker, attempt.worker); + assert_eq!(saved.runtime_run_id, attempt.runtime_run_id); + assert_eq!(saved.failure_category, attempt.failure_category); + if saved.worker_cleanup_state != Some(BackendJobWorkerCleanupState::Completed) + && !api.backend_job_recovery.cleanup_owner_ids().contains(&saved.attempt_id) + { + panic!("cleanup owner exited before completion: job={} state={:?} cause={:?}: {:?}", + saved.job_id, saved.worker_cleanup_state, + saved.worker_cleanup_failure_category, saved.worker_cleanup_failure_detail); + } + complete &= + saved.worker_cleanup_state == Some(BackendJobWorkerCleanupState::Completed); + } + if complete { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| { + let pending: Vec<_> = attempts.iter().filter_map(|attempt| { + let saved = api.store.get_backend_job_attempt(TEST_WORKSPACE_ID, &attempt.job_id, &attempt.attempt_id) + .unwrap().unwrap(); + (saved.worker_cleanup_state != Some(BackendJobWorkerCleanupState::Completed)).then_some(( + saved.job_id, saved.worker_cleanup_state, + saved.worker_cleanup_failure_category, saved.worker_cleanup_failure_detail)) + }).collect(); + panic!("cleanup page did not finish: live_owners={}, pending={pending:?}", + api.backend_job_recovery.cleanup_owner_ids().len()); + }); + assert_eq!( + execution.take_inputs().len(), + count, + "cleanup never repeats model input" + ); + api.stop_backend_job_recovery().await; +} + #[tokio::test] async fn backend_job_recovery_drains_more_than_one_reserved_page() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index c91ce632d..a6be930d1 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -2131,15 +2131,10 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { failure_detail: &str, now: &str, ) -> Result; - fn recover_backend_job_worker_cleanups( - &self, - workspace_id: &str, - now: &str, - limit: usize, - ) -> Result>; fn list_backend_job_worker_cleanups( &self, workspace_id: &str, + after: Option<(&str, &str)>, limit: usize, ) -> Result>; fn claim_backend_job_worker_cleanup( @@ -8731,28 +8726,20 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } - fn recover_backend_job_worker_cleanups( - &self, - workspace_id: &str, - _now: &str, - limit: usize, - ) -> Result> { - // Reopening observes retained owners; it does not revoke an executing - // cleanup or authorize a new one by resetting its state. - self.list_backend_job_worker_cleanups(workspace_id, limit) - } - fn list_backend_job_worker_cleanups( &self, workspace_id: &str, + after: Option<(&str, &str)>, limit: usize, ) -> Result> { self.with_conn(|conn| { let mut statement = conn.prepare( - "SELECT a.job_id, a.attempt_id FROM backend_job_attempts a JOIN backend_jobs j USING(workspace_id,job_id) WHERE a.workspace_id = ?1 AND a.state IN ('completed', 'failed') AND a.runtime_id IS NOT NULL AND (a.worker_cleanup_state IN ('pending', 'failed', 'executing') OR (a.state='completed' AND a.worker_cleanup_state='completed' AND a.attempt=j.current_attempt AND j.input_disposition='available' AND json_type(j.request_json,'$.grants.subjektiv_journal')='object')) AND NOT EXISTS(SELECT 1 FROM backend_job_upgrade_history h WHERE h.workspace_id=a.workspace_id AND h.job_id=a.job_id) ORDER BY a.worker_cleanup_updated_at, a.created_at LIMIT ?2", + "SELECT a.job_id, a.attempt_id FROM backend_job_attempts a JOIN backend_jobs j USING(workspace_id,job_id) WHERE a.workspace_id = ?1 AND (?3 IS NULL OR (a.job_id,a.attempt_id) > (?3,?4)) AND a.state IN ('completed', 'failed') AND a.runtime_id IS NOT NULL AND (a.worker_cleanup_state IN ('pending', 'failed', 'executing') OR (a.state='completed' AND a.worker_cleanup_state='completed' AND a.attempt=j.current_attempt AND j.input_disposition='available' AND json_type(j.request_json,'$.grants.subjektiv_journal')='object')) AND NOT EXISTS(SELECT 1 FROM backend_job_upgrade_history h WHERE h.workspace_id=a.workspace_id AND h.job_id=a.job_id) ORDER BY a.job_id, a.attempt_id LIMIT ?2", )?; + // Immutable identities keep the boundary stable while admitted + // owners change cleanup state/timestamps. Selection grants no claim. let identities = statement - .query_map(params![workspace_id, limit as i64], |row| { + .query_map(params![workspace_id, limit as i64, after.map(|(job, _)| job), after.map(|(_, attempt)| attempt)], |row| { Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) })? .collect::>>()?;