diff --git a/crates/cos/src/document.rs b/crates/cos/src/document.rs --- a/crates/cos/src/document.rs +++ b/crates/cos/src/document.rs @@ -23,10 +23,11 @@ pub(crate) use reader::find; #[cfg(test)] use reader::rfind; pub(crate) use state::ObjectReader; +pub(crate) use state::owned_bytes; use state::{ObjStm, Slot}; pub use state::{ReadState, SourceIdentity}; #[cfg(test)] -use state::{SECTION_READS, objstm_member, owned_bytes}; +use state::{SECTION_READS, objstm_member}; use xref::XrefIndex; pub use xref::{Revision, XrefEntry}; diff --git a/crates/cos/src/document/objects.rs b/crates/cos/src/document/objects.rs --- a/crates/cos/src/document/objects.rs +++ b/crates/cos/src/document/objects.rs @@ -1,5 +1,5 @@ use super::reader::find; -use super::state::{LENGTHS, LoadGuard, ObjStm, ObjectReader, ReadState, Slot, SourceIdentity, objstm_member, owned_bytes}; +use super::state::{CACHE_INSERTS, LENGTHS, LoadGuard, ObjStm, ObjectReader, ReadState, Slot, SourceIdentity, objstm_member, owned_bytes}; use super::xref::{Revision, XrefEntry}; use super::*; @@ -181,6 +181,7 @@ impl Document { // The editing cache is bounded like the scan reader's: re-parsing a file-backed view // is cheap, and returned objects outlive eviction. self.cache.lock().map_err(|_| CosError::Poisoned)?.insert(num, obj.clone()); + CACHE_INSERTS.with(|c| c.set(c.get().wrapping_add(1))); self.evict_editor_cache_if_over_budget(num, &obj)?; Ok(obj) } @@ -307,7 +308,7 @@ impl Document { let mut document = self.clone(); document.cache = Arc::default(); document.objstms = Arc::default(); - ObjectReader { document, decoded_count: 0 } + ObjectReader { document, decoded_count: 0, counted_objects: 0, counted_bytes: 0 } } fn load(&self, num: u32, depth: u32) -> Result { diff --git a/crates/cos/src/document/state.rs b/crates/cos/src/document/state.rs --- a/crates/cos/src/document/state.rs +++ b/crates/cos/src/document/state.rs @@ -122,6 +122,10 @@ impl std::fmt::Debug for ReadState { pub(crate) struct ObjectReader { pub(crate) document: Document, pub(super) decoded_count: usize, + /// Objects in the parsed-object cache and the bytes of their own, as last measured. Kept + /// as a running total so a read does not measure the whole cache again. + pub(super) counted_objects: usize, + pub(super) counted_bytes: usize, } impl ObjectReader { @@ -133,7 +137,9 @@ impl ObjectReader { } pub(crate) fn try_get(&mut self, reference: ObjRef) -> Result, CosError> { + let inserts_before = CACHE_INSERTS.with(|c| c.get()); let object = self.document.try_get(reference.num); + let inserted = CACHE_INSERTS.with(|c| c.get()).wrapping_sub(inserts_before); // A missed eviction here only costs a re-parse, so never fail the read on a poisoned // lock. The editing-cache evictors propagate instead; they guard user edits. if let Ok(mut streams) = self.document.objstms.lock() { @@ -165,14 +171,41 @@ impl ObjectReader { // have cached large objects. if let Ok(mut cache) = self.document.cache.lock() { let data = &self.document.data; - if cache.len() >= Self::OBJECTS || cache.values().fold(0usize, |n, o| n.saturating_add(owned_bytes(o, data, 0))) > Self::STREAM_BYTES { + // A hit inserted nothing; a read that inserted exactly this object (no nested + // /Length or object stream load) added only it. Cache length and pointer equality + // alone cannot tell that from a read that evicted and replaced entries (a large + // nested object makes the editor cache keep only the newcomer), so count the + // inserts. Anything else, such as a failed load, a /Length that cached further + // objects or a cache changed from outside, measures the whole cache again. + let grew_by_this_object = inserted == 1 + && cache.len() == self.counted_objects.saturating_add(1) + && matches!((&object, cache.get(&reference.num)), (Ok(read), Some(cached)) if Arc::ptr_eq(read, cached)); + if object.is_ok() && inserted == 0 && cache.len() == self.counted_objects { + // Unchanged: a hit, or an object held by the overlay. + } else if let (true, Ok(read)) = (grew_by_this_object, &object) { + self.counted_bytes = self.counted_bytes.saturating_add(owned_bytes(read, data, 0)); + self.counted_objects = cache.len(); + } else { + self.counted_bytes = cache.values().fold(0usize, |n, o| n.saturating_add(owned_bytes(o, data, 0))); + self.counted_objects = cache.len(); + } + if cache.len() >= Self::OBJECTS || self.counted_bytes > Self::STREAM_BYTES { cache.clear(); + self.counted_objects = 0; + self.counted_bytes = 0; } } object } } +thread_local! { + /// Objects this thread has put into a document's parsed-object cache. A reader compares + /// it before and after a read: a read that inserted exactly one object cannot have + /// evicted or replaced anything else, whatever the cache length says. + pub(super) static CACHE_INSERTS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + #[cfg(test)] thread_local! { /// Cross-reference sections read on this thread. @@ -183,21 +216,24 @@ thread_local! { /// is measured in full. const OWNED_BYTES_DEPTH: usize = 256; -/// Bytes `object` holds of its own: string and stream data that is not a view into `data` -/// (the input file). An object nested deeper than [`OWNED_BYTES_DEPTH`] cannot be measured -/// and counts as too large to keep. -pub(super) fn owned_bytes(object: &Object, data: &Arc>, depth: usize) -> usize { +/// Bytes `object` holds of its own: string, name and dictionary-key data, and stream data that +/// is not a view into `data` (the input file). A dictionary that indexes its keys holds each +/// one twice. An object nested deeper than [`OWNED_BYTES_DEPTH`] cannot be measured and counts +/// as too large to keep. +pub(crate) fn owned_bytes(object: &Object, data: &Arc>, depth: usize) -> usize { if depth > OWNED_BYTES_DEPTH { return usize::MAX; } + let dict_bytes = |own: usize, d: &crate::object::Dict| { + let copies = if d.is_indexed() { 2 } else { 1 }; + d.iter().fold(own, |n, (k, v)| n.saturating_add(k.len().saturating_mul(copies)).saturating_add(owned_bytes(v, data, depth + 1))) + }; match object { Object::String(s) => s.bytes.len(), - Object::Stream(s) => { - let own = if s.raw.shares(data) { 0 } else { s.raw.len() }; - s.dict.iter().fold(own, |n, (_, v)| n.saturating_add(owned_bytes(v, data, depth + 1))) - } + Object::Name(n) => n.len(), + Object::Stream(s) => dict_bytes(if s.raw.shares(data) { 0 } else { s.raw.len() }, &s.dict), Object::Array(a) => a.iter().fold(0usize, |n, v| n.saturating_add(owned_bytes(v, data, depth + 1))), - Object::Dict(d) => d.iter().fold(0usize, |n, (_, v)| n.saturating_add(owned_bytes(v, data, depth + 1))), + Object::Dict(d) => dict_bytes(0, d), _ => 0, } } diff --git a/crates/cos/src/document/tests.rs b/crates/cos/src/document/tests.rs --- a/crates/cos/src/document/tests.rs +++ b/crates/cos/src/document/tests.rs @@ -693,11 +693,43 @@ fn temporary_reader_keeps_file_views_and_evicts_data_of_its_own() { // Data of its own (decrypted streams, strings) counts: past the budget the cache is // cleared after the next read, including a failed one. let owned = Object::Stream(crate::Stream::from_raw(Dict::new(), vec![b'y'; ObjectReader::STREAM_BYTES + 1])); - reader.document.cache.lock().unwrap().insert(3, Arc::new(owned)); + reader.document.cache.lock().unwrap().insert(5, Arc::new(owned)); assert_eq!(*reader.get(ObjRef::new(99, 0)), Object::Null); assert!(reader.document.cache.lock().unwrap().is_empty(), "owned data over budget is evicted"); } +#[test] +fn temporary_reader_running_total_matches_the_cache() { + let mut bodies = vec!["<< /Type /Catalog /Pages 2 0 R >>".to_owned(), "<< /Type /Pages /Kids [] /Count 0 >>".to_owned()]; + bodies.extend((0..300).map(|i| format!("<< /Value {i} /Name ({}) >>", "s".repeat(i % 40)))); + let refs: Vec<_> = bodies.iter().map(String::as_str).collect(); + let doc = Document::open(Arc::new(build(&refs, "/Root 1 0 R"))).unwrap(); + let mut reader = doc.object_reader(); + for num in (1..=302).chain([5, 301, 99_999, 7]) { + let _ = reader.try_get(ObjRef::new(num, 0)); + let cache = reader.document.cache.lock().unwrap(); + let measured = cache.values().fold(0usize, |n, o| n.saturating_add(owned_bytes(o, &reader.document.data, 0))); + assert_eq!((reader.counted_objects, reader.counted_bytes), (cache.len(), measured), "after object {num}"); + assert!(cache.len() < ObjectReader::OBJECTS); + } +} + +#[test] +fn temporary_reader_recounts_when_a_read_replaces_cache_entries() { + // One cached object, then a read of a string over the editor cache's byte budget: the + // editor eviction keeps only the newcomer, so the cache length stays at one. The + // string is data of its own and must not stay cached past the reader's budget. + let big = format!("({})", "a".repeat(Document::EDITOR_CACHE_BYTES + 1)); + let bytes = build(&["<< /Type /Catalog /Pages 2 0 R >>", "<< /Type /Pages /Kids [] /Count 0 >>", "<< /Small 1 >>", &big], "/Root 1 0 R"); + let doc = Document::open(Arc::new(bytes)).unwrap(); + let mut reader = doc.object_reader(); + let _ = reader.get(ObjRef::new(3, 0)); + assert_eq!((reader.document.cache.lock().unwrap().len(), reader.counted_objects), (1, 1)); + assert!(matches!(&*reader.get(ObjRef::new(4, 0)), Object::String(_))); + assert!(reader.document.cache.lock().unwrap().is_empty(), "a replaced entry must be counted: the string is over budget"); + assert_eq!((reader.counted_objects, reader.counted_bytes), (0, 0)); +} + #[test] fn temporary_reader_and_full_save_do_not_fill_snapshot_caches() { let mut source = Document::new_empty(); diff --git a/crates/cos/src/object.rs b/crates/cos/src/object.rs --- a/crates/cos/src/object.rs +++ b/crates/cos/src/object.rs @@ -172,6 +172,11 @@ impl Dict { self.entries.iter().map(|(k, v)| (k, v)) } + /// Whether the keys are held a second time, in the lookup index. + pub(crate) fn is_indexed(&self) -> bool { + self.index.is_some() + } + pub fn len(&self) -> usize { self.entries.len() } diff --git a/crates/cos/src/writer.rs b/crates/cos/src/writer.rs --- a/crates/cos/src/writer.rs +++ b/crates/cos/src/writer.rs @@ -15,7 +15,7 @@ use std::io::Write; use std::sync::Arc; use crate::CosError; -use crate::document::{Document, ObjectReader, XrefEntry}; +use crate::document::{Document, ObjectReader, XrefEntry, owned_bytes}; use crate::object::{Dict, ObjRef, Object, PdfString, Stream}; /// Serialize one object. Deterministic: same object → same bytes. @@ -280,6 +280,10 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result, CosErro let mut map: HashMap = HashMap::new(); // Record the type while traversing so classifying packed objects needs no second parse. let mut order: Vec<(ObjRef, bool)> = Vec::new(); + // The parsed object of each entry in `order`, kept for the write pass while the budget + // allows it. Past the budget the write pass reads the object again. + let mut kept: Vec>> = Vec::new(); + let mut kept_budget = KeptObjects::default(); let mut queue: VecDeque = VecDeque::new(); let visit_refs = |o: &Object, queue: &mut VecDeque| collect_refs(o, &mut |r| queue.push_back(r)); for key in [&b"Root"[..], b"Info", b"Encrypt"] { @@ -305,6 +309,7 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result, CosErro map.insert(r, output_number(order.len())?); order.push((r, matches!(*o, Object::Stream(_)))); visit_refs(&o, &mut queue); + kept.push(kept_budget.admit(&o, reader.document.bytes()).then_some(o)); } let version = reader.document.version().chars().take(3).collect::(); let version = if opts.object_streams && version.as_str() < "1.5" { "1.5".to_string() } else { version }; @@ -329,14 +334,14 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result, CosErro let mut rows: BTreeMap = BTreeMap::new(); // Keep only references, not renamed copies of the entire object graph. Standalone // objects still precede object streams, preserving the established output order. - let mut packed: Vec<(u32, ObjRef)> = Vec::new(); + let mut packed: Vec<(u32, ObjRef, Option>)> = Vec::new(); for (i, (r, is_stream)) in order.iter().enumerate() { let num = output_number(i)?; let in_stream = opts.object_streams && !is_stream && !encryption_objects.contains(r) && root_outside != Some(*r); if in_stream { - packed.push((num, *r)); + packed.push((num, *r, kept[i].take())); } else { - let o = renumber(&*reread(&mut reader, *r, *is_stream)?, &map); + let o = renumber(&*written_object(&mut reader, *r, *is_stream, kept[i].take())?, &map); rows.insert(num, Row::InFile(out.len() as u64, 0)); write_indirect(num, 0, &prepared(&reader.document, r.num, num, 0, &o), &mut out); } @@ -345,9 +350,9 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result, CosErro // individually: the whole stream is (§7.5.7, §7.6.2). let mut next = output_number(order.len())?; let mut batch = ObjectStreamBatch::default(); - for (num, reference) in packed { + for (num, reference, parsed) in packed { let mut encoded = Vec::new(); - serialize(&renumber(&*reread(&mut reader, reference, false)?, &map), &mut encoded); + serialize(&renumber(&*written_object(&mut reader, reference, false, parsed)?, &map), &mut encoded); if batch.count > 0 && batch.body.len().saturating_add(encoded.len()).saturating_add(1) > OBJSTM_BYTES { batch.flush(&reader.document, &mut next, &mut rows, &mut out)?; } @@ -400,6 +405,56 @@ pub fn write_full(doc: &Document, opts: &SaveOptions) -> Result, CosErro Ok(out) } +/// What `write_full` keeps of the objects it parsed while numbering them: owned bytes (see +/// [`owned_bytes`]; streams read from the file are views and cost only their dictionary) and +/// tree nodes. Beyond either limit objects are read again when written, which bounds the memory +/// of a save of a huge document. +#[derive(Default)] +struct KeptObjects { + bytes: usize, + nodes: usize, +} + +impl KeptObjects { + const BYTES: usize = 64 * 1024 * 1024; + const NODES: usize = 2_000_000; + + /// Whether `object` still fits the budget; if so it is counted against it. + fn admit(&mut self, object: &Object, data: &Arc>) -> bool { + let bytes = self.bytes.saturating_add(owned_bytes(object, data, 0)); + let nodes = self.nodes.saturating_add(node_count(object, 0)); + if bytes > Self::BYTES || nodes > Self::NODES { + return false; + } + (self.bytes, self.nodes) = (bytes, nodes); + true + } +} + +/// Values in `object`'s tree, counting the object itself. Past the parser's nesting limit +/// the count is `usize::MAX`, so nothing that deep is kept. +fn node_count(object: &Object, depth: usize) -> usize { + if depth > 256 { + return usize::MAX; + } + let children = match object { + Object::Stream(s) => s.dict.iter().fold(0usize, |n, (_, v)| n.saturating_add(node_count(v, depth + 1))), + Object::Array(a) => a.iter().fold(0usize, |n, v| n.saturating_add(node_count(v, depth + 1))), + Object::Dict(d) => d.iter().fold(0usize, |n, (_, v)| n.saturating_add(node_count(v, depth + 1))), + _ => 0, + }; + children.saturating_add(1) +} + +/// The object to write for `r`: the one parsed while numbering when it was kept, else `r` read +/// again. +fn written_object(reader: &mut ObjectReader, r: ObjRef, is_stream: bool, kept: Option>) -> Result, CosError> { + match kept { + Some(o) => Ok(o), + None => reread(reader, r, is_stream), + } +} + /// Object `r` read again for writing. The traversal read it once and only kept its number, so /// a second read that fails, finds nothing or finds a different kind of object must fail the /// save rather than write `null` in its place. @@ -674,6 +729,50 @@ mod tests { Lexer::new(&out, 0).object().expect("parses") } + #[test] + fn kept_objects_stop_at_their_byte_and_node_budget() { + let data = Arc::new(Vec::new()); + let string = Object::String(PdfString::literal(vec![b'x'; 1000])); + let array = Object::Array(vec![Object::Int(1), Object::Null]); + let mut kept = KeptObjects::default(); + assert!(kept.admit(&string, &data) && kept.admit(&array, &data)); + assert_eq!((kept.bytes, kept.nodes), (1000, 4)); + let mut full = KeptObjects { bytes: KeptObjects::BYTES - 999, nodes: 0 }; + assert!(!full.admit(&string, &data), "one byte over the byte budget"); + assert_eq!((full.bytes, full.nodes), (KeptObjects::BYTES - 999, 0), "a refused object is not counted"); + assert!(full.admit(&Object::Int(1), &data)); + let mut full = KeptObjects { bytes: 0, nodes: KeptObjects::NODES - 2 }; + assert!(!full.admit(&array, &data), "a node past the node budget"); + assert!(full.admit(&Object::Int(1), &data)); + let mut deep = Object::Null; + for _ in 0..300 { + deep = Object::Array(vec![deep]); + } + assert!(!KeptObjects::default().admit(&deep, &data), "nesting past the limit is never kept"); + } + + #[test] + fn kept_objects_count_names_and_dictionary_keys() { + let data = Arc::new(Vec::new()); + let name = Object::Name(vec![b'n'; KeptObjects::BYTES + 1]); + let mut kept = KeptObjects::default(); + assert!(!kept.admit(&name, &data), "a name over the byte budget is not kept"); + assert_eq!((kept.bytes, kept.nodes), (0, 0)); + let mut dict = Dict::new(); + dict.set(vec![b'k'; KeptObjects::BYTES + 1], Object::Null); + assert!(!kept.admit(&Object::Dict(dict), &data), "nor a dictionary key"); + let mut small = Dict::new(); + small.set(b"Key".to_vec(), Object::name("Value")); + assert!(kept.admit(&Object::Dict(small), &data)); + assert_eq!(kept.bytes, 3 + 5); + // A dictionary large enough to index its keys holds each one twice. + let mut big = Dict::new(); + for i in 0..64 { + big.set(format!("K{i:03}").into_bytes(), Object::Null); + } + assert_eq!(owned_bytes(&Object::Dict(big), &data, 0), 64 * 4 * 2); + } + #[test] fn a_second_read_that_differs_fails_the_save_instead_of_writing_null() { let mut doc = Document::new_empty();