fix: resume exact reads within large segments
This commit is contained in:
@@ -600,41 +600,56 @@ pub fn read_session_public_index_entry(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let page = read_session_public_index_page(
|
||||
session_root,
|
||||
SessionPublicIndexLimits {
|
||||
max_segments: 1,
|
||||
..limits
|
||||
},
|
||||
expected_generation,
|
||||
Some(SessionPublicIndexScanPosition {
|
||||
segment_index,
|
||||
byte_offset: 0,
|
||||
entry_offset: 0,
|
||||
lineage: None,
|
||||
}),
|
||||
2,
|
||||
|entry| entry.entry_ref == entry_ref,
|
||||
)?;
|
||||
let mut matches = page.entries.into_iter();
|
||||
let item = matches.next();
|
||||
if matches.next().is_some() {
|
||||
return Err(SessionPublicIndexReadError::Corrupt);
|
||||
let mut position = SessionPublicIndexScanPosition {
|
||||
segment_index,
|
||||
byte_offset: 0,
|
||||
entry_offset: 0,
|
||||
lineage: None,
|
||||
};
|
||||
let mut generation = expected_generation.map(str::to_string);
|
||||
let mut selected = None;
|
||||
let mut scanned_bytes = 0_u64;
|
||||
let mut scanned_segments = 0_usize;
|
||||
let mut scanned_entries = 0_usize;
|
||||
loop {
|
||||
let page = read_session_public_index_page(
|
||||
session_root,
|
||||
SessionPublicIndexLimits {
|
||||
max_segments: 1,
|
||||
..limits
|
||||
},
|
||||
generation.as_deref(),
|
||||
Some(position.clone()),
|
||||
2,
|
||||
|entry| entry.entry_ref == entry_ref,
|
||||
)?;
|
||||
generation.get_or_insert_with(|| page.generation.clone());
|
||||
scanned_bytes = scanned_bytes.saturating_add(page.scanned_bytes);
|
||||
scanned_segments = scanned_segments.saturating_add(page.scanned_segments);
|
||||
scanned_entries = scanned_entries.saturating_add(page.scanned_entries);
|
||||
for item in page.entries {
|
||||
if selected.replace(item).is_some() {
|
||||
return Err(SessionPublicIndexReadError::Corrupt);
|
||||
}
|
||||
}
|
||||
let Some(next_position) = page.next_position else {
|
||||
break;
|
||||
};
|
||||
if next_position.segment_index != segment_index {
|
||||
break;
|
||||
}
|
||||
if next_position == position {
|
||||
return Err(SessionPublicIndexReadError::ResourceLimit);
|
||||
}
|
||||
position = next_position;
|
||||
}
|
||||
if page
|
||||
.next_position
|
||||
.as_ref()
|
||||
.is_some_and(|position| position.segment_index == segment_index)
|
||||
{
|
||||
return Err(SessionPublicIndexReadError::ResourceLimit);
|
||||
}
|
||||
Ok(item.map(|item| SessionPublicIndexEntryRead {
|
||||
session_id: page.session_id,
|
||||
generation: page.generation,
|
||||
Ok(selected.map(|item| SessionPublicIndexEntryRead {
|
||||
session_id: session_id.to_string(),
|
||||
generation: generation.expect("a successful index page always has a generation"),
|
||||
item,
|
||||
scanned_bytes: page.scanned_bytes,
|
||||
scanned_segments: page.scanned_segments,
|
||||
scanned_entries: page.scanned_entries,
|
||||
scanned_bytes,
|
||||
scanned_segments,
|
||||
scanned_entries,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -1810,6 +1825,79 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_entry_read_resumes_within_selected_segment_across_byte_and_entry_budgets() {
|
||||
let (_temp, root, store) = create_store();
|
||||
let session_id = crate::new_session_id();
|
||||
let segment_id = segment_id(1_700);
|
||||
store
|
||||
.create_segment(
|
||||
session_id,
|
||||
segment_id,
|
||||
&[start(
|
||||
session_id,
|
||||
vec![message("seed", LoggedRole::User, "ordinary seed")],
|
||||
None,
|
||||
None,
|
||||
)],
|
||||
)
|
||||
.unwrap();
|
||||
for index in 0..8 {
|
||||
let id = if index == 7 {
|
||||
"target".to_string()
|
||||
} else {
|
||||
format!("ordinary-{index}")
|
||||
};
|
||||
store
|
||||
.append(
|
||||
session_id,
|
||||
segment_id,
|
||||
&LogEntry::AnnotatedAssistantItem {
|
||||
ts: 2 + index,
|
||||
entry: message(
|
||||
&id,
|
||||
LoggedRole::Assistant,
|
||||
if index == 7 { "needle" } else { "ordinary" },
|
||||
),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
let (_, first_record_bytes, _) = store
|
||||
.read_first_log_record_read_only_bounded(session_id, segment_id, u64::MAX)
|
||||
.unwrap();
|
||||
|
||||
let by_bytes = read_session_public_index_entry(
|
||||
&root,
|
||||
SessionPublicIndexLimits {
|
||||
max_bytes: first_record_bytes,
|
||||
..SessionPublicIndexLimits::default()
|
||||
},
|
||||
None,
|
||||
&segment_id.to_string(),
|
||||
"Etarget",
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(by_bytes.item.entry.full_text, "needle");
|
||||
assert!(by_bytes.scanned_bytes > first_record_bytes);
|
||||
|
||||
let by_entries = read_session_public_index_entry(
|
||||
&root,
|
||||
SessionPublicIndexLimits {
|
||||
max_entries: 2,
|
||||
..SessionPublicIndexLimits::default()
|
||||
},
|
||||
None,
|
||||
&segment_id.to_string(),
|
||||
"Etarget",
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(by_entries.item.entry.full_text, "needle");
|
||||
assert!(by_entries.scanned_entries > 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_scan_rejects_one_indivisible_record_over_the_byte_budget() {
|
||||
let (_temp, root, store) = create_store();
|
||||
|
||||
@@ -534,4 +534,89 @@ mod tests {
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn search_then_read_crosses_entry_budget_within_one_segment() {
|
||||
let temp = tempfile::TempDir::new().unwrap();
|
||||
let root = temp.path().join("session");
|
||||
let store = WorkerSessionStore::new(&root).unwrap();
|
||||
let session_id = session_store::new_session_id();
|
||||
let segment_id = uuid::Uuid::from_u128(20_000);
|
||||
store
|
||||
.create_segment(
|
||||
session_id,
|
||||
segment_id,
|
||||
&[LogEntry::AnnotatedSegmentStart {
|
||||
ts: 1,
|
||||
session_id,
|
||||
system_prompt: None,
|
||||
config: RequestConfig::default(),
|
||||
history: vec![message("seed", "ordinary")],
|
||||
forked_from: None,
|
||||
compacted_from: None,
|
||||
}],
|
||||
)
|
||||
.unwrap();
|
||||
for index in 0..4 {
|
||||
let id = if index == 3 {
|
||||
"target".to_string()
|
||||
} else {
|
||||
format!("ordinary-{index}")
|
||||
};
|
||||
store
|
||||
.append(
|
||||
session_id,
|
||||
segment_id,
|
||||
&LogEntry::AnnotatedAssistantItem {
|
||||
ts: 2 + index,
|
||||
entry: message(&id, if index == 3 { "needle" } else { "ordinary" }),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let mut generation = None;
|
||||
let mut scan_cursor = None;
|
||||
let found = loop {
|
||||
let mut search_request = request(session_id, generation.clone(), scan_cursor.clone());
|
||||
search_request.max_entries = 2;
|
||||
let page = match search(&root, &search_request, None) {
|
||||
SessionPublicSearchAvailability::Page { page } => page,
|
||||
other => panic!("unexpected search result: {other:?}"),
|
||||
};
|
||||
if let Some(item) = page.items.into_iter().next() {
|
||||
break (page.generation, item);
|
||||
}
|
||||
assert!(page.has_more);
|
||||
generation = Some(page.generation);
|
||||
scan_cursor = page.next_scan_cursor;
|
||||
};
|
||||
assert_eq!(found.1.entry_ref, "Etarget");
|
||||
|
||||
let read_page = match read(
|
||||
&root,
|
||||
&SessionPublicReadRequest {
|
||||
workspace_id: "workspace-1".into(),
|
||||
source: runtime_api::SessionPublicSource::Retained {
|
||||
worker_id: crate::identity::WorkerId::now_v7(),
|
||||
},
|
||||
expected_session_id: session_id.to_string(),
|
||||
expected_generation: Some(found.0),
|
||||
segment_id: segment_id.to_string(),
|
||||
entry_ref: found.1.entry_ref,
|
||||
mode: SessionPublicReadMode::Full,
|
||||
byte_offset: 0,
|
||||
max_content_bytes: 16 * 1024,
|
||||
max_scan_bytes: 64 * 1024 * 1024,
|
||||
max_segments: 64,
|
||||
max_entries: 2,
|
||||
},
|
||||
None,
|
||||
) {
|
||||
SessionPublicReadAvailability::Page { page } => page,
|
||||
other => panic!("unexpected direct read result: {other:?}"),
|
||||
};
|
||||
assert_eq!(read_page.content, "needle");
|
||||
assert!(read_page.scanned_entries > 2);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user