mirror of
https://github.com/openai/codex.git
synced 2026-10-23 18:08:45 +00:00
Charlie Marsh
a9bc7bebfa
## What changed Move full rollout history reading and parsing into one blocking worker for both plain `.jsonl` and compressed `.jsonl.zst` files. Use the shared cancellation-aware line iterator so canceled loads stop before processing another record, while preserving rollout items, thread IDs, and parse-error counts. ## Testing Add regression tests for cancellation between records and equivalent history-loading results across plain and compressed files, including Unicode, CRLF line endings, malformed records, empty files, and invalid UTF-8. GitOrigin-RevId: 65e1bd7ebcdbbd0ab5f560e73b2c7f074f627f8f
100 lines
3.0 KiB
Rust
100 lines
3.0 KiB
Rust
//! Exercises cancellation during reads and between records.
|
|
|
|
use std::io;
|
|
use std::io::BufRead;
|
|
use std::io::Read;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::AtomicUsize;
|
|
use std::sync::atomic::Ordering;
|
|
|
|
use pretty_assertions::assert_eq;
|
|
use tempfile::TempDir;
|
|
|
|
use super::ReadMetrics;
|
|
use super::read_rollout_lines;
|
|
use super::scan_lines;
|
|
|
|
struct PausedRead {
|
|
contents: io::Cursor<&'static [u8]>,
|
|
started: Option<tokio::sync::oneshot::Sender<()>>,
|
|
resume: std::sync::mpsc::Receiver<()>,
|
|
finished: Option<tokio::sync::oneshot::Sender<()>>,
|
|
}
|
|
|
|
impl Read for PausedRead {
|
|
fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
|
|
if let Some(started) = self.started.take() {
|
|
let _ = started.send(());
|
|
self.resume.recv().map_err(io::Error::other)?;
|
|
}
|
|
self.contents.read(buffer)
|
|
}
|
|
}
|
|
|
|
impl Drop for PausedRead {
|
|
fn drop(&mut self) {
|
|
if let Some(finished) = self.finished.take() {
|
|
let _ = finished.send(());
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn cancellation_during_read_does_not_deliver_the_record() -> anyhow::Result<()> {
|
|
let (started, waiting) = tokio::sync::oneshot::channel();
|
|
let (resume, paused) = std::sync::mpsc::channel();
|
|
let (finished, done) = tokio::sync::oneshot::channel();
|
|
let reader = Box::new(PausedRead {
|
|
contents: io::Cursor::new(b"first\n".as_slice()),
|
|
started: Some(started),
|
|
resume: paused,
|
|
finished: Some(finished),
|
|
}) as Box<dyn Read + Send>;
|
|
let calls = Arc::new(AtomicUsize::new(0));
|
|
let seen = Arc::clone(&calls);
|
|
let task = tokio::spawn(scan_lines(
|
|
io::BufReader::new(reader).lines(),
|
|
ReadMetrics::default(),
|
|
move |lines| {
|
|
for line in lines {
|
|
line?;
|
|
seen.fetch_add(1, Ordering::SeqCst);
|
|
}
|
|
Ok(())
|
|
},
|
|
));
|
|
waiting.await?;
|
|
task.abort();
|
|
assert!(task.await.unwrap_err().is_cancelled());
|
|
resume.send(())?;
|
|
done.await?;
|
|
assert_eq!(calls.load(Ordering::SeqCst), 0);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn cancelling_full_scan_stops_before_the_next_record() -> anyhow::Result<()> {
|
|
let home = TempDir::new()?;
|
|
let path = home.path().join("rollout.jsonl");
|
|
std::fs::write(&path, "first\nsecond\n")?;
|
|
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
|
|
let (continue_tx, continue_rx) = std::sync::mpsc::channel();
|
|
let (finished_tx, finished_rx) = tokio::sync::oneshot::channel();
|
|
let task = tokio::spawn(async move {
|
|
read_rollout_lines(&path, move |lines| {
|
|
assert_eq!(lines.next().transpose()?, Some("first".to_string()));
|
|
let _ = started_tx.send(());
|
|
continue_rx.recv().map_err(io::Error::other)?;
|
|
let _ = finished_tx.send(lines.next().transpose());
|
|
Ok(())
|
|
})
|
|
.await
|
|
});
|
|
started_rx.await?;
|
|
task.abort();
|
|
assert!(task.await.unwrap_err().is_cancelled());
|
|
continue_tx.send(())?;
|
|
assert_eq!(finished_rx.await??, None);
|
|
Ok(())
|
|
}
|