Optimize unbounded byte scans with memchr (#26265)

## Summary

This PR adds `memchr` for some low-hanging performance improvements
(namely, in MCP stdio, Ollama streaming, and full message-history
newline counts).

Codex produced the following release benchmarks:

| Operation | Before | After | Speedup |
| --- | ---: | ---: | ---: |
| MCP 1 MiB chunked line | 2.172 s | 3.984 ms | 545x |
| Ollama 1 MiB chunked line | 1.673 s | 2.790 ms | 600x |
| Count newlines in 10 MiB history | 132.83 ms | 20.05 ms | 6.6x |

With a "real" MCP setup (`ExecutorStdioServerLauncher` started a Python
MCP server, completed `initialize`, requested `tools/list`, and
deserialized a 1 MiB tool description over newline-delimited stdio),
it's about 16x faster end-to-end:

| Branch | 50 calls | Per call |
| --- | ---: | ---: |
| `main` | 862.53 ms | 17.25 ms |
| this branch | 53.89 ms | 1.08 ms |

`memchr` is already in our dependency tree and extremely widely used for
this kind of optimized scanning.
This commit is contained in:
Charlie Marsh
2026-06-04 09:53:08 -04:00
committed by GitHub
parent d46a98d31a
commit 7da4af622f
13 changed files with 284 additions and 31 deletions
@@ -20,11 +20,11 @@
use std::future::Future;
use std::io;
use std::mem::take;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use bytes::BytesMut;
use codex_exec_server::ExecOutputStream;
use codex_exec_server::ExecProcess;
use codex_exec_server::ExecProcessEvent;
@@ -32,6 +32,7 @@ use codex_exec_server::ExecProcessEventReceiver;
use codex_exec_server::ProcessId;
use codex_exec_server::ProcessOutputChunk;
use codex_exec_server::WriteStatus;
use memchr::memchr;
use rmcp::service::RoleClient;
use rmcp::service::RxJsonRpcMessage;
use rmcp::service::TxJsonRpcMessage;
@@ -46,6 +47,42 @@ use tracing::warn;
static PROCESS_COUNTER: AtomicUsize = AtomicUsize::new(1);
#[derive(Default)]
#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
struct LineBuffer {
bytes: BytesMut,
/// Prefix already scanned and known not to contain a newline.
scanned_len: usize,
}
impl LineBuffer {
fn extend_from_slice(&mut self, bytes: &[u8]) {
self.bytes.extend_from_slice(bytes);
}
fn take_line(&mut self) -> Option<BytesMut> {
let Some(relative_index) = memchr(b'\n', &self.bytes[self.scanned_len..]) else {
self.scanned_len = self.bytes.len();
return None;
};
let newline_index = self.scanned_len + relative_index;
let mut line = self.bytes.split_to(newline_index + 1);
line.truncate(newline_index);
self.scanned_len = 0;
Some(line)
}
fn take_remaining(&mut self) -> Option<BytesMut> {
if self.bytes.is_empty() {
return None;
}
self.scanned_len = 0;
Some(self.bytes.split())
}
}
// Remote public implementation.
/// A client-side rmcp transport backed by an executor-managed process.
@@ -73,10 +110,10 @@ pub(super) struct ExecutorProcessTransport {
/// Buffered child stdout bytes that have not yet formed a complete
/// newline-delimited JSON-RPC message.
stdout: Vec<u8>,
stdout: LineBuffer,
/// Buffered stderr bytes for diagnostic logging.
stderr: Vec<u8>,
stderr: LineBuffer,
/// Whether the executor has reported process closure or a terminal
/// subscription failure. Once closed, any remaining partial stdout line is
@@ -105,8 +142,8 @@ impl ExecutorProcessTransport {
process,
events,
program_name,
stdout: Vec::new(),
stderr: Vec::new(),
stdout: LineBuffer::default(),
stderr: LineBuffer::default(),
closed: false,
terminated: false,
last_seq: 0,
@@ -288,15 +325,10 @@ impl ExecutorProcessTransport {
// so EOF after a complete JSON object behaves like local rmcp's
// `decode_eof` handling.
loop {
let line_end = self.stdout.iter().position(|byte| *byte == b'\n');
let line = match (line_end, allow_partial && !self.stdout.is_empty()) {
(Some(index), _) => {
let mut line = self.stdout.drain(..=index).collect::<Vec<_>>();
line.pop();
line
}
(None, true) => self.stdout.drain(..).collect(),
(None, false) => return None,
let line = match self.stdout.take_line() {
Some(line) => line,
None if allow_partial => self.stdout.take_remaining()?,
None => return None,
};
let line = Self::trim_trailing_carriage_return(line);
match from_slice::<RxJsonRpcMessage<RoleClient>>(&line) {
@@ -315,12 +347,8 @@ impl ExecutorProcessTransport {
// Keep stderr line-oriented in logs so a chatty MCP server does not
// produce one log record per byte chunk.
self.stderr.extend_from_slice(bytes);
while let Some(index) = self.stderr.iter().position(|byte| *byte == b'\n') {
let mut line = self.stderr.drain(..=index).collect::<Vec<_>>();
line.pop();
if line.last() == Some(&b'\r') {
line.pop();
}
while let Some(line) = self.stderr.take_line() {
let line = Self::trim_trailing_carriage_return(line);
info!(
"MCP server stderr ({}): {}",
self.program_name,
@@ -330,10 +358,9 @@ impl ExecutorProcessTransport {
}
fn flush_stderr(&mut self) {
if self.stderr.is_empty() {
let Some(line) = self.stderr.take_remaining() else {
return;
}
let line = take(&mut self.stderr);
};
info!(
"MCP server stderr ({}): {}",
self.program_name,
@@ -341,14 +368,18 @@ impl ExecutorProcessTransport {
);
}
fn trim_trailing_carriage_return(mut line: Vec<u8>) -> Vec<u8> {
fn trim_trailing_carriage_return(mut line: BytesMut) -> BytesMut {
if line.last() == Some(&b'\r') {
line.pop();
line.truncate(line.len() - 1);
}
line
}
}
#[cfg(test)]
#[path = "executor_process_transport_tests.rs"]
mod tests;
impl Drop for ExecutorProcessTransport {
fn drop(&mut self) {
if self.terminated {
@@ -0,0 +1,72 @@
use bytes::BytesMut;
use pretty_assertions::assert_eq;
use super::LineBuffer;
#[test]
fn searches_only_new_bytes_after_partial_line() {
let mut buffer = LineBuffer::default();
buffer.extend_from_slice(b"partial");
assert_eq!(buffer.take_line(), None);
assert_eq!(
buffer,
LineBuffer {
bytes: BytesMut::from(&b"partial"[..]),
scanned_len: 7,
}
);
buffer.extend_from_slice(b" line");
assert_eq!(buffer.take_line(), None);
assert_eq!(
buffer,
LineBuffer {
bytes: BytesMut::from(&b"partial line"[..]),
scanned_len: 12,
}
);
buffer.extend_from_slice(b"\nnext");
assert_eq!(
buffer.take_line(),
Some(BytesMut::from(&b"partial line"[..]))
);
assert_eq!(
buffer,
LineBuffer {
bytes: BytesMut::from(&b"next"[..]),
scanned_len: 0,
}
);
}
#[test]
fn splits_multiple_lines_and_retains_partial_tail() {
let mut buffer = LineBuffer::default();
buffer.extend_from_slice(b"first\nsecond\npartial");
assert_eq!(buffer.take_line(), Some(BytesMut::from(&b"first"[..])));
assert_eq!(buffer.take_line(), Some(BytesMut::from(&b"second"[..])));
assert_eq!(buffer.take_line(), None);
assert_eq!(
buffer,
LineBuffer {
bytes: BytesMut::from(&b"partial"[..]),
scanned_len: 7,
}
);
}
#[test]
fn takes_unterminated_remaining_bytes_at_eof() {
let mut buffer = LineBuffer::default();
buffer.extend_from_slice(b"remaining");
assert_eq!(buffer.take_line(), None);
assert_eq!(
buffer.take_remaining(),
Some(BytesMut::from(&b"remaining"[..]))
);
assert_eq!(buffer, LineBuffer::default());
}