spill_resource_probe/
spill_resource_probe.rs

1// Copyright 2026 Nowledge
2// SPDX-License-Identifier: Apache-2.0
3
4//! Manual resource probe. Run identical source on both revisions in release mode.
5//! Heap counters include all worker threads; obtain process peak RSS externally.
6
7use hawdb_core::{Catalog, Result, RuntimeTaskContext, Value};
8use hawdb_executor::binding::Binding;
9use hawdb_executor::blocking::stream_sort_batches;
10use hawdb_executor::observer::QueryExecutionObserver;
11use hawdb_executor::pipeline::{
12    BatchControl, BatchExecutionContext, BindingBatch, BindingBatchSource,
13};
14use hawdb_executor::{ExecutionLimit, ExecutionMemoryConfig, QueryMemoryLedger};
15use hawdb_plan_cypher::{PhysicalPlan, SortDirection, SortItem, SortKey};
16use hawdb_qos::{
17    IoConcurrencyBudget, ProcessMemoryCapabilities, ProcessMemoryPolicy, ProcessMemoryPolicyConfig,
18    ProcessMemorySnapshot, RuntimeGovernor, RuntimeGovernorConfig, RuntimeMemorySnapshot,
19    RuntimeResourceBudget, RuntimeResourceSnapshot, RuntimeWorkRequest,
20};
21use std::alloc::{GlobalAlloc, Layout, System};
22use std::collections::BTreeMap;
23use std::num::{NonZeroU64, NonZeroUsize};
24use std::sync::atomic::{AtomicUsize, Ordering};
25use std::sync::Barrier;
26use std::time::Instant;
27
28struct CountingAllocator;
29static LIVE: AtomicUsize = AtomicUsize::new(0);
30static PEAK: AtomicUsize = AtomicUsize::new(0);
31static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0);
32static ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
33
34fn account(bytes: usize) {
35    let live = LIVE.fetch_add(bytes, Ordering::Relaxed) + bytes;
36    PEAK.fetch_max(live, Ordering::Relaxed);
37    ALLOCATIONS.fetch_add(1, Ordering::Relaxed);
38    ALLOCATED_BYTES.fetch_add(bytes, Ordering::Relaxed);
39}
40
41// SAFETY: Every operation delegates its unchanged layout and pointer to System.
42// Counters neither allocate nor alter allocation ownership or alignment.
43unsafe impl GlobalAlloc for CountingAllocator {
44    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
45        let pointer = unsafe { System.alloc(layout) };
46        if !pointer.is_null() {
47            account(layout.size());
48        }
49        pointer
50    }
51    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
52        let pointer = unsafe { System.alloc_zeroed(layout) };
53        if !pointer.is_null() {
54            account(layout.size());
55        }
56        pointer
57    }
58    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
59        LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
60        unsafe { System.dealloc(pointer, layout) };
61    }
62    unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, bytes: usize) -> *mut u8 {
63        let next = unsafe { System.realloc(pointer, layout, bytes) };
64        if !next.is_null() {
65            LIVE.fetch_sub(layout.size(), Ordering::Relaxed);
66            account(bytes);
67        }
68        next
69    }
70}
71
72#[global_allocator]
73static ALLOCATOR: CountingAllocator = CountingAllocator;
74
75struct Rows {
76    count: usize,
77    width: usize,
78}
79impl BindingBatchSource for Rows {
80    fn execute(
81        &mut self,
82        _: &PhysicalPlan,
83        _: ExecutionLimit,
84        emit: &mut dyn FnMut(BindingBatch) -> Result<BatchControl>,
85    ) -> Result<BatchControl> {
86        for index in 0..self.count {
87            let binding = Binding::values(BTreeMap::from([
88                ("key".into(), Value::Int(((index * 7919) % 97) as i64)),
89                ("ordinal".into(), Value::Int(index as i64)),
90                (
91                    "payload".into(),
92                    Value::List(vec![Value::Map(BTreeMap::from([(
93                        "body".into(),
94                        Value::String("x".repeat(self.width)),
95                    )]))]),
96                ),
97            ]));
98            if emit(vec![binding])? == BatchControl::Stop {
99                return Ok(BatchControl::Stop);
100            }
101        }
102        Ok(BatchControl::Continue)
103    }
104}
105
106#[derive(Default)]
107struct Measurement {
108    rows: usize,
109    spill_bytes: u64,
110    runs: usize,
111    ledger_peak: usize,
112    checksum: u64,
113}
114
115fn query(
116    index: usize,
117    rows: usize,
118    width: usize,
119    blocking: usize,
120    barrier: &Barrier,
121    governor: Option<&RuntimeGovernor>,
122) -> Result<Measurement> {
123    let directory =
124        std::env::temp_dir().join(format!("hawdb-spill-probe-{}-{index}", std::process::id()));
125    let memory = ExecutionMemoryConfig {
126        query_memory_bytes: NonZeroUsize::new(16 * 1024 * 1024).unwrap(),
127        blocking_operator_bytes: NonZeroUsize::new(blocking).unwrap(),
128        batch_payload_bytes: NonZeroUsize::new(64 * 1024).unwrap(),
129        max_spill_runs: NonZeroUsize::new(4096).unwrap(),
130        max_total_spill_runs: NonZeroUsize::new(16384).unwrap(),
131        max_spill_bytes: NonZeroU64::new(1024 * 1024 * 1024).unwrap(),
132        min_spill_free_bytes: NonZeroU64::MIN,
133        spill_directory: directory.clone(),
134        ..ExecutionMemoryConfig::default()
135    };
136    let permit = governor.map(|governor| {
137        governor
138            .try_admit(RuntimeWorkRequest::foreground_query(16 * 1024 * 1024, 0))
139            .unwrap()
140    });
141    let task = permit
142        .as_ref()
143        .map(|permit| permit.bind_task_context(RuntimeTaskContext::default()));
144    let ledger = QueryMemoryLedger::for_query(&memory, task.as_ref())?;
145    let observer = QueryExecutionObserver::default();
146    let catalog = Catalog::default();
147    let items = [SortItem {
148        key: SortKey::Column("key".into()),
149        direction: SortDirection::Asc,
150    }];
151    let mut source = Rows { count: rows, width };
152    let mut measurement = Measurement::default();
153    let mut previous = None;
154    barrier.wait();
155    barrier.wait();
156    stream_sort_batches(
157        &PhysicalPlan::EmptyExec,
158        &items,
159        &mut source,
160        BatchExecutionContext {
161            catalog: &catalog,
162            memory: &memory,
163            memory_ledger: &ledger,
164            task_context: task.as_ref(),
165            observer: &observer,
166            graph: None,
167        },
168        ExecutionLimit::unlimited(),
169        &mut |batch| {
170            for binding in batch {
171                let Value::Int(key) = binding.values["key"] else {
172                    panic!("missing key")
173                };
174                let Value::Int(ordinal) = binding.values["ordinal"] else {
175                    panic!("missing ordinal")
176                };
177                assert!(
178                    previous.is_none_or(|previous| previous < (key, ordinal)),
179                    "stable ordering changed"
180                );
181                previous = Some((key, ordinal));
182                let Value::List(payload) = &binding.values["payload"] else {
183                    panic!("missing payload")
184                };
185                let Value::Map(payload) = &payload[0] else {
186                    panic!("missing payload map")
187                };
188                let Value::String(body) = &payload["body"] else {
189                    panic!("missing body")
190                };
191                assert_eq!(body.len(), width);
192                measurement.checksum = measurement.checksum.wrapping_add(ordinal as u64);
193                measurement.rows += 1;
194            }
195            Ok(BatchControl::Continue)
196        },
197    )?;
198    assert_eq!(measurement.rows, rows);
199    assert_eq!(measurement.checksum, (rows as u64) * (rows as u64 - 1) / 2);
200    assert_eq!(ledger.snapshot().used_bytes, 0);
201    assert_eq!(memory.spill_pool_snapshot()?.active_bytes, 0);
202    measurement.ledger_peak = ledger.snapshot().peak_bytes;
203    for report in observer.into_reports().blocking_memory {
204        measurement.spill_bytes += report.spilled_bytes;
205        measurement.runs += report.spill_run_count;
206    }
207    if directory.exists() {
208        std::fs::remove_dir_all(directory).unwrap();
209    }
210    Ok(measurement)
211}
212
213fn shared_governor() -> RuntimeGovernor {
214    let policy = ProcessMemoryPolicy::new(
215        ProcessMemoryPolicyConfig::new(NonZeroU64::new(2 * 1024 * 1024 * 1024).unwrap())
216            .with_sample_max_age(std::time::Duration::from_secs(60)),
217    );
218    // A fixed host sample isolates governed execution from an RSS sampler.
219    policy.update(ProcessMemorySnapshot {
220        capabilities: ProcessMemoryCapabilities {
221            resident_memory: true,
222            ..Default::default()
223        },
224        resident_bytes: 128 * 1024 * 1024,
225        peak_resident_bytes: 128 * 1024 * 1024,
226        total_page_faults: None,
227        minor_page_faults: None,
228        major_page_faults: None,
229    });
230    RuntimeGovernor::new_with_process_memory_policy(
231        RuntimeGovernorConfig {
232            memory_budget_bytes: Some(512 * 1024 * 1024),
233            ..RuntimeGovernorConfig::shared_host()
234        },
235        RuntimeResourceSnapshot::from_parts(
236            RuntimeResourceBudget::from_limits(NonZeroUsize::new(16).unwrap(), None, None),
237            RuntimeMemorySnapshot::from_limits(Some(4 << 30), Some(4 << 30), None, None, None),
238        ),
239        IoConcurrencyBudget::new(16, 1),
240        policy,
241    )
242}
243
244fn main() {
245    let scenario = std::env::args().nth(1).unwrap_or_else(|| "wide".into());
246    let (workers, rows, width, blocking) = match scenario.as_str() {
247        "small" => (1, 32, 64, 128 * 1024),
248        "wide" => (1, 2048, 4096, 128 * 1024),
249        "many" => (1, 8192, 128, 8 * 1024),
250        "concurrent" => (4, 2048, 4096, 128 * 1024),
251        _ => panic!("use small, wide, many, or concurrent"),
252    };
253    let governed = std::env::args().nth(2).is_some_and(|option| {
254        assert_eq!(option, "--governed");
255        true
256    });
257    let governor = governed.then(shared_governor);
258    let barrier = Barrier::new(workers + 1);
259    std::thread::scope(|scope| {
260        let handles: Vec<_> = (0..workers)
261            .map(|worker| {
262                let barrier = &barrier;
263                let governor = governor.as_ref();
264                scope.spawn(move || query(worker, rows, width, blocking, barrier, governor))
265            })
266            .collect();
267        barrier.wait();
268        let baseline_live = LIVE.load(Ordering::Relaxed);
269        PEAK.store(baseline_live, Ordering::Relaxed);
270        let allocations = ALLOCATIONS.load(Ordering::Relaxed);
271        let allocated_bytes = ALLOCATED_BYTES.load(Ordering::Relaxed);
272        let start = Instant::now();
273        barrier.wait();
274        let results: Vec<_> = handles
275            .into_iter()
276            .map(|handle| handle.join().unwrap().unwrap())
277            .collect();
278        println!("scenario={scenario} governed={} workers={workers} rows={} width={width} blocking_bytes={blocking} elapsed_us={} heap_peak_bytes={} heap_growth_bytes={} allocations={} allocated_bytes={} spill_bytes={} spill_runs={} ledger_peak_bytes={} checksum={}",
279            usize::from(governed), rows * workers, start.elapsed().as_micros(), PEAK.load(Ordering::Relaxed),
280            PEAK.load(Ordering::Relaxed).saturating_sub(baseline_live), ALLOCATIONS.load(Ordering::Relaxed) - allocations,
281            ALLOCATED_BYTES.load(Ordering::Relaxed) - allocated_bytes,
282            results.iter().map(|result| result.spill_bytes).sum::<u64>(), results.iter().map(|result| result.runs).sum::<usize>(),
283            results.iter().map(|result| result.ledger_peak).max().unwrap(), results.iter().map(|result| result.checksum).sum::<u64>(),
284        );
285    });
286    if let Some(governor) = governor {
287        let snapshot = governor.snapshot();
288        assert_eq!(snapshot.admissions, workers as u64);
289        assert_eq!(snapshot.completions, workers as u64);
290        assert_eq!(snapshot.admitted_memory_bytes, 0);
291    }
292}
spill_resource_probe/spill_resource_probe.rs

spill_resource_probe/
spill_resource_probe.rs
