diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index a1a57d50fd1..31a87a944f8 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -528,6 +528,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult, pub(super) thread_state_manager: ThreadStateManager, pub(super) outgoing: Arc, - pub(super) pending_thread_unloads: Arc>>, + pub(super) pending_thread_unloads: PendingThreadUnloads, pub(super) thread_watch_manager: ThreadWatchManager, pub(super) thread_list_state_permit: Arc, pub(super) fallback_model_provider: String, @@ -124,41 +124,45 @@ impl UnloadingState { clippy::await_holding_invalid_type, reason = "idle unload admission holds both lifecycle guards through its final rechecks" )] -async fn claim_idle_unload_if_permitted( +async fn try_start_idle_thread_unload( unloading_state: &mut UnloadingState, thread_list_state_permit: Arc, - pending_thread_unloads: &Arc>>, + pending_thread_unloads: &PendingThreadUnloads, thread_state_manager: &ThreadStateManager, conversation_id: ThreadId, is_running: impl std::future::Future, -) -> bool { + teardown: F, +) -> Option +where + F: std::future::Future + Send + 'static, +{ let Ok(permit) = thread_list_state_permit.acquire_owned().await else { - return false; + return None; }; - let mut pending = pending_thread_unloads.lock().await; - if pending.contains(&conversation_id) { - return false; + let admission = pending_thread_unloads.lock_admission().await; + if pending_thread_unloads.contains(conversation_id) { + return None; } if !unloading_state.should_unload_now() { - return false; + return None; } if is_running.await { unloading_state.note_thread_activity_observed(); - return false; + return None; } if thread_state_manager .has_connections_or_reservations(conversation_id) .await { - return false; + return None; } if !unloading_state.should_unload_now() { - return false; + return None; } - pending.insert(conversation_id); - drop(pending); + let start = pending_thread_unloads.try_start(conversation_id, teardown); + drop(admission); drop(permit); - true + Some(start) } #[cfg(test)] @@ -199,8 +203,14 @@ pub(super) async fn ensure_conversation_listener( } }; let thread_state = { - let pending_thread_unloads = listener_task_context.pending_thread_unloads.lock().await; - if pending_thread_unloads.contains(&conversation_id) { + let _admission = listener_task_context + .pending_thread_unloads + .lock_admission() + .await; + if listener_task_context + .pending_thread_unloads + .contains(conversation_id) + { return Err(invalid_request(format!( "thread {conversation_id} is closing; retry after the thread is closed" ))); @@ -305,8 +315,8 @@ pub(super) async fn ensure_listener_task_running( .. } = listener_task_context; - let pending_thread_unloads_guard = pending_thread_unloads.lock().await; - if pending_thread_unloads_guard.contains(&conversation_id) { + let pending_thread_unloads_guard = pending_thread_unloads.lock_admission().await; + if pending_thread_unloads.contains(conversation_id) { return Err(invalid_request(format!( "thread {conversation_id} is closing; retry after the thread is closed" ))); @@ -336,7 +346,7 @@ pub(super) async fn ensure_listener_task_running( }; let listener_command_tx_for_task = listener_command_tx.clone(); let thread_state_for_task = Arc::clone(&thread_state); - let pending_thread_unloads_for_task = Arc::clone(&pending_thread_unloads); + let pending_thread_unloads_for_task = pending_thread_unloads.clone(); let listener_registry = thread_state_manager.clone(); let outgoing_for_task = Arc::clone(&outgoing); let listener_task = async move { @@ -429,7 +439,7 @@ pub(super) async fn ensure_listener_task_running( unloading_state.note_thread_activity_observed(); continue; } - if !claim_idle_unload_if_permitted( + let Some(start) = try_start_idle_thread_unload( &mut unloading_state, thread_list_state_permit.clone(), &pending_thread_unloads, @@ -438,21 +448,26 @@ pub(super) async fn ensure_listener_task_running( async { matches!(conversation.agent_status().await, AgentStatus::Running) }, + unload_thread_without_subscribers( + thread_manager.clone(), + outgoing_for_task.clone(), + thread_state_manager.clone(), + thread_watch_manager.clone(), + conversation_id, + conversation.clone(), + ), ) .await - { + else { continue; + }; + match start { + PendingThreadUnloadStartResult::Started(_completed) => {} + PendingThreadUnloadStartResult::Pending(completed) => { + wait_for_thread_unload(completed).await; + } + PendingThreadUnloadStartResult::Closing => {} } - unload_thread_without_subscribers( - thread_manager.clone(), - outgoing_for_task.clone(), - pending_thread_unloads.clone(), - thread_state_manager.clone(), - thread_watch_manager.clone(), - conversation_id, - conversation.clone(), - ) - .await; break; } } @@ -501,7 +516,6 @@ pub(super) async fn wait_for_thread_shutdown(thread: &Arc) -> Threa pub(super) async fn unload_thread_without_subscribers( thread_manager: Arc, outgoing: Arc, - pending_thread_unloads: Arc>>, thread_state_manager: ThreadStateManager, thread_watch_manager: ThreadWatchManager, thread_id: ThreadId, @@ -514,40 +528,36 @@ pub(super) async fn unload_thread_without_subscribers( outgoing .cancel_requests_for_thread(thread_id, /*error*/ None) .await; - let _ = thread_state_manager.remove_thread_state(thread_id).await; - - tokio::spawn(async move { - match wait_for_thread_shutdown(&thread).await { - ThreadShutdownResult::Complete => { - if thread_manager.remove_thread(&thread_id).await.is_none() { - info!("thread {thread_id} was already removed before teardown finalized"); - thread_watch_manager - .remove_thread(&thread_id.to_string()) - .await; - pending_thread_unloads.lock().await.remove(&thread_id); - return; - } + if let Some(listener) = thread_state_manager.remove_thread_state(thread_id).await { + listener.wait_until_closed().await; + } + + match wait_for_thread_shutdown(&thread).await { + ThreadShutdownResult::Complete => { + if thread_manager.remove_thread(&thread_id).await.is_none() { + info!("thread {thread_id} was already removed before teardown finalized"); thread_watch_manager .remove_thread(&thread_id.to_string()) .await; - let notification = ThreadClosedNotification { - thread_id: thread_id.to_string(), - }; - outgoing - .send_server_notification(ServerNotification::ThreadClosed(notification)) - .await; - pending_thread_unloads.lock().await.remove(&thread_id); - } - ThreadShutdownResult::SubmitFailed => { - pending_thread_unloads.lock().await.remove(&thread_id); - warn!("failed to submit Shutdown to thread {thread_id}"); - } - ThreadShutdownResult::TimedOut => { - pending_thread_unloads.lock().await.remove(&thread_id); - warn!("thread {thread_id} shutdown timed out; leaving thread loaded"); + return; } + thread_watch_manager + .remove_thread(&thread_id.to_string()) + .await; + let notification = ThreadClosedNotification { + thread_id: thread_id.to_string(), + }; + outgoing + .send_server_notification(ServerNotification::ThreadClosed(notification)) + .await; + } + ThreadShutdownResult::SubmitFailed => { + warn!("failed to submit Shutdown to thread {thread_id}"); } - }); + ThreadShutdownResult::TimedOut => { + warn!("thread {thread_id} shutdown timed out; leaving thread loaded"); + } + } } #[allow(clippy::too_many_arguments)] @@ -559,7 +569,7 @@ pub(super) async fn handle_thread_listener_command( thread_state: &Arc>, thread_watch_manager: &ThreadWatchManager, outgoing: &Arc, - pending_thread_unloads: &Arc>>, + pending_thread_unloads: &PendingThreadUnloads, listener_command: ThreadListenerCommand, ) { match listener_command { @@ -623,7 +633,7 @@ pub(super) async fn handle_pending_thread_resume_request( thread_state: &Arc>, thread_watch_manager: &ThreadWatchManager, outgoing: &Arc, - pending_thread_unloads: &Arc>>, + pending_thread_unloads: &PendingThreadUnloads, pending: crate::thread_state::PendingThreadResumeRequest, ) { let active_turn = { @@ -691,11 +701,11 @@ pub(super) async fn handle_pending_thread_resume_request( } } - if pending_thread_unloads - .lock() - .await - .contains(&conversation_id) - { + let unload_pending = { + let _admission = pending_thread_unloads.lock_admission().await; + pending_thread_unloads.contains(conversation_id) + }; + if unload_pending { outgoing .send_error( request_id, diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs index 1ab13e80c73..fec727ce9cb 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle_tests.rs @@ -13,7 +13,7 @@ async fn resume_reservation_prevents_idle_unload_before_listener_commit() { thread_status_rx, is_active: (false, now), }; - let pending = Arc::new(Mutex::new(HashSet::new())); + let pending = PendingThreadUnloads::default(); let thread_state_manager = ThreadStateManager::new(); let thread_id = ThreadId::new(); let connection_id = ConnectionId(1); @@ -32,17 +32,19 @@ async fn resume_reservation_prevents_idle_unload_before_listener_commit() { .expect("live connection should reserve the thread"); drop(resume_permit); assert!( - !claim_idle_unload_if_permitted( + try_start_idle_thread_unload( &mut unloading_state, permit, &pending, &thread_state_manager, thread_id, std::future::ready(false), + std::future::ready(()), ) .await + .is_none() ); - assert!(pending.lock().await.is_empty()); + assert!(!pending.contains(thread_id)); assert!( thread_state_manager @@ -92,6 +94,7 @@ async fn resume_reservation_prevents_idle_unload_before_listener_commit() { .await .is_empty() ); + pending.close_and_wait().await; thread_state_manager.clear_all_listeners().await; } diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 9bc85ba36e8..6834d95e497 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -355,7 +355,7 @@ pub(crate) struct ThreadRequestProcessor { pub(super) config: Arc, pub(super) config_manager: ConfigManager, pub(super) thread_store: Arc, - pub(super) pending_thread_unloads: Arc>>, + pub(super) pending_thread_unloads: PendingThreadUnloads, pub(super) thread_state_manager: ThreadStateManager, pub(super) thread_watch_manager: ThreadWatchManager, pub(super) thread_list_state_permit: Arc, @@ -394,7 +394,7 @@ impl ThreadRequestProcessor { config: Arc, config_manager: ConfigManager, thread_store: Arc, - pending_thread_unloads: Arc>>, + pending_thread_unloads: PendingThreadUnloads, thread_state_manager: ThreadStateManager, thread_watch_manager: ThreadWatchManager, thread_list_state_permit: Arc, @@ -798,7 +798,6 @@ impl ThreadRequestProcessor { } async fn finalize_thread_teardown(&self, thread_id: ThreadId) { - self.pending_thread_unloads.lock().await.remove(&thread_id); self.outgoing .cancel_requests_for_thread(thread_id, /*error*/ None) .await; @@ -867,7 +866,7 @@ impl ThreadRequestProcessor { thread_manager: Arc::clone(&self.thread_manager), thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), - pending_thread_unloads: Arc::clone(&self.pending_thread_unloads), + pending_thread_unloads: self.pending_thread_unloads.clone(), thread_watch_manager: self.thread_watch_manager.clone(), thread_list_state_permit: self.thread_list_state_permit.clone(), fallback_model_provider: self.config.model_provider_id.clone(), @@ -969,7 +968,7 @@ impl ThreadRequestProcessor { thread_manager: Arc::clone(&self.thread_manager), thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), - pending_thread_unloads: Arc::clone(&self.pending_thread_unloads), + pending_thread_unloads: self.pending_thread_unloads.clone(), thread_watch_manager: self.thread_watch_manager.clone(), thread_list_state_permit: self.thread_list_state_permit.clone(), fallback_model_provider: self.config.model_provider_id.clone(), @@ -1027,6 +1026,10 @@ impl ThreadRequestProcessor { self.thread_state_manager.clear_all_listeners().await; } + pub(crate) async fn drain_thread_teardowns(&self) { + self.pending_thread_unloads.close_and_wait().await; + } + pub(crate) async fn shutdown_threads(&self) { let report = self .thread_manager @@ -2648,22 +2651,22 @@ impl ThreadRequestProcessor { app_server_client_version: Option, supports_openai_form_elicitation: bool, ) -> Result<(), JSONRPCErrorError> { - if let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) - && self - .pending_thread_unloads - .lock() - .await - .contains(&thread_id) - { - self.outgoing - .send_error( - request_id, - invalid_request(format!( - "thread {thread_id} is closing; retry thread/resume after the thread is closed" - )), - ) - .await; - return Ok(()); + if let Ok(thread_id) = ThreadId::from_string(¶ms.thread_id) { + let unload_pending = { + let _admission = self.pending_thread_unloads.lock_admission().await; + self.pending_thread_unloads.contains(thread_id) + }; + if unload_pending { + self.outgoing + .send_error( + request_id, + invalid_request(format!( + "thread {thread_id} is closing; retry thread/resume after the thread is closed" + )), + ) + .await; + return Ok(()); + } } if params.sandbox.is_some() && params.permissions.is_some() { diff --git a/codex-rs/app-server/src/request_processors/thread_teardown.rs b/codex-rs/app-server/src/request_processors/thread_teardown.rs new file mode 100644 index 00000000000..fb3c95ceba7 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_teardown.rs @@ -0,0 +1,169 @@ +use super::*; +use std::pin::Pin; +use std::sync::Mutex as StdMutex; + +type ThreadTeardownFuture = Pin + Send + 'static>>; + +#[derive(Clone, Default)] +pub(crate) struct PendingThreadUnloads { + inner: Arc, +} + +#[derive(Default)] +struct PendingThreadUnloadsInner { + admission: Mutex<()>, + registry: StdMutex, + tasks: TaskTracker, +} + +#[derive(Default)] +struct PendingThreadUnloadRegistry { + closing: bool, + entries: HashMap>, +} + +pub(super) enum PendingThreadUnloadClaimResult { + Claimed(PendingThreadUnloadClaim), + Pending(watch::Receiver), + Closing, +} + +pub(super) enum PendingThreadUnloadStartResult { + Started(watch::Receiver), + Pending(watch::Receiver), + Closing, +} + +pub(super) struct PendingThreadUnloadClaim { + start_tx: Option>, + completed: watch::Receiver, +} + +impl PendingThreadUnloadClaim { + pub(super) fn start(mut self, teardown: F) -> watch::Receiver + where + F: std::future::Future + Send + 'static, + { + if let Some(start_tx) = self.start_tx.take() { + let _ = start_tx.send(Box::pin(teardown)); + } + self.completed.clone() + } +} + +struct PendingThreadUnloadOwner { + pending: PendingThreadUnloads, + thread_id: ThreadId, + completed: watch::Sender, +} + +impl Drop for PendingThreadUnloadOwner { + fn drop(&mut self) { + self.pending.release(self.thread_id, &self.completed); + } +} + +impl PendingThreadUnloads { + pub(super) async fn lock_admission(&self) -> tokio::sync::MutexGuard<'_, ()> { + self.inner.admission.lock().await + } + + fn lock_registry(&self) -> std::sync::MutexGuard<'_, PendingThreadUnloadRegistry> { + self.inner + .registry + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + pub(super) fn contains(&self, thread_id: ThreadId) -> bool { + self.lock_registry().entries.contains_key(&thread_id) + } + + pub(super) fn try_claim(&self, thread_id: ThreadId) -> PendingThreadUnloadClaimResult { + let mut registry = self.lock_registry(); + if registry.closing { + return PendingThreadUnloadClaimResult::Closing; + } + if let Some(completed) = registry.entries.get(&thread_id) { + return PendingThreadUnloadClaimResult::Pending(completed.subscribe()); + } + + let (completed, completed_rx) = watch::channel(false); + registry.entries.insert(thread_id, completed.clone()); + let owner = PendingThreadUnloadOwner { + pending: self.clone(), + thread_id, + completed, + }; + let (start_tx, start_rx) = oneshot::channel::(); + let task = self.inner.tasks.spawn(async move { + if let Ok(teardown) = start_rx.await { + teardown.await; + } + drop(owner); + }); + drop(task); + PendingThreadUnloadClaimResult::Claimed(PendingThreadUnloadClaim { + start_tx: Some(start_tx), + completed: completed_rx, + }) + } + + pub(super) fn try_start( + &self, + thread_id: ThreadId, + teardown: F, + ) -> PendingThreadUnloadStartResult + where + F: std::future::Future + Send + 'static, + { + match self.try_claim(thread_id) { + PendingThreadUnloadClaimResult::Claimed(claim) => { + PendingThreadUnloadStartResult::Started(claim.start(teardown)) + } + PendingThreadUnloadClaimResult::Pending(completed) => { + PendingThreadUnloadStartResult::Pending(completed) + } + PendingThreadUnloadClaimResult::Closing => PendingThreadUnloadStartResult::Closing, + } + } + + fn release(&self, thread_id: ThreadId, completed: &watch::Sender) { + let mut registry = self.lock_registry(); + if registry + .entries + .get(&thread_id) + .is_some_and(|current| current.same_channel(completed)) + { + registry.entries.remove(&thread_id); + } + drop(registry); + let _ = completed.send(true); + } + + pub(super) async fn close_and_wait(&self) { + { + let mut registry = self.lock_registry(); + registry.closing = true; + self.inner.tasks.close(); + } + self.inner.tasks.wait().await; + } + + #[cfg(test)] + fn is_empty(&self) -> bool { + self.lock_registry().entries.is_empty() + } +} + +pub(super) async fn wait_for_thread_unload(mut completed: watch::Receiver) { + while !*completed.borrow_and_update() { + if completed.changed().await.is_err() { + break; + } + } +} + +#[cfg(test)] +#[path = "thread_teardown_tests.rs"] +mod tests; diff --git a/codex-rs/app-server/src/request_processors/thread_teardown_tests.rs b/codex-rs/app-server/src/request_processors/thread_teardown_tests.rs new file mode 100644 index 00000000000..7e2f466121d --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_teardown_tests.rs @@ -0,0 +1,59 @@ +use super::*; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +#[tokio::test] +async fn teardown_is_singleflight_and_outlives_its_first_waiter() { + let pending = PendingThreadUnloads::default(); + let provisional_id = ThreadId::new(); + let PendingThreadUnloadClaimResult::Claimed(provisional) = pending.try_claim(provisional_id) + else { + panic!("new thread id should be claimable"); + }; + let completion = provisional.completed.clone(); + drop(provisional); + wait_for_thread_unload(completion).await; + + let thread_id = ThreadId::new(); + let calls = Arc::new(AtomicUsize::new(0)); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let first_calls = Arc::clone(&calls); + let first_completed = match pending.try_start(thread_id, async move { + first_calls.fetch_add(1, Ordering::SeqCst); + let _ = started_tx.send(()); + let _ = release_rx.await; + }) { + PendingThreadUnloadStartResult::Started(completed) => completed, + _ => panic!("first teardown should start"), + }; + let first = tokio::spawn(wait_for_thread_unload(first_completed)); + started_rx.await.expect("teardown should start"); + let second_calls = Arc::clone(&calls); + let second_completed = match pending.try_start(thread_id, async move { + second_calls.fetch_add(1, Ordering::SeqCst); + }) { + PendingThreadUnloadStartResult::Pending(completed) => completed, + _ => panic!("second teardown should join the first"), + }; + let second = tokio::spawn(wait_for_thread_unload(second_completed)); + let pending_for_drain = pending.clone(); + let drain = tokio::spawn(async move { pending_for_drain.close_and_wait().await }); + tokio::task::yield_now().await; + assert!(!drain.is_finished()); + + first.abort(); + assert!( + first + .await + .expect_err("first waiter should be cancelled") + .is_cancelled() + ); + release_tx + .send(()) + .expect("teardown should still be running"); + second.await.expect("second waiter should complete"); + drain.await.expect("teardown drain should complete"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(pending.is_empty()); +} diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index 3ff64cd34e3..63d590f4b08 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -74,7 +74,7 @@ pub(crate) struct TurnRequestProcessor { arg0_paths: Arg0DispatchPaths, config: Arc, config_manager: ConfigManager, - pending_thread_unloads: Arc>>, + pending_thread_unloads: PendingThreadUnloads, thread_state_manager: ThreadStateManager, thread_watch_manager: ThreadWatchManager, thread_list_state_permit: Arc, @@ -130,7 +130,7 @@ impl TurnRequestProcessor { arg0_paths: Arg0DispatchPaths, config: Arc, config_manager: ConfigManager, - pending_thread_unloads: Arc>>, + pending_thread_unloads: PendingThreadUnloads, thread_state_manager: ThreadStateManager, thread_watch_manager: ThreadWatchManager, thread_list_state_permit: Arc, @@ -1413,7 +1413,7 @@ impl TurnRequestProcessor { thread_manager: Arc::clone(&self.thread_manager), thread_state_manager: self.thread_state_manager.clone(), outgoing: Arc::clone(&self.outgoing), - pending_thread_unloads: Arc::clone(&self.pending_thread_unloads), + pending_thread_unloads: self.pending_thread_unloads.clone(), thread_watch_manager: self.thread_watch_manager.clone(), thread_list_state_permit: self.thread_list_state_permit.clone(), fallback_model_provider: self.config.model_provider_id.clone(),