Codex Handbook
core/src/unified_exec/head_tail_buffer.rs 183 lines
use crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES;use std::collections::VecDeque;/// A capped buffer that preserves a stable prefix ("head") and suffix ("tail"),/// dropping the middle once it exceeds the configured maximum. The buffer is/// symmetric meaning 50% of the capacity is allocated to the head and 50% is/// allocated to the tail.#[derive(Debug)]pub(crate) struct HeadTailBuffer {    max_bytes: usize,    head_budget: usize,    tail_budget: usize,    head: VecDeque<Vec<u8>>,    tail: VecDeque<Vec<u8>>,    head_bytes: usize,    tail_bytes: usize,    omitted_bytes: usize,}impl Default for HeadTailBuffer {    fn default() -> Self {        Self::new(UNIFIED_EXEC_OUTPUT_MAX_BYTES)    }}impl HeadTailBuffer {    /// Create a new buffer that retains at most `max_bytes` of output.    ///    /// The retained output is split across a prefix ("head") and suffix ("tail")    /// budget, dropping bytes from the middle once the limit is exceeded.    pub(crate) fn new(max_bytes: usize) -> Self {        let head_budget = max_bytes / 2;        let tail_budget = max_bytes.saturating_sub(head_budget);        Self {            max_bytes,            head_budget,            tail_budget,            head: VecDeque::new(),            tail: VecDeque::new(),            head_bytes: 0,            tail_bytes: 0,            omitted_bytes: 0,        }    }    // Used for tests.    #[allow(dead_code)]    /// Total bytes currently retained by the buffer (head + tail).    pub(crate) fn retained_bytes(&self) -> usize {        self.head_bytes.saturating_add(self.tail_bytes)    }    // Used for tests.    #[allow(dead_code)]    /// Total bytes that were dropped from the middle due to the size cap.    pub(crate) fn omitted_bytes(&self) -> usize {        self.omitted_bytes    }    /// Append a chunk of bytes to the buffer.    ///    /// 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 self.max_bytes == 0 {            self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len());            return;        }        // Fill the head budget first, then keep a capped tail.        if self.head_bytes < self.head_budget {            let remaining_head = self.head_budget.saturating_sub(self.head_bytes);            if chunk.len() <= remaining_head {                self.head_bytes = self.head_bytes.saturating_add(chunk.len());                self.head.push_back(chunk);                return;            }            // Split the chunk: part goes to head, remainder goes to tail.            let (head_part, tail_part) = chunk.split_at(remaining_head);            if !head_part.is_empty() {                self.head_bytes = self.head_bytes.saturating_add(head_part.len());                self.head.push_back(head_part.to_vec());            }            self.push_to_tail(tail_part.to_vec());            return;        }        self.push_to_tail(chunk);    }    /// Snapshot the retained output as a list of chunks.    ///    /// The returned chunks are ordered as: head chunks first, then tail chunks.    /// Omitted bytes are not represented in the snapshot.    pub(crate) fn snapshot_chunks(&self) -> Vec<Vec<u8>> {        let mut out = Vec::new();        out.extend(self.head.iter().cloned());        out.extend(self.tail.iter().cloned());        out    }    /// Return the retained output as a single byte vector.    ///    /// The output is formed by concatenating head chunks, then tail chunks.    /// Omitted bytes are not represented in the returned value.    pub(crate) fn to_bytes(&self) -> Vec<u8> {        let mut out = Vec::with_capacity(self.retained_bytes());        for chunk in self.head.iter() {            out.extend_from_slice(chunk);        }        for chunk in self.tail.iter() {            out.extend_from_slice(chunk);        }        out    }    /// Drain all retained chunks from the buffer and reset its state.    ///    /// The drained chunks are returned in head-then-tail order. Omitted bytes    /// are discarded along with the retained content.    pub(crate) fn drain_chunks(&mut self) -> Vec<Vec<u8>> {        let mut out: Vec<Vec<u8>> = self.head.drain(..).collect();        out.extend(self.tail.drain(..));        self.head_bytes = 0;        self.tail_bytes = 0;        self.omitted_bytes = 0;        out    }    fn push_to_tail(&mut self, chunk: Vec<u8>) {        if self.tail_budget == 0 {            self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len());            return;        }        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..].to_vec();            let dropped = chunk.len().saturating_sub(kept.len());            self.omitted_bytes = self                .omitted_bytes                .saturating_add(self.tail_bytes)                .saturating_add(dropped);            self.tail.clear();            self.tail_bytes = kept.len();            self.tail.push_back(kept);            return;        }        self.tail_bytes = self.tail_bytes.saturating_add(chunk.len());        self.tail.push_back(chunk);        self.trim_tail_to_budget();    }    fn trim_tail_to_budget(&mut self) {        let mut excess = self.tail_bytes.saturating_sub(self.tail_budget);        while excess > 0 {            match self.tail.front_mut() {                Some(front) if excess >= front.len() => {                    excess -= front.len();                    self.tail_bytes = self.tail_bytes.saturating_sub(front.len());                    self.omitted_bytes = self.omitted_bytes.saturating_add(front.len());                    self.tail.pop_front();                }                Some(front) => {                    front.drain(..excess);                    self.tail_bytes = self.tail_bytes.saturating_sub(excess);                    self.omitted_bytes = self.omitted_bytes.saturating_add(excess);                    break;                }                None => break,            }        }    }}#[cfg(test)]#[path = "head_tail_buffer_tests.rs"]mod tests;