1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
//! Allocation SITES of one warm frame, as `alloc_census_warm` is its sizes.
//!
//! The census says a frame allocates in a shape that looks like regrowth; this
//! says which buffers. The profiler is started only after a first frame has
//! run, so everything a compressor builds once and keeps is already allocated
//! and out of the picture, and what it records is what a steady-state frame
//! costs.
//!
//! Run: `cargo run --release -p ffi-bench --example alloc_sites_warm
//! --features dhat-heap -- <level> [corpus path]`
//! then read `dhat-heap.json` (the viewer, or `jq` over its `pps` array).
use std::env;
use structured_zstd::encoding::{CompressionLevel, FrameCompressor};
#[cfg(feature = "dhat-heap")]
#[global_allocator]
static ALLOC: dhat::Alloc = dhat::Alloc;
fn main() {
let args: Vec<String> = env::args().collect();
let level: i32 = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(3);
let corpus = args
.get(2)
.map(String::as_str)
.unwrap_or("zstd/decodecorpus_files/z000033");
let data = std::fs::read(corpus)
.unwrap_or_else(|e| panic!("alloc_sites_warm: cannot read {corpus}: {e}"));
let mut out = Vec::new();
let mut compressor: FrameCompressor<&[u8], &mut Vec<u8>> =
FrameCompressor::new(CompressionLevel::Level(level));
compressor.set_source(&data[..]);
compressor.set_drain(&mut out);
compressor.compress();
// Everything above is warm-up. From here the compressor is reused, which is
// the shape under audit.
// `fresh` audits a compressor built inside the profiled region, which is
// what a one-shot caller pays and what the upstream one-shot entry point
// (`ZSTD_compress`, which creates and frees a context per call) is
// measured against. The default reuses the warmed one.
let fresh = args.get(3).map(|s| s == "fresh").unwrap_or(false);
// Sized before the profiler starts, as the C loop sizes its destination
// outside its loop: the output buffer's growth is the caller's cost, not
// the frame's, and counting it here would attribute it to the encoder.
let mut out2 = Vec::with_capacity(structured_zstd::encoding::compress_bound(data.len()));
#[cfg(feature = "dhat-heap")]
let profiler = dhat::Profiler::new_heap();
// Declared outside the branch so it outlives the profiler: dropping a
// compressor inside the profiled region reports its teardown as part of
// the frame.
let mut fresh_compressor: Option<FrameCompressor<&[u8], &mut Vec<u8>>> = None;
if fresh {
// The one-shot contract, not the streaming one: a context per frame
// compressing straight from the caller's slice, which is what
// `ZSTD_compress` does and what `compress_slice_to_vec` wraps.
let mut one_shot: FrameCompressor<&[u8], &mut Vec<u8>> =
FrameCompressor::new(CompressionLevel::Level(level));
one_shot.compress_independent_frame_into(&data[..], &mut out2);
fresh_compressor = Some(one_shot);
} else {
compressor.set_source(&data[..]);
compressor.set_drain(&mut out2);
compressor.compress();
}
// Dropped explicitly so the report covers the frame and not the teardown
// of the buffers above it.
#[cfg(feature = "dhat-heap")]
drop(profiler);
drop(fresh_compressor);
println!("level {level}: {} bytes in, {} out", data.len(), out2.len());
}