mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Shut down the in-process outbound router explicitly (#35523)
## Why Detached processor work can retain outgoing message senders, so waiting for channel closure can keep the in-process app-server's outbound router alive during shutdown. ## What changed Add an explicit shutdown signal for the outbound router and prioritize it over pending outgoing messages. Send that signal after processor tasks are stopped so runtime shutdown no longer depends on every sender being dropped. ## Testing Add a paused-time test that verifies the router exits and closes its channel even when an outgoing sender is still retained. GitOrigin-RevId: a8f88d88a0950f50339b1c56909fedb781f178d4
This commit is contained in:
@@ -383,6 +383,25 @@ pub async fn start(mut args: InProcessStartArgs) -> IoResult<InProcessClientHand
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
async fn run_outbound_router(
|
||||
mut outgoing_rx: mpsc::Receiver<OutgoingEnvelope>,
|
||||
mut outbound_connections: HashMap<ConnectionId, OutboundConnectionState>,
|
||||
mut shutdown_rx: oneshot::Receiver<()>,
|
||||
) {
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = &mut shutdown_rx => break,
|
||||
envelope = outgoing_rx.recv() => {
|
||||
let Some(envelope) = envelope else {
|
||||
break;
|
||||
};
|
||||
route_outgoing_envelope(&mut outbound_connections, envelope).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
|
||||
let channel_capacity = args.channel_capacity.max(1);
|
||||
let installation_id = resolve_installation_id(&args.config.codex_home).await?;
|
||||
@@ -390,7 +409,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
|
||||
let (event_tx, event_rx) = mpsc::channel::<InProcessServerEvent>(channel_capacity);
|
||||
|
||||
let runtime_handle = tokio::spawn(async move {
|
||||
let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<OutgoingEnvelope>(channel_capacity);
|
||||
let (outgoing_tx, outgoing_rx) = mpsc::channel::<OutgoingEnvelope>(channel_capacity);
|
||||
let auth_manager =
|
||||
AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env)
|
||||
.await;
|
||||
@@ -418,11 +437,12 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
let mut outbound_handle = tokio::spawn(async move {
|
||||
while let Some(envelope) = outgoing_rx.recv().await {
|
||||
route_outgoing_envelope(&mut outbound_connections, envelope).await;
|
||||
}
|
||||
});
|
||||
let (outbound_shutdown_tx, outbound_shutdown_rx) = oneshot::channel();
|
||||
let mut outbound_handle = tokio::spawn(run_outbound_router(
|
||||
outgoing_rx,
|
||||
outbound_connections,
|
||||
outbound_shutdown_rx,
|
||||
));
|
||||
|
||||
let processor_outgoing = Arc::clone(&outgoing_message_sender);
|
||||
let config_manager = ConfigManager::new(
|
||||
@@ -709,8 +729,8 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
|
||||
"in-process app-server runtime is shutting down",
|
||||
)))
|
||||
.await;
|
||||
// Drop the runtime's last sender before awaiting the router task so
|
||||
// `outgoing_rx.recv()` can observe channel closure and exit cleanly.
|
||||
// Detached processor work can retain outgoing senders, so channel
|
||||
// closure alone cannot be used to shut down the outbound router.
|
||||
drop(outgoing_message_sender);
|
||||
for (_, response_tx) in pending_request_responses {
|
||||
let _ = response_tx.send(Err(internal_error(
|
||||
@@ -722,6 +742,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
|
||||
processor_handle.abort();
|
||||
let _ = processor_handle.await;
|
||||
}
|
||||
let _ = outbound_shutdown_tx.send(());
|
||||
if let Err(_elapsed) = timeout(SHUTDOWN_TIMEOUT, &mut outbound_handle).await {
|
||||
outbound_handle.abort();
|
||||
let _ = outbound_handle.await;
|
||||
@@ -896,6 +917,30 @@ mod tests {
|
||||
.expect("in-process runtime should shutdown cleanly");
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn in_process_outbound_router_shutdown_does_not_wait_for_retained_sender() {
|
||||
let (outgoing_tx, outgoing_rx) = mpsc::channel(/*buffer*/ 1);
|
||||
let retained_outgoing_tx = outgoing_tx.clone();
|
||||
drop(outgoing_tx);
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
let mut outbound_handle = tokio::spawn(run_outbound_router(
|
||||
outgoing_rx,
|
||||
HashMap::new(),
|
||||
shutdown_rx,
|
||||
));
|
||||
|
||||
assert!(!retained_outgoing_tx.is_closed());
|
||||
shutdown_tx
|
||||
.send(())
|
||||
.expect("outbound router should accept explicit shutdown");
|
||||
timeout(SHUTDOWN_TIMEOUT, &mut outbound_handle)
|
||||
.await
|
||||
.expect("outbound router should not wait for its retained sender")
|
||||
.expect("outbound router should complete successfully");
|
||||
assert!(retained_outgoing_tx.is_closed());
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn in_process_shutdown_waits_for_analytics_flush_budget() {
|
||||
let (client_tx, mut client_rx) = mpsc::channel(/*buffer*/ 1);
|
||||
|
||||
Reference in New Issue
Block a user