mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
Track app-server running turns incrementally (#49084)
## Why Every thread state mutation scanned all tracked runtimes to recount running turns while holding the state lock. This count supports graceful restart draining introduced in https://github.com/openai/codex/pull/12517. ## What changed Maintain `running_turn_count` in `ThreadWatchState`, adjusting it when a runtime's `running` flag changes or a running thread is removed. Publish the cached count while preserving notifications only when the count changes. ## Testing Extend the running-turn watch test to cover duplicate starts, two simultaneous running threads, and removal of a running thread, with counts checked against the runtime map. GitOrigin-RevId: e7d7ebb33b9ff4d24dc70d8f01ae981378fed5f8
This commit is contained in:
committed by
copyberry
parent
46d2585ea4
commit
88e9a8329d
@@ -227,11 +227,7 @@ impl ThreadWatchManager {
|
||||
let notification = {
|
||||
let mut state = self.state.lock().await;
|
||||
let notification = mutate(&mut state);
|
||||
let running_turn_count = state
|
||||
.runtime_by_thread_id
|
||||
.values()
|
||||
.filter(|runtime| runtime.running)
|
||||
.count();
|
||||
let running_turn_count = state.running_turn_count;
|
||||
self.running_turn_count_tx.send_if_modified(|current| {
|
||||
if *current == running_turn_count {
|
||||
false
|
||||
@@ -311,6 +307,7 @@ pub(crate) fn resolve_thread_status(
|
||||
struct ThreadWatchState {
|
||||
runtime_by_thread_id: HashMap<String, RuntimeFacts>,
|
||||
status_watcher_by_thread_id: HashMap<String, watch::Sender<ThreadStatus>>,
|
||||
running_turn_count: usize,
|
||||
}
|
||||
|
||||
impl ThreadWatchState {
|
||||
@@ -335,7 +332,13 @@ impl ThreadWatchState {
|
||||
|
||||
fn remove_thread(&mut self, thread_id: &str) -> Option<ThreadStatusChangedNotification> {
|
||||
let previous_status = self.status_for(thread_id);
|
||||
self.runtime_by_thread_id.remove(thread_id);
|
||||
if self
|
||||
.runtime_by_thread_id
|
||||
.remove(thread_id)
|
||||
.is_some_and(|runtime| runtime.running)
|
||||
{
|
||||
self.running_turn_count -= 1;
|
||||
}
|
||||
self.update_status_watcher(thread_id, &ThreadStatus::NotLoaded);
|
||||
if previous_status.is_some() && previous_status != Some(ThreadStatus::NotLoaded) {
|
||||
Some(ThreadStatusChangedNotification {
|
||||
@@ -360,8 +363,14 @@ impl ThreadWatchState {
|
||||
.runtime_by_thread_id
|
||||
.entry(thread_id.to_string())
|
||||
.or_default();
|
||||
let was_running = runtime.running;
|
||||
runtime.is_loaded = true;
|
||||
mutate(runtime);
|
||||
match (was_running, runtime.running) {
|
||||
(false, true) => self.running_turn_count += 1,
|
||||
(true, false) => self.running_turn_count -= 1,
|
||||
(false, false) | (true, true) => {}
|
||||
}
|
||||
self.update_status_watcher_for_thread(thread_id);
|
||||
self.status_changed_notification(thread_id.to_string(), previous_status)
|
||||
}
|
||||
@@ -712,11 +721,24 @@ mod tests {
|
||||
assert!(count.has_changed().expect("watch remains open"));
|
||||
assert_eq!(*count.borrow_and_update(), 1);
|
||||
|
||||
manager.note_turn_started(INTERACTIVE_THREAD_ID).await;
|
||||
assert!(!count.has_changed().expect("watch remains open"));
|
||||
|
||||
manager.note_turn_started(NON_INTERACTIVE_THREAD_ID).await;
|
||||
assert!(count.has_changed().expect("watch remains open"));
|
||||
assert_eq!(*count.borrow_and_update(), 2);
|
||||
assert_eq!(manager.running_turn_count().await, 2);
|
||||
|
||||
let _permission_guard = manager
|
||||
.note_permission_requested(INTERACTIVE_THREAD_ID)
|
||||
.await;
|
||||
assert!(!count.has_changed().expect("watch remains open"));
|
||||
|
||||
manager.remove_thread(NON_INTERACTIVE_THREAD_ID).await;
|
||||
assert!(count.has_changed().expect("watch remains open"));
|
||||
assert_eq!(*count.borrow_and_update(), 1);
|
||||
assert_eq!(manager.running_turn_count().await, 1);
|
||||
|
||||
manager.note_thread_shutdown(INTERACTIVE_THREAD_ID).await;
|
||||
assert!(count.has_changed().expect("watch remains open"));
|
||||
assert_eq!(*count.borrow_and_update(), 0);
|
||||
|
||||
Reference in New Issue
Block a user