openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reduce unified exec output buffer allocations (#39712)

Browse Source 
## What changed

- Accept borrowed byte slices when appending output chunks, avoiding clones before
  forwarding those chunks to streaming consumers.
- Merge drained head-tail buffers by reusing their owned storage while preserving
  the retained prefix, latest suffix, and omitted-byte count.
- Make unified exec output handles capacity-generic so bounded-output tests can
  exercise repeated drains with small buffers.

## Testing

- Update head-tail buffer and process manager tests to cover chunk retention,
  omission accounting, and repeated output drains with the new APIs.

GitOrigin-RevId: caa9813d3c9d9921cff28888263a212def381640
This commit is contained in:
Tamir Duberstein
2026-08-20 13:17:20 +00:00
committed by copyberry
parent 59f7da58d6
commit 4a942885c8
7 changed files with 118 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/unified_exec/async_watcher.rs
+1 -1
View File
@@ -249,7 +249,7 @@ async fn process_chunk(
while let Some(prefix) = split_valid_utf8_prefix(pending) {
{
let mut guard = transcript.lock().await;
guard.push_chunk(prefix.to_vec());
guard.push_chunk(&prefix);
}
if *emitted_deltas >= MAX_EXEC_OUTPUT_DELTAS_PER_CALL {
codex-rs/core/src/unified_exec/head_tail_buffer.rs
+75 -49
View File
@@ -42,22 +42,52 @@ impl<const MAX_BYTES: usize> HeadTailBuffer<MAX_BYTES> {
/// Bytes are first added to the head until the head budget is full; any
/// remaining bytes are added to the tail, with older tail bytes being
/// dropped to preserve the tail budget.
pub(crate) fn push_chunk(&mut self, chunk: Vec<u8>) {
if chunk.is_empty() {
return;
}
if MAX_BYTES == 0 {
self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len());
return;
}
pub(crate) fn push_chunk(&mut self, chunk: &[u8]) {
let chunk = self.fill_head(chunk);
self.push_tail(chunk);
}
// Fill the head budget first, then keep a capped tail.
let remaining_head = Self::HEAD_BUDGET.saturating_sub(self.head.len());
let head_len = remaining_head.min(chunk.len());
if head_len > 0 {
self.head.extend_from_slice(&chunk[..head_len]);
}
self.push_to_tail(&chunk[head_len..]);
/// Fill the stable prefix and return the bytes that did not fit.
fn fill_head<'a>(&mut self, chunk: &'a [u8]) -> &'a [u8] {
let Self {
head,
tail: _,
omitted_bytes: _,
} = self;
let remaining_head = Self::HEAD_BUDGET.saturating_sub(head.len());
// A shorter chunk fits entirely in the head.
let (chunk_head, chunk_tail) = chunk
.split_at_checked(remaining_head)
.unwrap_or((chunk, &[]));
head.extend_from_slice(chunk_head);
chunk_tail
}
/// Append bytes known not to belong in the head, keeping the newest tail bytes.
fn push_tail(&mut self, chunk: &[u8]) {
let Self {
head: _,
tail,
omitted_bytes,
} = self;
let remaining_tail = Self::TAIL_BUDGET.saturating_sub(tail.len());
let excess_tail = chunk.len().saturating_sub(remaining_tail);
*omitted_bytes = omitted_bytes.saturating_add(excess_tail);
// Discard old tail bytes first, then skip any excess incoming bytes.
let chunk = match excess_tail.checked_sub(tail.len()) {
None => {
tail.drain(..excess_tail);
chunk
}
Some(skip) => {
tail.clear();
&chunk[skip..]
}
};
tail.extend(chunk);
}
/// Return the retained output as a single byte vector.
@@ -95,45 +125,41 @@ impl<const MAX_BYTES: usize> HeadTailBuffer<MAX_BYTES> {
/// Append a later buffer with the same budget. This preserves the summary
/// of the original concatenated output, including its omission count.
pub(crate) fn push_buffer(&mut self, mut buffer: Self) {
self.push_chunk(std::mem::take(&mut buffer.head));
self.push_chunk(buffer.tail.drain(..).collect());
self.omitted_bytes = self.omitted_bytes.saturating_add(buffer.omitted_bytes);
}
pub(crate) fn push_buffer(&mut self, buffer: Self) {
let Self {
head,
tail,
omitted_bytes,
} = buffer;
fn push_to_tail(&mut self, chunk: &[u8]) {
if chunk.is_empty() {
return;
}
if Self::TAIL_BUDGET == 0 {
self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len());
return;
}
self.omitted_bytes = self.omitted_bytes.saturating_add(omitted_bytes);
if chunk.len() >= Self::TAIL_BUDGET {
// This single chunk is larger than the whole tail budget. Keep only the last
// tail_budget bytes and drop everything else.
let start = chunk.len().saturating_sub(Self::TAIL_BUDGET);
let kept = &chunk[start..];
let dropped = chunk.len().saturating_sub(kept.len());
// Preserve an existing prefix; otherwise reuse the source head.
let overflow = if self.head.is_empty() {
self.head = head;
&[]
} else {
self.fill_head(&head)
};
// A full source tail displaces both the old tail and the unused source head.
if tail.len() == Self::TAIL_BUDGET {
self.omitted_bytes = self
.omitted_bytes
.saturating_add(self.tail.len())
.saturating_add(dropped);
self.tail.clear();
self.tail.extend(kept);
return;
}
self.tail.extend(chunk);
self.trim_tail_to_budget();
}
fn trim_tail_to_budget(&mut self) {
let excess = self.tail.len().saturating_sub(Self::TAIL_BUDGET);
if excess > 0 {
drop(self.tail.drain(..excess));
self.omitted_bytes = self.omitted_bytes.saturating_add(excess);
.saturating_add(overflow.len());
self.tail = tail;
} else {
self.push_tail(overflow);
// An empty destination can take a partial source tail without copying it.
if self.tail.is_empty() {
self.tail = tail;
} else {
// A nonempty source tail means its head, and now ours, is full.
let (first, second) = tail.as_slices();
self.push_tail(first);
self.push_tail(second);
}
}
}
}
codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs
+18 -18
View File
@@ -6,11 +6,11 @@ use pretty_assertions::assert_eq;
fn keeps_prefix_and_suffix_when_over_budget() {
let mut buf = HeadTailBuffer::<10>::default();
buf.push_chunk(b"0123456789".to_vec());
buf.push_chunk(b"0123456789");
assert_eq!(buf.omitted_bytes(), 0);
// Exceeds max by 2; we should keep head+tail and omit the middle.
buf.push_chunk(b"ab".to_vec());
buf.push_chunk(b"ab");
assert!(buf.omitted_bytes() > 0);
let rendered = String::from_utf8_lossy(&buf.to_bytes()).to_string();
@@ -25,30 +25,30 @@ fn keeps_prefix_and_suffix_when_over_budget() {
#[test]
fn max_bytes_zero_drops_everything() {
let mut buf = HeadTailBuffer::<0>::default();
buf.push_chunk(b"abc".to_vec());
buf.push_chunk(b"abc");
assert_eq!(buf.retained_bytes(), 0);
assert_eq!(buf.omitted_bytes(), 3);
assert_eq!(buf.to_bytes(), b"".to_vec());
assert_eq!(buf.to_bytes(), b"");
}
#[test]
fn head_budget_zero_keeps_only_last_byte_in_tail() {
let mut buf = HeadTailBuffer::<1>::default();
buf.push_chunk(b"abc".to_vec());
buf.push_chunk(b"abc");
assert_eq!(buf.retained_bytes(), 1);
assert_eq!(buf.omitted_bytes(), 2);
assert_eq!(buf.to_bytes(), b"c".to_vec());
assert_eq!(buf.to_bytes(), b"c");
}
#[test]
fn chunk_larger_than_tail_budget_keeps_only_tail_end() {
let mut buf = HeadTailBuffer::<10>::default();
buf.push_chunk(b"0123456789".to_vec());
buf.push_chunk(b"0123456789");
// Tail budget is 5 bytes. This chunk should replace the tail and keep only its last 5 bytes.
buf.push_chunk(b"ABCDEFGHIJK".to_vec());
buf.push_chunk(b"ABCDEFGHIJK");
let out = String::from_utf8_lossy(&buf.to_bytes()).to_string();
assert!(out.starts_with("01234"));
@@ -61,19 +61,19 @@ fn fills_head_then_tail_across_multiple_chunks() {
let mut buf = HeadTailBuffer::<10>::default();
// Fill the 5-byte head budget across multiple chunks.
buf.push_chunk(b"01".to_vec());
buf.push_chunk(b"234".to_vec());
assert_eq!(buf.to_bytes(), b"01234".to_vec());
buf.push_chunk(b"01");
buf.push_chunk(b"234");
assert_eq!(buf.to_bytes(), b"01234");
// Then fill the 5-byte tail budget.
buf.push_chunk(b"567".to_vec());
buf.push_chunk(b"89".to_vec());
assert_eq!(buf.to_bytes(), b"0123456789".to_vec());
buf.push_chunk(b"567");
buf.push_chunk(b"89");
assert_eq!(buf.to_bytes(), b"0123456789");
assert_eq!(buf.omitted_bytes(), 0);
// One more byte causes the tail to drop its oldest byte.
buf.push_chunk(b"a".to_vec());
assert_eq!(buf.to_bytes(), b"012346789a".to_vec());
buf.push_chunk(b"a");
assert_eq!(buf.to_bytes(), b"012346789a");
assert_eq!(buf.omitted_bytes(), 1);
}
@@ -82,8 +82,8 @@ fn empty_and_tiny_chunks_have_bounded_metadata() {
let mut buf = HeadTailBuffer::<10>::default();
for byte in b"0123456789ab" {
buf.push_chunk(Vec::new());
buf.push_chunk(vec![*byte]);
buf.push_chunk(&[]);
buf.push_chunk(&[*byte]);
}
assert_eq!(buf.retained_bytes(), 10);
codex-rs/core/src/unified_exec/mod_tests.rs
-29
View File
@@ -1,4 +1,3 @@
use super::head_tail_buffer::HeadTailBuffer;
use super::*;
use crate::codex_thread::BackgroundTerminalInfo;
use crate::environment_selection::TurnEnvironmentState;
@@ -322,34 +321,6 @@ async fn write_stdin(
.await
}
#[test]
fn push_chunk_preserves_prefix_and_suffix() {
let mut buffer: HeadTailBuffer = Default::default();
buffer.push_chunk(vec![b'a'; UNIFIED_EXEC_OUTPUT_MAX_BYTES]);
buffer.push_chunk(vec![b'b']);
buffer.push_chunk(vec![b'c']);
assert_eq!(buffer.retained_bytes(), UNIFIED_EXEC_OUTPUT_MAX_BYTES);
let snapshot = buffer.to_bytes();
let head_bytes = UNIFIED_EXEC_OUTPUT_MAX_BYTES / 2;
let tail_bytes = UNIFIED_EXEC_OUTPUT_MAX_BYTES - head_bytes;
let expected = std::iter::repeat_n(b'a', head_bytes + tail_bytes - 2)
.chain(b"bc".iter().copied())
.collect::<Vec<_>>();
assert_eq!(snapshot, expected);
}
#[test]
fn head_tail_buffer_default_preserves_prefix_and_suffix() {
let mut buffer: HeadTailBuffer = Default::default();
buffer.push_chunk(vec![b'a'; UNIFIED_EXEC_OUTPUT_MAX_BYTES]);
buffer.push_chunk(b"bc".to_vec());
let rendered = buffer.to_bytes();
assert_eq!(rendered.first(), Some(&b'a'));
assert!(rendered.ends_with(b"bc"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unified_exec_persists_across_requests() -> anyhow::Result<()> {
skip_if_sandbox!(Ok(()));
codex-rs/core/src/unified_exec/process.rs
+6 -6
View File
@@ -29,6 +29,7 @@ use codex_utils_pty::ExecCommandSession;
use codex_utils_pty::ProcessSignal as PtyProcessSignal;
use codex_utils_pty::SpawnedPty;
use super::UNIFIED_EXEC_OUTPUT_MAX_BYTES;
use super::UNIFIED_EXEC_OUTPUT_MAX_TOKENS;
use super::UnifiedExecError;
use super::head_tail_buffer::HeadTailBuffer;
@@ -56,11 +57,10 @@ pub(crate) struct NoopSpawnLifecycle;
impl SpawnLifecycle for NoopSpawnLifecycle {}
pub(crate) type OutputBuffer = Arc<Mutex<HeadTailBuffer>>;
/// Shared output state exposed to polling and streaming consumers.
#[derive(Clone)]
pub(crate) struct OutputHandles {
pub(crate) output_buffer: OutputBuffer,
pub(crate) struct OutputHandles<const MAX_BYTES: usize = UNIFIED_EXEC_OUTPUT_MAX_BYTES> {
pub(crate) output_buffer: Arc<Mutex<HeadTailBuffer<MAX_BYTES>>>,
pub(crate) output_notify: Arc<Notify>,
pub(crate) output_closed: Arc<AtomicBool>,
pub(crate) output_closed_notify: Arc<Notify>,
@@ -501,7 +501,7 @@ impl UnifiedExecProcess {
for chunk in chunks.into_iter().filter(|chunk| chunk.seq > last_seq) {
let bytes = chunk.chunk.into_inner();
let mut guard = output_buffer.lock().await;
guard.push_chunk(bytes.clone());
guard.push_chunk(&bytes);
drop(guard);
let _ = output_tx.send(bytes);
output_notify.notify_waiters();
@@ -544,7 +544,7 @@ impl UnifiedExecProcess {
last_seq = chunk.seq;
let bytes = chunk.chunk.into_inner();
let mut guard = output_buffer.lock().await;
guard.push_chunk(bytes.clone());
guard.push_chunk(&bytes);
drop(guard);
let _ = output_tx.send(bytes);
output_notify.notify_waiters();
@@ -605,7 +605,7 @@ impl UnifiedExecProcess {
match receiver.recv().await {
Ok(chunk) => {
let mut guard = output_buffer.lock().await;
guard.push_chunk(chunk.clone());
guard.push_chunk(&chunk);
drop(guard);
let _ = output_tx.send(chunk);
output_notify.notify_waiters();
codex-rs/core/src/unified_exec/process_manager.rs
+4 -4
View File
@@ -1315,11 +1315,11 @@ impl UnifiedExecProcessManager {
})
}
pub(super) async fn collect_output_until_deadline(
output: &OutputHandles,
pub(super) async fn collect_output_until_deadline<const MAX_BYTES: usize>(
output: &OutputHandles<MAX_BYTES>,
mut pause_state: Option<watch::Receiver<bool>>,
mut deadline: Instant,
) -> HeadTailBuffer {
) -> HeadTailBuffer<MAX_BYTES> {
const POST_EXIT_CLOSE_WAIT_CAP: Duration = Duration::from_millis(50);
let OutputHandles {
@@ -1339,7 +1339,7 @@ impl UnifiedExecProcessManager {
&mut post_exit_deadline,
)
.await;
let drained_output: HeadTailBuffer;
let drained_output: HeadTailBuffer<MAX_BYTES>;
let has_drained_output: bool;
let mut wait_for_output = None;
{
codex-rs/core/src/unified_exec/process_manager_tests.rs
+14 -27
View File
@@ -283,7 +283,8 @@ fn initial_exec_yield_time_has_no_platform_floor() {
#[tokio::test]
async fn output_collection_stays_bounded_across_repeated_drains() {
let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
let chunks: [&[u8]; 4] = [b"01234567", b"89ABCDEF", b"ghijklmnopq", b"rs"];
let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::<10>::default()));
let output_notify = Arc::new(Notify::new());
let output_closed = Arc::new(AtomicBool::new(false));
let output_closed_notify = Arc::new(Notify::new());
@@ -302,10 +303,8 @@ async fn output_collection_stays_bounded_across_repeated_drains() {
Instant::now() + Duration::from_secs(5),
);
let produce = async {
for byte in [b'a', b'b', b'c'] {
output_buffer.lock().await.push_chunk(
vec![byte; crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES],
);
for chunk in chunks {
output_buffer.lock().await.push_chunk(chunk);
output_notify.notify_one();
tokio::time::timeout(Duration::from_secs(1), async {
loop {
@@ -326,30 +325,21 @@ async fn output_collection_stays_bounded_across_repeated_drains() {
};
let (collected, ()) = tokio::join!(collect, produce);
let mut expected = HeadTailBuffer::default();
for byte in [b'a', b'b', b'c'] {
expected.push_chunk(vec![
byte;
crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES
]);
let mut expected = HeadTailBuffer::<10>::default();
for chunk in chunks {
expected.push_chunk(chunk);
}
assert_eq!(collected, expected);
}
#[tokio::test]
async fn output_collection_preserves_omissions_from_drained_buffer() {
let mut buffered_output = HeadTailBuffer::default();
buffered_output.push_chunk(vec![
b'a';
crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES
]);
buffered_output.push_chunk(b"overflow".to_vec());
let mut expected = HeadTailBuffer::default();
expected.push_chunk(vec![
b'a';
crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES
]);
expected.push_chunk(b"overflow".to_vec());
let mut buffered_output = HeadTailBuffer::<10>::default();
buffered_output.push_chunk(&[b'a'; 10]);
buffered_output.push_chunk(b"overflow");
let mut expected = HeadTailBuffer::<10>::default();
expected.push_chunk(&[b'a'; 10]);
expected.push_chunk(b"overflow");
let output_buffer = Arc::new(tokio::sync::Mutex::new(buffered_output));
let output_notify = Arc::new(Notify::new());
let output_closed = Arc::new(AtomicBool::new(true));
@@ -435,10 +425,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() {
};
let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
transcript
.lock()
.await
.push_chunk(b"PARTIAL_TRANSCRIPT".to_vec());
transcript.lock().await.push_chunk(b"PARTIAL_TRANSCRIPT");
emit_failed_initial_exec_end_if_unstored(
/*process_started_alive*/ false,