From d25c114d494ddb693290b76bf5e5f64ecbdb38fc Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Fri, 2 Oct 2026 03:23:11 +0000 Subject: [PATCH] Bound in-flight file opens in exec-server (#50162) ## Why The per-connection limit of 128 file read handles was checked only after opening the file, leaving pending opens outside the limit. ## What changed Reserve a semaphore permit before polling the file-open future, counting pending opens and registered handles against the same limit. Release capacity when an open fails or is cancelled, or when a handle closes, even if a read still retains the file. Use synchronous handle-table locks that never cross an await, and recheck duplicate IDs after opening to preserve the winning handle during concurrent opens. ## Testing Add tests for stalled-open progress, pending-open admission limits, permit cleanup after failure or cancellation, capacity reuse while a file remains in use, and concurrent duplicate IDs. GitOrigin-RevId: caabf38c74dda5951f2fc75234a7149daaa01263 --- codex-rs/exec-server/src/file_handle.rs | 90 +++++++++--- codex-rs/exec-server/src/file_handle_tests.rs | 134 ++++++++++++++++++ codex-rs/exec-server/src/local_file_system.rs | 18 +-- .../exec-server/src/sandboxed_file_system.rs | 4 +- .../src/server/file_system_handler.rs | 15 +- 5 files changed, 217 insertions(+), 44 deletions(-) create mode 100644 codex-rs/exec-server/src/file_handle_tests.rs diff --git a/codex-rs/exec-server/src/file_handle.rs b/codex-rs/exec-server/src/file_handle.rs index 1e5789884c..dc1011b28f 100644 --- a/codex-rs/exec-server/src/file_handle.rs +++ b/codex-rs/exec-server/src/file_handle.rs @@ -1,10 +1,18 @@ +//! Connection-scoped file handles with bounded positional reads. +//! Semaphore permits bound open handles and in-flight opens; table locks never cross an await. + use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::fs::File; +use std::future::Future; use std::io; use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; use codex_file_system::FILE_READ_CHUNK_SIZE; -use tokio::sync::Mutex; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; const MAX_OPEN_FILES: usize = 128; @@ -14,32 +22,51 @@ pub(crate) struct FileReadBlock { pub(crate) eof: bool, } -#[derive(Clone, Default)] +#[derive(Clone)] pub(crate) struct FileHandleManager { - handles: Arc>>>, + handles: Arc>>, + slots: Arc, +} + +impl Default for FileHandleManager { + fn default() -> Self { + Self { + handles: Arc::default(), + slots: Arc::new(Semaphore::new(MAX_OPEN_FILES)), + } + } } impl FileHandleManager { pub(crate) async fn open( &self, handle_id: String, - file: tokio::fs::File, + open_file: impl Future>, ) -> io::Result { - let file = Arc::new(file.into_std().await); - let mut handles = self.handles.lock().await; - if handles.contains_key(&handle_id) { + if self.lock_handles().contains_key(&handle_id) { return Err(io::Error::new( io::ErrorKind::InvalidInput, format!("file read handle `{handle_id}` already exists"), )); } - if handles.len() >= MAX_OPEN_FILES { - return Err(io::Error::new( + let permit = Arc::clone(&self.slots).try_acquire_owned().map_err(|_| { + io::Error::new( io::ErrorKind::InvalidInput, format!("at most {MAX_OPEN_FILES} file reads may be open per connection"), + ) + })?; + let file = Arc::new(open_file.await?.into_std().await); + let mut handles = self.lock_handles(); + let Entry::Vacant(entry) = handles.entry(handle_id.clone()) else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("file read handle `{handle_id}` already exists"), )); - } - handles.insert(handle_id.clone(), file); + }; + entry.insert(FileHandleEntry { + file, + _permit: permit, + }); Ok(handle_id) } @@ -50,13 +77,7 @@ impl FileHandleManager { len: usize, ) -> io::Result { validate_read_block_len(len)?; - let file = { - let handles = self.handles.lock().await; - handles - .get(handle_id) - .cloned() - .ok_or_else(|| unknown_handle_error(handle_id))? - }; + let file = self.get(handle_id)?; let result = match tokio::task::spawn_blocking(move || read_block_at(&file, offset, len)).await { Ok(result) => result, @@ -65,18 +86,37 @@ impl FileHandleManager { ))), }; if result.is_err() { - self.close(handle_id).await; + self.close(handle_id); } result } - pub(crate) async fn close(&self, handle_id: &str) { - self.handles.lock().await.remove(handle_id); + fn get(&self, handle_id: &str) -> io::Result> { + self.lock_handles() + .get(handle_id) + .map(|entry| Arc::clone(&entry.file)) + .ok_or_else(|| unknown_handle_error(handle_id)) } - pub(crate) async fn close_all(&self) { - self.handles.lock().await.clear(); + pub(crate) fn close(&self, handle_id: &str) { + self.lock_handles().remove(handle_id); } + + pub(crate) fn close_all(&self) { + self.lock_handles().clear(); + } + + fn lock_handles(&self) -> MutexGuard<'_, HashMap> { + self.handles + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +struct FileHandleEntry { + file: Arc, + // Closing an entry releases capacity even if a read still holds the file. + _permit: OwnedSemaphorePermit, } fn read_block_at(file: &File, offset: u64, len: usize) -> io::Result { @@ -126,3 +166,7 @@ fn unknown_handle_error(handle_id: &str) -> io::Error { format!("unknown file read handle `{handle_id}`"), ) } + +#[cfg(test)] +#[path = "file_handle_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/file_handle_tests.rs b/codex-rs/exec-server/src/file_handle_tests.rs new file mode 100644 index 0000000000..35a0aca1d8 --- /dev/null +++ b/codex-rs/exec-server/src/file_handle_tests.rs @@ -0,0 +1,134 @@ +//! Exercises permit cleanup and progress without relying on I/O timing. + +use std::future::pending; +use std::future::ready; +use std::io; +use std::io::Write; +use std::task::Poll; + +use anyhow::Result; +use pretty_assertions::assert_eq; +use tokio::fs::File; +use tokio::sync::oneshot; + +use super::FileHandleManager; +use super::FileReadBlock; +use super::MAX_OPEN_FILES; + +/// A stalled open must let unrelated handles open and read without waiting for it. +#[tokio::test] +async fn stalled_open_does_not_block_other_handles() -> Result<()> { + let manager = FileHandleManager::default(); + let mut stalled = Box::pin(manager.open("stalled".to_string(), pending())); + assert!(futures::poll!(stalled.as_mut()).is_pending()); + + let mut other = Box::pin(manager.open( + "other".to_string(), + ready(Ok(File::from_std(tempfile::tempfile()?))), + )); + let Poll::Ready(result) = futures::poll!(other.as_mut()) else { + anyhow::bail!("an unrelated open was blocked by the stalled open"); + }; + assert_eq!(result?, "other"); + manager.read_block("other", /*offset*/ 0, /*len*/ 1).await?; + Ok(()) +} + +/// Pending opens consume capacity before polling another file open. +#[tokio::test] +async fn pending_opens_reserve_capacity() -> Result<()> { + let manager = FileHandleManager::default(); + let mut opens = Vec::new(); + for index in 0..MAX_OPEN_FILES { + let mut open = Box::pin(manager.open(index.to_string(), pending())); + assert!(futures::poll!(open.as_mut()).is_pending()); + opens.push(open); + } + let mut rejected = Box::pin(manager.open("overflow".to_string(), pending())); + let Poll::Ready(Err(error)) = futures::poll!(rejected.as_mut()) else { + anyhow::bail!("admission must fail before polling the file open"); + }; + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); + Ok(()) +} + +/// Failed and cancelled opens must return their permits for reuse. +#[test_case::test_case(true; "cancelled")] +#[test_case::test_case(false; "failed")] +#[tokio::test] +async fn unfinished_opens_release_capacity(cancelled: bool) -> Result<()> { + let manager = FileHandleManager::default(); + for _ in 0..=MAX_OPEN_FILES { + let mut open = Box::pin(manager.open("reused".to_string(), async { + if cancelled { + pending::<()>().await; + } + Err(io::Error::other("open failed")) + })); + let result = futures::poll!(open.as_mut()); + if cancelled { + assert!(result.is_pending()); + } else { + let Poll::Ready(Err(error)) = result else { + anyhow::bail!("the open error must reach the caller"); + }; + assert_eq!(error.kind(), io::ErrorKind::Other); + } + drop(open); + } + Ok(()) +} + +/// Closing an entry frees capacity even while an operation retains its file descriptor. +#[tokio::test] +async fn close_releases_capacity_while_the_file_is_still_in_use() -> Result<()> { + let manager = FileHandleManager::default(); + let file = tempfile::tempfile()?; + for index in 0..MAX_OPEN_FILES { + manager + .open( + index.to_string(), + ready(Ok(File::from_std(file.try_clone()?))), + ) + .await?; + } + let retained = manager.get("0")?; + manager.close("0"); + manager + .open("next".to_string(), ready(Ok(File::from_std(file)))) + .await?; + drop(retained); + Ok(()) +} + +/// Concurrent opens of one ID must reject the loser without changing the winning file. +#[tokio::test] +async fn concurrent_opens_reject_duplicate_ids_without_replacing_the_file() -> Result<()> { + let manager = FileHandleManager::default(); + let (send, receive) = oneshot::channel(); + let mut late = Box::pin(manager.open("shared".to_string(), async { + receive.await.map_err(io::Error::other) + })); + assert!(futures::poll!(late.as_mut()).is_pending()); + + let mut file = tempfile::tempfile()?; + file.write_all(b"first")?; + manager + .open("shared".to_string(), ready(Ok(File::from_std(file)))) + .await?; + send.send(File::from_std(tempfile::tempfile()?)) + .map_err(|_| anyhow::anyhow!("the pending open dropped its receiver"))?; + + let error = late.await.expect_err("the duplicate open must fail"); + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); + assert_eq!( + manager + .read_block("shared", /*offset*/ 0, /*len*/ 6) + .await?, + FileReadBlock { + bytes: b"first".to_vec(), + eof: true, + } + ); + Ok(()) +} diff --git a/codex-rs/exec-server/src/local_file_system.rs b/codex-rs/exec-server/src/local_file_system.rs index a95f60f851..7e3c7ef899 100644 --- a/codex-rs/exec-server/src/local_file_system.rs +++ b/codex-rs/exec-server/src/local_file_system.rs @@ -118,7 +118,7 @@ impl LocalFileSystem { } impl LocalFileSystem { - pub(crate) async fn open_file_for_read( + pub(crate) async fn open_file( &self, path: &PathUri, sandbox: Option<&FileSystemSandboxContext>, @@ -127,11 +127,9 @@ impl LocalFileSystem { sandbox.validate_file_system_paths_for_current_host()?; } if sandbox.is_some_and(FileSystemSandboxContext::should_read_from_sandbox) { - return self.sandboxed()?.open_file_for_read(path, sandbox).await; + return self.sandboxed()?.open_file(path, sandbox).await; } - self.unsandboxed - .open_file_for_read(path, /*sandbox*/ None) - .await + self.unsandboxed.open_file(path, /*sandbox*/ None).await } async fn canonicalize( @@ -350,15 +348,13 @@ impl ExecutorFileSystem for LocalFileSystem { } impl UnsandboxedFileSystem { - async fn open_file_for_read( + async fn open_file( &self, path: &PathUri, sandbox: Option<&FileSystemSandboxContext>, ) -> FileSystemResult { reject_platform_sandbox_context(sandbox)?; - self.file_system - .open_file_for_read(path, /*sandbox*/ None) - .await + self.file_system.open_file(path, /*sandbox*/ None).await } async fn canonicalize( @@ -581,7 +577,7 @@ impl ExecutorFileSystem for UnsandboxedFileSystem { } impl DirectFileSystem { - async fn open_file_for_read( + async fn open_file( &self, path: &PathUri, sandbox: Option<&FileSystemSandboxContext>, @@ -619,7 +615,7 @@ impl DirectFileSystem { path: &PathUri, sandbox: Option<&FileSystemSandboxContext>, ) -> FileSystemResult { - let file = self.open_file_for_read(path, sandbox).await?; + let file = self.open_file(path, sandbox).await?; Ok(FileSystemReadStream::new(ReaderStream::with_capacity( file, FILE_READ_CHUNK_SIZE, diff --git a/codex-rs/exec-server/src/sandboxed_file_system.rs b/codex-rs/exec-server/src/sandboxed_file_system.rs index 858359444c..24b8a240f1 100644 --- a/codex-rs/exec-server/src/sandboxed_file_system.rs +++ b/codex-rs/exec-server/src/sandboxed_file_system.rs @@ -80,7 +80,7 @@ impl SandboxedFileSystem { .map_err(map_sandbox_error) } - pub(crate) async fn open_file_for_read( + pub(crate) async fn open_file( &self, path: &PathUri, sandbox: Option<&FileSystemSandboxContext>, @@ -367,7 +367,7 @@ impl ExecutorFileSystem for SandboxedFileSystem { sandbox: Option<&'a FileSystemSandboxContext>, ) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> { Box::pin(async move { - let file = self.open_file_for_read(path, sandbox).await?; + let file = self.open_file(path, sandbox).await?; Ok(FileSystemReadStream::new(ReaderStream::with_capacity( file, FILE_READ_CHUNK_SIZE, diff --git a/codex-rs/exec-server/src/server/file_system_handler.rs b/codex-rs/exec-server/src/server/file_system_handler.rs index 997c9558da..0d416a0bcd 100644 --- a/codex-rs/exec-server/src/server/file_system_handler.rs +++ b/codex-rs/exec-server/src/server/file_system_handler.rs @@ -70,7 +70,7 @@ impl FileSystemHandler { } pub(crate) async fn shutdown(&self) { - self.file_handles.close_all().await; + self.file_handles.close_all(); } pub(crate) async fn discover_capability_roots( @@ -131,14 +131,13 @@ impl FileSystemHandler { "exec-server does not support writable file streams".to_string(), )); } - let file = self - .file_system - .open_file_for_read(¶ms.path, params.sandbox.as_ref()) - .await - .map_err(map_fs_error)?; let handle_id = self .file_handles - .open(params.handle_id, file) + .open( + params.handle_id, + self.file_system + .open_file(¶ms.path, params.sandbox.as_ref()), + ) .await .map_err(map_fs_error)?; Ok(FsOpenResponse { handle_id }) @@ -176,7 +175,7 @@ impl FileSystemHandler { params: FsCloseParams, ) -> Result { validate_file_handle_id(¶ms.handle_id)?; - self.file_handles.close(¶ms.handle_id).await; + self.file_handles.close(¶ms.handle_id); Ok(FsCloseResponse {}) }