diff --git a/crates/infrastructure/src/store/background_task.rs b/crates/infrastructure/src/store/background_task.rs index a691114..e4706c5 100644 --- a/crates/infrastructure/src/store/background_task.rs +++ b/crates/infrastructure/src/store/background_task.rs @@ -6,10 +6,11 @@ //! /.ideai/background-tasks/.json //! ``` //! -//! Each file is a small JSON document `{ version, projectId, tasks }`. Writes are -//! atomic: serialize to `.json.tmp`, then rename over the target. The -//! in-memory registry indexes open tasks by owner agent and task id; it is -//! rebuilt lazily from disk on first use and updated on every successful mutation. +//! Each file is a small JSON document `{ version, projectId, tasks }`. Mutations +//! are serialized per project and writes are atomic: serialize to a unique tmp +//! path, then rename over the target. The in-memory registry indexes open tasks +//! by owner agent and task id; it is rebuilt lazily from disk on first use and +//! updated on every successful mutation. //! //! Boot reconcile rule for B2: because the runner registry is not wired yet, a //! task in `Running` or `Waiting` whose id is not present in the caller-provided @@ -19,10 +20,12 @@ use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::path::{Path, PathBuf}; +use std::sync::Arc; use async_trait::async_trait; use serde::{Deserialize, Serialize}; -use tokio::sync::RwLock; +use tokio::sync::{Mutex, RwLock}; +use uuid::Uuid; use domain::{ AgentId, BackgroundTask, BackgroundTaskPortError, BackgroundTaskResult, BackgroundTaskState, @@ -37,6 +40,7 @@ const TASK_DOC_VERSION: u32 = 1; pub struct FsBackgroundTaskStore { dir: PathBuf, registry: RwLock, + project_locks: Mutex>>>, } #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -78,6 +82,7 @@ impl FsBackgroundTaskStore { Self { dir, registry: RwLock::new(RegistryIndex::default()), + project_locks: Mutex::new(HashMap::new()), } } @@ -148,10 +153,31 @@ impl FsBackgroundTaskStore { } for project_id in changed_projects { - if let Some(doc) = state.docs.get(&project_id) { - self.write_doc(doc).await?; - } + let live = &live; + self.mutate_project_doc(project_id, |doc| { + for task in &mut doc.tasks { + if matches!( + task.state, + BackgroundTaskState::Running | BackgroundTaskState::Waiting + ) && !live.contains(&task.id) + { + let finished_at_ms = now_ms.max(task.updated_at_ms); + let result = BackgroundTaskResult::Failure { + finished_at_ms, + exit_code: None, + error: "background task lost its runtime handle during IdeA restart" + .into(), + stdout_tail: None, + stderr_tail: None, + }; + *task = task.complete(result).map_err(invalid_task)?; + } + } + Ok(()) + }) + .await?; } + let state = self.read_all_docs().await?; self.replace_registry_from_state(&state).await; Ok(report) } @@ -161,7 +187,10 @@ impl FsBackgroundTaskStore { return Ok(()); } let state = self.read_all_docs().await?; - self.replace_registry_from_state(&state).await; + let mut registry = self.registry.write().await; + if !registry.loaded { + *registry = RegistryIndex::from_state(&state); + } Ok(()) } @@ -175,7 +204,30 @@ impl FsBackgroundTaskStore { } fn tmp_path_for_project(&self, project_id: ProjectId) -> PathBuf { - self.dir.join(format!("{project_id}.json.tmp")) + self.dir + .join(format!("{project_id}.json.{}.tmp", Uuid::new_v4())) + } + + async fn project_lock(&self, project_id: ProjectId) -> Arc> { + let mut locks = self.project_locks.lock().await; + Arc::clone( + locks + .entry(project_id) + .or_insert_with(|| Arc::new(Mutex::new(()))), + ) + } + + async fn mutate_project_doc( + &self, + project_id: ProjectId, + mutate: impl FnOnce(&mut TaskDoc) -> Result<(), BackgroundTaskPortError>, + ) -> Result { + let lock = self.project_lock(project_id).await; + let _guard = lock.lock().await; + let mut doc = self.read_doc(project_id).await?; + mutate(&mut doc)?; + self.write_doc(&doc).await?; + Ok(doc) } async fn read_all_docs(&self) -> Result { @@ -303,12 +355,14 @@ impl BackgroundTaskStore for FsBackgroundTaskStore { return Err(BackgroundTaskPortError::AlreadyExists); } - let mut doc = self.read_doc(task.project_id).await?; - if doc.tasks.iter().any(|existing| existing.id == task.id) { - return Err(BackgroundTaskPortError::AlreadyExists); - } - doc.tasks.push(task.clone()); - self.write_doc(&doc).await?; + self.mutate_project_doc(task.project_id, |doc| { + if doc.tasks.iter().any(|existing| existing.id == task.id) { + return Err(BackgroundTaskPortError::AlreadyExists); + } + doc.tasks.push(task.clone()); + Ok(()) + }) + .await?; self.update_registry_for_task(task).await; Ok(()) } @@ -323,13 +377,15 @@ impl BackgroundTaskStore for FsBackgroundTaskStore { async fn save(&self, task: &BackgroundTask) -> Result<(), BackgroundTaskPortError> { self.ensure_registry_loaded().await?; - let mut doc = self.read_doc(task.project_id).await?; - if let Some(slot) = doc.tasks.iter_mut().find(|existing| existing.id == task.id) { - *slot = task.clone(); - } else { - return Err(BackgroundTaskPortError::NotFound); - } - self.write_doc(&doc).await?; + self.mutate_project_doc(task.project_id, |doc| { + if let Some(slot) = doc.tasks.iter_mut().find(|existing| existing.id == task.id) { + *slot = task.clone(); + Ok(()) + } else { + Err(BackgroundTaskPortError::NotFound) + } + }) + .await?; self.update_registry_for_task(task).await; Ok(()) } diff --git a/crates/infrastructure/tests/background_task_store.rs b/crates/infrastructure/tests/background_task_store.rs index 595fb95..85a2ba1 100644 --- a/crates/infrastructure/tests/background_task_store.rs +++ b/crates/infrastructure/tests/background_task_store.rs @@ -1,6 +1,7 @@ //! B2 integration tests for [`FsBackgroundTaskStore`]. use std::path::PathBuf; +use std::sync::Arc; use domain::{ AgentId, BackgroundTask, BackgroundTaskKind, BackgroundTaskResult, BackgroundTaskState, @@ -33,6 +34,17 @@ impl TempDir { fn task_tmp_file(&self, project_id: ProjectId) -> PathBuf { self.task_dir().join(format!("{project_id}.json.tmp")) } + + fn tmp_files(&self) -> Vec { + let Ok(entries) = std::fs::read_dir(self.task_dir()) else { + return Vec::new(); + }; + entries + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("tmp")) + .collect() + } } impl Drop for TempDir { @@ -198,6 +210,93 @@ async fn persistence_is_segmented_by_project_id() { assert!(!raw_p2.contains(&p1.id.to_string())); } +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +async fn concurrent_creates_for_same_project_do_not_lose_tasks() { + let tmp = TempDir::new(); + let store = Arc::new(FsBackgroundTaskStore::new(&tmp.project_path())); + let project = project_id(10); + let owner = agent_id(100); + + let mut handles = Vec::new(); + for n in 1..=64 { + let store = Arc::clone(&store); + handles.push(tokio::spawn(async move { + let task = queued_task(n, 10, 100, 1_000 + n as u64); + store.create(&task).await.map(|()| task.id) + })); + } + + let mut expected = Vec::new(); + for handle in handles { + expected.push(handle.await.unwrap().unwrap()); + } + expected.sort(); + + let reborn = FsBackgroundTaskStore::new(&tmp.project_path()); + for id in &expected { + assert!( + reborn.get(*id).await.unwrap().is_some(), + "created task {id} must be durable" + ); + } + + assert_eq!( + store.registry_open_task_ids_for_agent(owner).await, + expected, + "registry must retain every concurrent create" + ); + assert!(tmp.tmp_files().is_empty(), "tmp files left behind"); + assert!(tmp.task_file(project).exists()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +async fn concurrent_saves_for_same_project_do_not_overwrite_each_other() { + let tmp = TempDir::new(); + let store = Arc::new(FsBackgroundTaskStore::new(&tmp.project_path())); + + for n in 1..=64 { + store + .create(&running_task(n, 10, 100, 1_000)) + .await + .unwrap(); + } + + let mut handles = Vec::new(); + for n in 1..=64 { + let store = Arc::clone(&store); + handles.push(tokio::spawn(async move { + let task = completed_task(n, 10, 100, 2_000 + n as u64); + store.save(&task).await.map(|()| task.id) + })); + } + + let mut expected = Vec::new(); + for handle in handles { + expected.push(handle.await.unwrap().unwrap()); + } + expected.sort(); + + let reborn = FsBackgroundTaskStore::new(&tmp.project_path()); + for id in &expected { + let task = reborn.get(*id).await.unwrap().unwrap(); + assert_eq!(task.state, BackgroundTaskState::Completed); + } + + assert!( + reborn + .list_open_for_agent(agent_id(100)) + .await + .unwrap() + .is_empty(), + "all concurrently saved tasks must be terminal" + ); + assert_eq!( + reborn.list_undelivered_completions().await.unwrap().len(), + expected.len() + ); + assert!(tmp.tmp_files().is_empty(), "tmp files left behind"); +} + #[tokio::test] async fn reconcile_marks_running_without_live_handle_failed_and_pending() { let tmp = TempDir::new();