| 1 | use super::*; |
| 2 | use std::io::Write; |
| 3 | use std::process::{Child, Command, Stdio}; |
| 4 | use std::sync::atomic::{AtomicUsize, Ordering}; |
| 5 | |
| 6 | fn config(root: &Path, scope: &str) -> TaskManagerConfig { |
| 7 | TaskManagerConfig { |
| 8 | data_dir: root.to_path_buf(), |
| 9 | worker_count: 1, |
| 10 | default_workspace: root.join(scope), |
| 11 | default_model: format!("{scope}-model"), |
| 12 | default_mode: "plan".into(), |
| 13 | allow_shell: false, |
| 14 | trust_mode: false, |
| 15 | execution_limits: TaskExecutionLimits::default(), |
| 16 | } |
| 17 | } |
| 18 | |
| 19 | struct RecordingExecutor { |
| 20 | root: PathBuf, |
| 21 | scope: String, |
| 22 | } |
| 23 | |
| 24 | #[async_trait] |
| 25 | impl TaskExecutor for RecordingExecutor { |
| 26 | async fn execute( |
| 27 | &self, |
| 28 | task: ExecutionTask, |
| 29 | events: mpsc::Sender<TaskExecutionEvent>, |
| 30 | cancel: CancellationToken, |
| 31 | ) -> TaskExecutionResult { |
| 32 | let mut receipt = fs::OpenOptions::new() |
| 33 | .create(true) |
| 34 | .append(true) |
| 35 | .open(self.root.join(format!("{}.executions", self.scope))) |
| 36 | .unwrap(); |
| 37 | writeln!( |
| 38 | receipt, |
| 39 | "{}", |
| 40 | serde_json::json!({ |
| 41 | "id": task.id, "scope": self.scope, "model": task.model, |
| 42 | "provider": task.model_provider, "provider_id": task.model_provider_id, |
| 43 | "workspace": task.workspace, "prompt": task.prompt, |
| 44 | "allow_shell": task.allow_shell, "trust_mode": task.trust_mode, |
| 45 | }) |
| 46 | ) |
| 47 | .unwrap(); |
| 48 | if task.prompt.ends_with("hold") { |
| 49 | cancel.cancelled().await; |
| 50 | fs::write(self.root.join(format!("{}.canceled", self.scope)), &task.id).unwrap(); |
| 51 | } |
| 52 | events |
| 53 | .send(TaskExecutionEvent::MessageDelta { |
| 54 | content: format!("{} result", self.scope), |
| 55 | }) |
| 56 | .await |
| 57 | .ok(); |
| 58 | TaskExecutionResult { |
| 59 | status: TaskStatus::Completed, |
| 60 | result_text: Some(format!("{} completed", self.scope)), |
| 61 | error: None, |
| 62 | terminal_reason: TaskTerminalReason::Completed, |
| 63 | } |
| 64 | } |
| 65 | } |
| 66 | |
| 67 | async fn manager(root: &Path, scope: &str) -> Result<SharedTaskManager> { |
| 68 | TaskManager::start_with_executor_in_scope( |
| 69 | config(root, scope), |
| 70 | Arc::new(RecordingExecutor { |
| 71 | root: root.to_path_buf(), |
| 72 | scope: scope.into(), |
| 73 | }), |
| 74 | scope, |
| 75 | ) |
| 76 | .await |
| 77 | } |
| 78 | |
| 79 | async fn wait_file(path: &Path) -> Result<Vec<u8>> { |
| 80 | tokio::time::timeout(Duration::from_secs(15), async { |
| 81 | loop { |
| 82 | match fs::read(path) { |
| 83 | Ok(bytes) if !bytes.is_empty() => return Ok(bytes), |
| 84 | Ok(_) => {} |
| 85 | Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} |
| 86 | Err(error) => return Err(error.into()), |
| 87 | } |
| 88 | sleep(Duration::from_millis(10)).await; |
| 89 | } |
| 90 | }) |
| 91 | .await |
| 92 | .context("fixture receipt timeout")? |
| 93 | } |
| 94 | |
| 95 | fn executions(root: &Path, scope: &str) -> Vec<Value> { |
| 96 | fs::read_to_string(root.join(format!("{scope}.executions"))) |
| 97 | .unwrap_or_default() |
| 98 | .lines() |
| 99 | .map(|line| serde_json::from_str(line).unwrap()) |
| 100 | .collect() |
| 101 | } |
| 102 | |
| 103 | struct ChildOwner(Child); |
| 104 | impl Drop for ChildOwner { |
| 105 | fn drop(&mut self) { |
| 106 | let _ = self.0.kill(); |
| 107 | let _ = self.0.wait(); |
| 108 | } |
| 109 | } |
| 110 | fn child_owner(root: &Path) -> Result<ChildOwner> { |
| 111 | let log = fs::File::create(root.join("child.log"))?; |
| 112 | Ok(ChildOwner( |
| 113 | Command::new(std::env::current_exe()?) |
| 114 | .args([ |
| 115 | "--exact", |
| 116 | "task_manager::ownership_tests::ownership_process_child", |
| 117 | "--ignored", |
| 118 | "--nocapture", |
| 119 | ]) |
| 120 | .env("CW_TASK_OWNER_FIXTURE_ROOT", root) |
| 121 | .stdout(Stdio::from(log.try_clone()?)) |
| 122 | .stderr(Stdio::from(log)) |
| 123 | .spawn()?, |
| 124 | )) |
| 125 | } |
| 126 | |
| 127 | #[tokio::test] |
| 128 | #[ignore = "subprocess entry; invoked by shared-store ownership fixtures"] |
| 129 | async fn ownership_process_child() -> Result<()> { |
| 130 | let root = PathBuf::from(std::env::var("CW_TASK_OWNER_FIXTURE_ROOT")?); |
| 131 | let owner = manager(&root, "A").await?; |
| 132 | let active = owner |
| 133 | .add_task(NewTaskRequest { |
| 134 | owner_session_id: Some("session-a".into()), |
| 135 | ..NewTaskRequest::from_prompt("A hold") |
| 136 | }) |
| 137 | .await?; |
| 138 | wait_file(&root.join("A.executions")).await?; |
| 139 | let queued = owner |
| 140 | .add_task(NewTaskRequest { |
| 141 | owner_session_id: Some("session-a".into()), |
| 142 | ..NewTaskRequest::from_prompt("A queued") |
| 143 | }) |
| 144 | .await?; |
| 145 | fs::write( |
| 146 | root.join("ready"), |
| 147 | serde_json::to_vec(&[active.id, queued.id])?, |
| 148 | )?; |
| 149 | wait_file(&root.join("stop")).await?; |
| 150 | owner.shutdown_and_wait().await?; |
| 151 | drop(owner); |
| 152 | Ok(()) |
| 153 | } |
| 154 | |
| 155 | #[tokio::test] |
| 156 | async fn second_process_preserves_live_owner_and_durable_cancel_reaches_that_owner() -> Result<()> { |
| 157 | let root = tempfile::tempdir()?; |
| 158 | let mut child = child_owner(root.path())?; |
| 159 | let ids: Vec<String> = serde_json::from_slice(&wait_file(&root.path().join("ready")).await?)?; |
| 160 | let active_path = root.path().join("tasks").join(format!("{}.json", ids[0])); |
| 161 | let before = fs::read(&active_path)?; |
| 162 | let other = manager(root.path(), "B").await?; |
| 163 | assert_eq!(other.get_task(&ids[0]).await?.status, TaskStatus::Running); |
| 164 | assert_eq!( |
| 165 | fs::read(&active_path)?, |
| 166 | before, |
| 167 | "opening a foreign scope cannot rewrite live work" |
| 168 | ); |
| 169 | assert!( |
| 170 | other |
| 171 | .get_task_for_owner(&ids[0], "session-b") |
| 172 | .await |
| 173 | .is_err() |
| 174 | ); |
| 175 | assert!( |
| 176 | other |
| 177 | .cancel_task_for_owner(&ids[0][..12], "session-b") |
| 178 | .await |
| 179 | .is_err() |
| 180 | ); |
| 181 | assert!( |
| 182 | other |
| 183 | .list_tasks_for_owner(None, None, "session-b") |
| 184 | .await? |
| 185 | .is_empty() |
| 186 | ); |
| 187 | let own = other |
| 188 | .add_task(NewTaskRequest::from_prompt("B quick")) |
| 189 | .await?; |
| 190 | let done = wait_for_terminal_state(&other, &own.id, Duration::from_secs(10)).await?; |
| 191 | assert_eq!(done.status, TaskStatus::Completed); |
| 192 | assert_eq!(executions(root.path(), "B").len(), 1); |
| 193 | assert_eq!(executions(root.path(), "B")[0]["model"], "B-model"); |
| 194 | assert_eq!( |
| 195 | executions(root.path(), "A").len(), |
| 196 | 1, |
| 197 | "B cannot claim A's queued request" |
| 198 | ); |
| 199 | other.cancel_task(&ids[1]).await?; |
| 200 | other.cancel_task(&ids[0]).await?; |
| 201 | wait_file(&root.path().join("A.canceled")).await?; |
| 202 | let canceled = wait_for_terminal_state(&other, &ids[0], Duration::from_secs(10)).await?; |
| 203 | assert_eq!(canceled.status, TaskStatus::Canceled); |
| 204 | assert!(canceled.cancel_requested_seq > 0); |
| 205 | assert_eq!(other.get_task(&ids[1]).await?.status, TaskStatus::Canceled); |
| 206 | assert_eq!(executions(root.path(), "A").len(), 1); |
| 207 | fs::write(root.path().join("stop"), "stop")?; |
| 208 | assert!(child.0.wait()?.success()); |
| 209 | other.shutdown_and_wait().await?; |
| 210 | Ok(()) |
| 211 | } |
| 212 | |
| 213 | #[tokio::test] |
| 214 | async fn killed_process_is_reconciled_once_and_only_its_scope_resumes_queued_work() -> Result<()> { |
| 215 | let root = tempfile::tempdir()?; |
| 216 | let mut child = child_owner(root.path())?; |
| 217 | let ids: Vec<String> = serde_json::from_slice(&wait_file(&root.path().join("ready")).await?)?; |
| 218 | let observer = manager(root.path(), "B").await?; |
| 219 | assert!( |
| 220 | manager(root.path(), "A").await.is_err(), |
| 221 | "live scope cannot have a second executor" |
| 222 | ); |
| 223 | child.0.kill()?; |
| 224 | child.0.wait()?; |
| 225 | let restarted = manager(root.path(), "A").await?; |
| 226 | let interrupted = restarted.get_task(&ids[0]).await?; |
| 227 | assert_eq!(interrupted.status, TaskStatus::Failed); |
| 228 | assert!( |
| 229 | interrupted |
| 230 | .error |
| 231 | .as_deref() |
| 232 | .unwrap() |
| 233 | .contains("Interrupted by process restart") |
| 234 | ); |
| 235 | let terminal = wait_for_terminal_state(&restarted, &ids[1], Duration::from_secs(10)).await?; |
| 236 | assert_eq!(terminal.status, TaskStatus::Completed); |
| 237 | let receipt = executions(root.path(), "A"); |
| 238 | assert_eq!(receipt.len(), 2); |
| 239 | assert_eq!( |
| 240 | receipt.iter().filter(|r| r["id"] == ids[0]).count(), |
| 241 | 1, |
| 242 | "accepted Running work cannot replay" |
| 243 | ); |
| 244 | assert_eq!(receipt[1]["id"], ids[1]); |
| 245 | assert_eq!(receipt[1]["model"], "A-model"); |
| 246 | assert!(executions(root.path(), "B").is_empty()); |
| 247 | restarted.shutdown_and_wait().await?; |
| 248 | drop(restarted); |
| 249 | let again = manager(root.path(), "A").await?; |
| 250 | assert_eq!( |
| 251 | again.get_task(&ids[0]).await?.lifecycle_seq, |
| 252 | interrupted.lifecycle_seq |
| 253 | ); |
| 254 | assert_eq!(executions(root.path(), "A").len(), 2); |
| 255 | again.shutdown_and_wait().await?; |
| 256 | observer.shutdown_and_wait().await?; |
| 257 | Ok(()) |
| 258 | } |
| 259 | |
| 260 | async fn stop_idle_workers(manager: &TaskManager) { |
| 261 | for worker in std::mem::take(&mut *manager.workers.lock().await) { |
| 262 | worker.abort(); |
| 263 | let _ = worker.await; |
| 264 | } |
| 265 | } |
| 266 | |
| 267 | struct CountExecutor(Arc<AtomicUsize>); |
| 268 | #[async_trait] |
| 269 | impl TaskExecutor for CountExecutor { |
| 270 | async fn execute( |
| 271 | &self, |
| 272 | _: ExecutionTask, |
| 273 | _: mpsc::Sender<TaskExecutionEvent>, |
| 274 | _: CancellationToken, |
| 275 | ) -> TaskExecutionResult { |
| 276 | self.0.fetch_add(1, Ordering::SeqCst); |
| 277 | TaskExecutionResult::from_reason(TaskTerminalReason::Completed, None) |
| 278 | } |
| 279 | } |
| 280 | |
| 281 | #[tokio::test] |
| 282 | async fn cancellation_after_claim_before_first_poll_never_invokes_executor() -> Result<()> { |
| 283 | let root = tempfile::tempdir()?; |
| 284 | let calls = Arc::new(AtomicUsize::new(0)); |
| 285 | let owner = TaskManager::start_with_executor_in_scope( |
| 286 | config(root.path(), "A"), |
| 287 | Arc::new(CountExecutor(calls.clone())), |
| 288 | "A", |
| 289 | ) |
| 290 | .await?; |
| 291 | stop_idle_workers(&owner).await; |
| 292 | let task = owner |
| 293 | .add_task(NewTaskRequest::from_prompt("claim barrier")) |
| 294 | .await?; |
| 295 | let (id, request, cancel) = owner.claim_next_task().await?.context("claim")?; |
| 296 | let other = manager(root.path(), "B").await?; |
| 297 | other.cancel_task(&task.id).await?; |
| 298 | owner.run_task(id, request, cancel).await; |
| 299 | assert_eq!(calls.load(Ordering::SeqCst), 0); |
| 300 | assert_eq!(owner.get_task(&task.id).await?.status, TaskStatus::Canceled); |
| 301 | let healthy = owner |
| 302 | .add_task(NewTaskRequest::from_prompt("uncanceled control")) |
| 303 | .await?; |
| 304 | let (id, request, cancel) = owner.claim_next_task().await?.context("control claim")?; |
| 305 | owner.run_task(id, request, cancel).await; |
| 306 | assert_eq!( |
| 307 | calls.load(Ordering::SeqCst), |
| 308 | 1, |
| 309 | "the same executor must run without cancellation" |
| 310 | ); |
| 311 | assert_eq!( |
| 312 | owner.get_task(&healthy.id).await?.status, |
| 313 | TaskStatus::Completed |
| 314 | ); |
| 315 | owner.shutdown_and_wait().await?; |
| 316 | other.shutdown_and_wait().await?; |
| 317 | Ok(()) |
| 318 | } |
| 319 | |
| 320 | #[tokio::test] |
| 321 | async fn pending_flush_and_finish_merge_fresh_metadata_and_monotonic_cancel() -> Result<()> { |
| 322 | let root = tempfile::tempdir()?; |
| 323 | let owner = manager(root.path(), "A").await?; |
| 324 | stop_idle_workers(&owner).await; |
| 325 | let task = owner |
| 326 | .add_task(NewTaskRequest::from_prompt("bounded deltas")) |
| 327 | .await?; |
| 328 | let (_, request, cancel) = owner.claim_next_task().await?.context("claim")?; |
| 329 | owner |
| 330 | .apply_execution_event( |
| 331 | &task.id, |
| 332 | TaskExecutionEvent::MessageDelta { |
| 333 | content: "accepted progress".into(), |
| 334 | }, |
| 335 | ) |
| 336 | .await?; |
| 337 | let other = manager(root.path(), "B").await?; |
| 338 | other |
| 339 | .record_tool_metadata( |
| 340 | &task.id, |
| 341 | &serde_json::json!({"task_updates": {"checklist": { |
| 342 | "items": [{"id": 1, "content": "merged checklist", "status": "done"}], |
| 343 | "completion_pct": 100, |
| 344 | "in_progress_id": null, |
| 345 | "updated_at": null |
| 346 | }}}), |
| 347 | ) |
| 348 | .await?; |
| 349 | other.cancel_task(&task.id).await?; |
| 350 | let seq = other.get_task(&task.id).await?.cancel_requested_seq; |
| 351 | owner.flush_task(&task.id).await?; |
| 352 | owner |
| 353 | .finish_task( |
| 354 | &task.id, |
| 355 | TaskExecutionResult::from_reason(TaskTerminalReason::Completed, None), |
| 356 | cancel, |
| 357 | &request.mode_label, |
| 358 | ) |
| 359 | .await?; |
| 360 | let final_task = other.get_task(&task.id).await?; |
| 361 | assert_eq!(final_task.status, TaskStatus::Canceled); |
| 362 | assert_eq!(final_task.cancel_requested_seq, seq); |
| 363 | assert_eq!(final_task.checklist.items.len(), 1); |
| 364 | assert_eq!(final_task.checklist.completion_pct, 100); |
| 365 | assert_eq!( |
| 366 | final_task |
| 367 | .timeline |
| 368 | .iter() |
| 369 | .filter(|event| event.summary == "accepted progress") |
| 370 | .count(), |
| 371 | 1 |
| 372 | ); |
| 373 | owner.shutdown_and_wait().await?; |
| 374 | other.shutdown_and_wait().await?; |
| 375 | Ok(()) |
| 376 | } |
| 377 | |
| 378 | #[cfg(unix)] |
| 379 | struct ReadOnlyFixtureDir(PathBuf); |
| 380 | #[cfg(unix)] |
| 381 | impl ReadOnlyFixtureDir { |
| 382 | fn new(path: &Path) -> Result<Self> { |
| 383 | use std::os::unix::fs::PermissionsExt; |
| 384 | fs::set_permissions(path, fs::Permissions::from_mode(0o500))?; |
| 385 | Ok(Self(path.to_path_buf())) |
| 386 | } |
| 387 | } |
| 388 | #[cfg(unix)] |
| 389 | impl Drop for ReadOnlyFixtureDir { |
| 390 | fn drop(&mut self) { |
| 391 | use std::os::unix::fs::PermissionsExt; |
| 392 | let _ = fs::set_permissions(&self.0, fs::Permissions::from_mode(0o700)); |
| 393 | } |
| 394 | } |
| 395 | |
| 396 | #[cfg(unix)] |
| 397 | #[tokio::test] |
| 398 | async fn failed_running_persistence_leaves_queued_work_without_executor_poll() -> Result<()> { |
| 399 | let root = tempfile::tempdir()?; |
| 400 | let calls = Arc::new(AtomicUsize::new(0)); |
| 401 | let owner = TaskManager::start_with_executor_in_scope( |
| 402 | config(root.path(), "A"), |
| 403 | Arc::new(CountExecutor(calls.clone())), |
| 404 | "A", |
| 405 | ) |
| 406 | .await?; |
| 407 | stop_idle_workers(&owner).await; |
| 408 | let task = owner |
| 409 | .add_task(NewTaskRequest::from_prompt("claim write failure")) |
| 410 | .await?; |
| 411 | let blocked = ReadOnlyFixtureDir::new(&owner.tasks_dir)?; |
| 412 | assert!(owner.claim_next_task().await.is_err()); |
| 413 | assert_eq!(calls.load(Ordering::SeqCst), 0); |
| 414 | assert_eq!(owner.get_task(&task.id).await?.status, TaskStatus::Queued); |
| 415 | drop(blocked); |
| 416 | let (id, request, cancel) = owner.claim_next_task().await?.context("retry claim")?; |
| 417 | owner.run_task(id, request, cancel).await; |
| 418 | assert_eq!(calls.load(Ordering::SeqCst), 1); |
| 419 | owner.shutdown_and_wait().await?; |
| 420 | Ok(()) |
| 421 | } |
| 422 | |
| 423 | #[cfg(unix)] |
| 424 | #[tokio::test] |
| 425 | async fn failed_event_persistence_is_bounded_and_retains_each_accepted_delta_once() -> Result<()> { |
| 426 | let root = tempfile::tempdir()?; |
| 427 | let owner = manager(root.path(), "A").await?; |
| 428 | stop_idle_workers(&owner).await; |
| 429 | let task = owner |
| 430 | .add_task(NewTaskRequest::from_prompt("bounded write failure")) |
| 431 | .await?; |
| 432 | let (_, request, cancel) = owner.claim_next_task().await?.context("claim")?; |
| 433 | let blocked = ReadOnlyFixtureDir::new(&owner.tasks_dir)?; |
| 434 | for seq in 0..TASK_EVENT_CHANNEL_CAPACITY { |
| 435 | owner |
| 436 | .apply_execution_event( |
| 437 | &task.id, |
| 438 | TaskExecutionEvent::RuntimeEvent { |
| 439 | seq: seq as u64, |
| 440 | event: "fixture".into(), |
| 441 | summary: seq.to_string(), |
| 442 | }, |
| 443 | ) |
| 444 | .await?; |
| 445 | } |
| 446 | let rejected = TaskExecutionEvent::RuntimeEvent { |
| 447 | seq: 999, |
| 448 | event: "fixture".into(), |
| 449 | summary: "after recovery".into(), |
| 450 | }; |
| 451 | assert!( |
| 452 | owner |
| 453 | .apply_execution_event(&task.id, rejected.clone()) |
| 454 | .await |
| 455 | .is_err() |
| 456 | ); |
| 457 | assert_eq!( |
| 458 | owner.state.lock().await.pending_events[&task.id].len(), |
| 459 | TASK_EVENT_CHANNEL_CAPACITY |
| 460 | ); |
| 461 | assert!( |
| 462 | cancel.is_cancelled(), |
| 463 | "failed persistence requests actual-owner cancellation" |
| 464 | ); |
| 465 | drop(blocked); |
| 466 | owner.apply_execution_event(&task.id, rejected).await?; |
| 467 | owner.flush_task(&task.id).await?; |
| 468 | assert_eq!( |
| 469 | owner.get_task(&task.id).await?.runtime_event_count, |
| 470 | TASK_EVENT_CHANNEL_CAPACITY + 1 |
| 471 | ); |
| 472 | assert!( |
| 473 | !owner |
| 474 | .state |
| 475 | .lock() |
| 476 | .await |
| 477 | .pending_events |
| 478 | .contains_key(&task.id) |
| 479 | ); |
| 480 | owner |
| 481 | .finish_task( |
| 482 | &task.id, |
| 483 | TaskExecutionResult::from_reason(TaskTerminalReason::Canceled, None), |
| 484 | cancel, |
| 485 | &request.mode_label, |
| 486 | ) |
| 487 | .await?; |
| 488 | owner.shutdown_and_wait().await?; |
| 489 | Ok(()) |
| 490 | } |
| 491 | |
| 492 | #[tokio::test] |
| 493 | async fn legacy_queued_and_running_records_remain_visible_unverified_and_byte_identical() |
| 494 | -> Result<()> { |
| 495 | let root = tempfile::tempdir()?; |
| 496 | let owner = manager(root.path(), "A").await?; |
| 497 | stop_idle_workers(&owner).await; |
| 498 | let first = owner |
| 499 | .add_task(NewTaskRequest { |
| 500 | owner_session_id: Some("legacy-owner".into()), |
| 501 | ..NewTaskRequest::from_prompt("legacy running") |
| 502 | }) |
| 503 | .await?; |
| 504 | owner.claim_next_task().await?.context("claim")?; |
| 505 | let second = owner |
| 506 | .add_task(NewTaskRequest::from_prompt("legacy queued")) |
| 507 | .await?; |
| 508 | let mut snapshots = Vec::new(); |
| 509 | for id in [&first.id, &second.id] { |
| 510 | let path = owner.tasks_dir.join(format!("{id}.json")); |
| 511 | let mut legacy: Value = serde_json::from_slice(&fs::read(&path)?)?; |
| 512 | let object = legacy.as_object_mut().unwrap(); |
| 513 | object.remove("execution_scope"); |
| 514 | object.remove("execution_generation"); |
| 515 | object.remove("cancel_requested_seq"); |
| 516 | object.insert("schema_version".into(), Value::from(3)); |
| 517 | let bytes = serde_json::to_vec(&legacy)?; |
| 518 | fs::write(&path, &bytes)?; |
| 519 | snapshots.push((path, bytes)); |
| 520 | } |
| 521 | let other = manager(root.path(), "B").await?; |
| 522 | sleep(STORE_REFRESH_INTERVAL * 2).await; |
| 523 | let rows = other.list_tasks(None).await?; |
| 524 | assert_eq!(rows.len(), 2); |
| 525 | assert!(rows.iter().all(|row| !row.execution_binding_known)); |
| 526 | for (path, bytes) in snapshots { |
| 527 | assert_eq!(fs::read(path)?, bytes); |
| 528 | } |
| 529 | assert_eq!( |
| 530 | other |
| 531 | .get_task_for_owner(&first.id, "legacy-owner") |
| 532 | .await? |
| 533 | .status, |
| 534 | TaskStatus::Running |
| 535 | ); |
| 536 | assert!(executions(root.path(), "B").is_empty()); |
| 537 | owner.shutdown_and_wait().await?; |
| 538 | other.shutdown_and_wait().await?; |
| 539 | Ok(()) |
| 540 | } |
| 541 | |
| 542 | #[tokio::test] |
| 543 | async fn unreadable_owner_before_first_poll_waits_for_storage_without_executing() -> Result<()> { |
| 544 | for already_canceled in [false, true] { |
| 545 | let root = tempfile::tempdir()?; |
| 546 | let calls = Arc::new(AtomicUsize::new(0)); |
| 547 | let owner = TaskManager::start_with_executor_in_scope( |
| 548 | config(root.path(), "A"), |
| 549 | Arc::new(CountExecutor(calls.clone())), |
| 550 | "A", |
| 551 | ) |
| 552 | .await?; |
| 553 | stop_idle_workers(&owner).await; |
| 554 | let task = owner |
| 555 | .add_task(NewTaskRequest::from_prompt("unavailable before poll")) |
| 556 | .await?; |
| 557 | let (id, request, cancel) = owner.claim_next_task().await?.context("claim")?; |
| 558 | if already_canceled { |
| 559 | cancel.cancel(); |
| 560 | } |
| 561 | let path = owner.tasks_dir.join(format!("{}.json", task.id)); |
| 562 | let stored = fs::read(&path)?; |
| 563 | fs::write(&path, b"{corrupt")?; |
| 564 | let running = tokio::spawn({ |
| 565 | let owner = owner.clone(); |
| 566 | async move { owner.run_task(id, request, cancel).await } |
| 567 | }); |
| 568 | sleep(STORE_REFRESH_INTERVAL * 2).await; |
| 569 | assert_eq!(calls.load(Ordering::SeqCst), 0); |
| 570 | assert!( |
| 571 | !running.is_finished(), |
| 572 | "unpersisted terminal state is still pending recovery" |
| 573 | ); |
| 574 | fs::write(&path, stored)?; |
| 575 | tokio::time::timeout(Duration::from_secs(5), running).await??; |
| 576 | assert_eq!(owner.get_task(&task.id).await?.status, TaskStatus::Canceled); |
| 577 | assert_eq!(calls.load(Ordering::SeqCst), 0); |
| 578 | owner.shutdown_and_wait().await?; |
| 579 | } |
| 580 | Ok(()) |
| 581 | } |
| 582 |