Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions wacore/binary/src/zlib_pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,14 @@ impl<'a> InflateReader<'a> {
self.total_out
}

/// Compressed input consumed so far plus the input's full length. The
/// ratio lets callers extrapolate totals (e.g. record counts) from a
/// prefix without a second pass over the blob.
#[inline]
pub fn compressed_progress(&self) -> (usize, usize) {
(self.in_pos, self.input.len())
}

/// Whether zlib reported a proper stream end (terminator + adler32
/// checksum). An EOF (`ensure` returning false) without this means the
/// input was truncated, not finished.
Expand Down
50 changes: 50 additions & 0 deletions wacore/src/history_sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,29 @@ impl<'a> FieldWalker<'a> {
self.reader.total_out()
}

/// Decompressed bytes already handed to the parser. Excludes the buffered
/// not-yet-parsed window, so a density observed over this prefix is not
/// diluted by data the inflater ran ahead on.
fn parsed_bytes(&self) -> u64 {
self.reader
.total_out()
.saturating_sub(self.reader.available().len() as u64)
Comment thread
jlucaso1 marked this conversation as resolved.
Outdated
}

/// Total decompressed size extrapolated from the zlib ratio observed so
/// far; exact once the stream has fully inflated.
fn estimated_total_out(&self) -> u64 {
let (in_pos, in_len) = self.reader.compressed_progress();
if in_pos == 0 {
return self.reader.total_out();
}
self.reader
.total_out()
.saturating_mul(in_len as u64)
.checked_div(in_pos as u64)
.unwrap_or(0)
}

Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
/// Re-borrow the payload of the field most recently yielded by
/// [`FieldWalker::next_field`] (it stays buffered until the next call).
fn pending_payload(&self, payload_start: usize) -> &[u8] {
Expand Down Expand Up @@ -250,6 +273,20 @@ fn process_history_sync_streaming(
decompressed_size: 0,
};

// One-shot capacity extrapolation for the secret-record accumulator. A
// full pre-count pass (see the field comment above) was measured at ~2.5%
// of the decode, so instead the record density observed over a prefix is
Comment thread
jlucaso1 marked this conversation as resolved.
// scaled to the blob's extrapolated decompressed size. Costs O(1), adapts
// to blobs with few secrets, and an off estimate just falls back to
// doubling. Extrapolating from the first conversation alone overshot to
// the clamp on measured blobs (doubling the transient peak), so the
// sample must first reach RESERVE_SAMPLE_RECORDS; the doubling ladder up
// to that point copies only ~2x its own bytes, which is noise. The clamp
// bounds over-allocation to ~2 MB.
const RESERVE_SAMPLE_RECORDS: usize = 128;
const RECORD_RESERVE_CAP: usize = 16384;
let mut density_reserved = false;

while let Some(field) = walker.next_field()? {
if field.wire_type != wire_type::LENGTH_DELIMITED {
continue;
Expand All @@ -264,6 +301,19 @@ fn process_history_sync_streaming(
{
result.tc_token_candidates.push(candidate);
}
if !density_reserved && result.msg_secret_records.len() >= RESERVE_SAMPLE_RECORDS {
density_reserved = true;
let parsed = walker.parsed_bytes().max(1);
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
let records = result.msg_secret_records.len();
let estimated = (records as u64)
.saturating_mul(walker.estimated_total_out())
.checked_div(parsed)
.map_or(records, |est| est as usize)
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
Outdated
.min(RECORD_RESERVE_CAP);
result
.msg_secret_records
.reserve(estimated.saturating_sub(records));
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}
// pushnames (repeated) — only our own is needed
tags::history_sync::PUSHNAMES => {
Expand Down
Loading