mirror of
https://github.com/openai/codex.git
synced 2026-11-13 22:47:20 +00:00
Limit archive rollout reads to requested threads (#43494)
## Why Archiving threads reads metadata from every active rollout, including unrelated compressed files. ## What changed Add `RolloutReferenceIndex::scan_unarchived_threads` and use it when archiving. Filter canonical filenames by the requested thread IDs before reading contents, while retaining metadata-based ownership checks. ## Testing Add tests covering reverted rollouts, requested descendants, compressed rollouts, exclusion of archived and unrelated files, and candidate ownership checks. GitOrigin-RevId: a64f7fa25e5493ed20312f1d8216422752272bae
This commit is contained in:
1 parent
d665e3bbc8
commit
d0a8dcd157
3 files changed
+141
-13
No files matched your search
@@ -1,6 +1,7 @@
|
|||||||
//! Indexes direct fork references found in local rollout files.
|
//! Indexes direct fork references found in local rollout files.
|
||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
use std::collections::HashSet;
|
||||||
use std::collections::hash_map::Entry;
|
use std::collections::hash_map::Entry;
|
||||||
use std::io;
|
use std::io;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
@@ -13,6 +14,7 @@ use codex_protocol::protocol::HistoryPosition;
|
|||||||
use crate::ARCHIVED_SESSIONS_SUBDIR;
|
use crate::ARCHIVED_SESSIONS_SUBDIR;
|
||||||
use crate::SESSIONS_SUBDIR;
|
use crate::SESSIONS_SUBDIR;
|
||||||
use crate::compression::RolloutFile;
|
use crate::compression::RolloutFile;
|
||||||
|
use crate::rollout_file_name::RolloutFileName;
|
||||||
|
|
||||||
/// Direct history-base edges discovered from local rollout metadata.
|
/// Direct history-base edges discovered from local rollout metadata.
|
||||||
///
|
///
|
||||||
@@ -34,10 +36,13 @@ struct IndexedRollout {
|
|||||||
impl RolloutReferenceIndex {
|
impl RolloutReferenceIndex {
|
||||||
/// Scans active and archived local rollout metadata.
|
/// Scans active and archived local rollout metadata.
|
||||||
pub async fn scan(codex_home: &Path) -> io::Result<Self> {
|
pub async fn scan(codex_home: &Path) -> io::Result<Self> {
|
||||||
Self::scan_paths(vec![
|
Self::scan_paths(
|
||||||
codex_home.join(ARCHIVED_SESSIONS_SUBDIR),
|
vec![
|
||||||
codex_home.join(SESSIONS_SUBDIR),
|
codex_home.join(ARCHIVED_SESSIONS_SUBDIR),
|
||||||
])
|
codex_home.join(SESSIONS_SUBDIR),
|
||||||
|
],
|
||||||
|
/*thread_ids*/ None,
|
||||||
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -46,10 +51,30 @@ impl RolloutReferenceIndex {
|
|||||||
/// Reference counts exclude archived history and must not be used to decide whether a
|
/// Reference counts exclude archived history and must not be used to decide whether a
|
||||||
/// rollout can be deleted or compressed.
|
/// rollout can be deleted or compressed.
|
||||||
pub async fn scan_unarchived(codex_home: &Path) -> io::Result<Self> {
|
pub async fn scan_unarchived(codex_home: &Path) -> io::Result<Self> {
|
||||||
Self::scan_paths(vec![codex_home.join(SESSIONS_SUBDIR)]).await
|
Self::scan_paths(
|
||||||
|
vec![codex_home.join(SESSIONS_SUBDIR)],
|
||||||
|
/*thread_ids*/ None,
|
||||||
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn scan_paths(mut stack: Vec<PathBuf>) -> io::Result<Self> {
|
/// Scans unarchived files whose canonical filenames belong to the requested threads.
|
||||||
|
///
|
||||||
|
/// Skips unrelated rollout contents, including compressed files. Metadata still determines
|
||||||
|
/// ownership among the candidates. Reference counts are partial and must not be used to
|
||||||
|
/// decide whether a rollout can be deleted or compressed.
|
||||||
|
pub async fn scan_unarchived_threads(
|
||||||
|
codex_home: &Path,
|
||||||
|
thread_ids: &[ThreadId],
|
||||||
|
) -> io::Result<Self> {
|
||||||
|
let thread_ids = thread_ids.iter().copied().collect();
|
||||||
|
Self::scan_paths(vec![codex_home.join(SESSIONS_SUBDIR)], Some(&thread_ids)).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn scan_paths(
|
||||||
|
mut stack: Vec<PathBuf>,
|
||||||
|
thread_ids: Option<&HashSet<ThreadId>>,
|
||||||
|
) -> io::Result<Self> {
|
||||||
let mut rollouts_by_id = HashMap::new();
|
let mut rollouts_by_id = HashMap::new();
|
||||||
while let Some(directory) = stack.pop() {
|
while let Some(directory) = stack.pop() {
|
||||||
let mut entries = match tokio::fs::read_dir(directory.as_path()).await {
|
let mut entries = match tokio::fs::read_dir(directory.as_path()).await {
|
||||||
@@ -73,9 +98,13 @@ impl RolloutReferenceIndex {
|
|||||||
let Some(rollout_file) = RolloutFile::from_path(path) else {
|
let Some(rollout_file) = RolloutFile::from_path(path) else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
let Some(rollout_id) = crate::rollout_id_from_path(rollout_file.path()) else {
|
let Some(file_name) = RolloutFileName::parse(rollout_file.plain_file_name()) else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
|
if thread_ids.is_some_and(|ids| !ids.contains(&file_name.thread_id())) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let rollout_id = file_name.rollout_id();
|
||||||
let Ok(meta) = crate::read_session_meta_line(rollout_file.path()).await else {
|
let Ok(meta) = crate::read_session_meta_line(rollout_file.path()).await else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -151,6 +151,102 @@ async fn unarchived_scan_finds_all_active_rollouts_owned_by_a_thread() -> anyhow
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn filtered_unarchived_scan_includes_reverts_and_requested_descendants() -> anyhow::Result<()>
|
||||||
|
{
|
||||||
|
let home = TempDir::new()?;
|
||||||
|
let owner_uuid = Uuid::from_u128(40);
|
||||||
|
let owner_id = thread_id(owner_uuid)?;
|
||||||
|
let replacement_id = thread_id(Uuid::from_u128(41))?;
|
||||||
|
let descendant_uuid = Uuid::from_u128(42);
|
||||||
|
let descendant_id = thread_id(descendant_uuid)?;
|
||||||
|
let unrelated_uuid = Uuid::from_u128(43);
|
||||||
|
let unrelated_id = thread_id(unrelated_uuid)?;
|
||||||
|
let original_path = active_rollout_path(home.path(), owner_uuid);
|
||||||
|
let replacement_path = original_path.with_file_name(format!(
|
||||||
|
"rollout-2025-01-03T12-00-00-{owner_id}_{replacement_id}.jsonl"
|
||||||
|
));
|
||||||
|
let descendant_path = active_rollout_path(home.path(), descendant_uuid);
|
||||||
|
|
||||||
|
write_rollout(original_path.clone(), owner_id, /*history_base*/ None)?;
|
||||||
|
compress_now(&original_path)?;
|
||||||
|
write_rollout(
|
||||||
|
replacement_path.clone(),
|
||||||
|
owner_id,
|
||||||
|
/*history_base*/ None,
|
||||||
|
)?;
|
||||||
|
// A leftover compressed sibling must not supersede the plain rollout.
|
||||||
|
compress_now(&replacement_path)?;
|
||||||
|
write_rollout(
|
||||||
|
replacement_path.clone(),
|
||||||
|
owner_id,
|
||||||
|
/*history_base*/ None,
|
||||||
|
)?;
|
||||||
|
write_rollout(
|
||||||
|
descendant_path.clone(),
|
||||||
|
descendant_id,
|
||||||
|
Some(history_position(owner_id)),
|
||||||
|
)?;
|
||||||
|
write_rollout(
|
||||||
|
archived_rollout_path(home.path(), owner_uuid),
|
||||||
|
owner_id,
|
||||||
|
/*history_base*/ None,
|
||||||
|
)?;
|
||||||
|
let unrelated_path = active_rollout_path(home.path(), unrelated_uuid);
|
||||||
|
write_rollout(
|
||||||
|
unrelated_path.clone(),
|
||||||
|
unrelated_id,
|
||||||
|
/*history_base*/ None,
|
||||||
|
)?;
|
||||||
|
compress_now(&unrelated_path)?;
|
||||||
|
|
||||||
|
let index =
|
||||||
|
RolloutReferenceIndex::scan_unarchived_threads(home.path(), &[owner_id, descendant_id])
|
||||||
|
.await?;
|
||||||
|
let mut owned: Vec<_> = index
|
||||||
|
.rollouts_for_thread(owner_id)
|
||||||
|
.map(|(id, path)| (id, path.to_path_buf()))
|
||||||
|
.collect();
|
||||||
|
owned.sort_by_key(|(id, _)| id.to_string());
|
||||||
|
assert_eq!(
|
||||||
|
owned,
|
||||||
|
vec![
|
||||||
|
(owner_id, original_path.with_extension("jsonl.zst")),
|
||||||
|
(replacement_id, replacement_path),
|
||||||
|
]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
index.rollouts_for_thread(descendant_id).collect::<Vec<_>>(),
|
||||||
|
vec![(descendant_id, descendant_path.as_path())]
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
index.rollouts_for_thread(unrelated_id).collect::<Vec<_>>(),
|
||||||
|
Vec::new()
|
||||||
|
);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn filtered_unarchived_scan_validates_candidate_ownership() -> anyhow::Result<()> {
|
||||||
|
let home = TempDir::new()?;
|
||||||
|
let requested_uuid = Uuid::from_u128(50);
|
||||||
|
let requested_id = thread_id(requested_uuid)?;
|
||||||
|
let actual_owner = thread_id(Uuid::from_u128(51))?;
|
||||||
|
write_rollout(
|
||||||
|
active_rollout_path(home.path(), requested_uuid),
|
||||||
|
actual_owner,
|
||||||
|
/*history_base*/ None,
|
||||||
|
)?;
|
||||||
|
|
||||||
|
let index =
|
||||||
|
RolloutReferenceIndex::scan_unarchived_threads(home.path(), &[requested_id]).await?;
|
||||||
|
assert_eq!(
|
||||||
|
index.rollouts_for_thread(requested_id).collect::<Vec<_>>(),
|
||||||
|
Vec::new()
|
||||||
|
);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn self_history_base_does_not_count_as_reference() -> anyhow::Result<()> {
|
async fn self_history_base_does_not_count_as_reference() -> anyhow::Result<()> {
|
||||||
let home = TempDir::new()?;
|
let home = TempDir::new()?;
|
||||||
|
|||||||
@@ -39,12 +39,15 @@ pub(super) async fn archive_threads(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
let _writer_guards = store.acquire_writer_locks(&lock_thread_ids).await?;
|
let _writer_guards = store.acquire_writer_locks(&lock_thread_ids).await?;
|
||||||
// Already-archived rollouts need no move. Avoid reading the entire archive on every request.
|
// Only inspect active files whose names belong to the threads being archived.
|
||||||
let reference_index = RolloutReferenceIndex::scan_unarchived(store.config.codex_home.as_path())
|
let reference_index = RolloutReferenceIndex::scan_unarchived_threads(
|
||||||
.await
|
store.config.codex_home.as_path(),
|
||||||
.map_err(|err| ThreadStoreError::Internal {
|
&thread_ids,
|
||||||
message: format!("failed to scan thread rollout files: {err}"),
|
)
|
||||||
})?;
|
.await
|
||||||
|
.map_err(|err| ThreadStoreError::Internal {
|
||||||
|
message: format!("failed to scan thread rollout files: {err}"),
|
||||||
|
})?;
|
||||||
|
|
||||||
let parent_thread_id = thread_ids[0];
|
let parent_thread_id = thread_ids[0];
|
||||||
let mut archived_thread_ids = Vec::new();
|
let mut archived_thread_ids = Vec::new();
|
||||||
|
|||||||
Reference in new issue
Block a user