Skip to main content

mz_timely_util/containers/
heap_size.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Heap size accounting for the containers that back arrangement batches.
11
12use columnar::{AsBytes, Borrow, Columnar};
13use columnation::Columnation;
14use differential_dataflow::columnar::layout::Coltainer;
15
16use crate::columnation::ColumnationStack;
17
18/// A container that can report the heap allocations backing it.
19pub trait HeapSize {
20    /// Calls `callback(size, capacity)`, in bytes, once per allocation backing `self`.
21    fn heap_size(&self, callback: impl FnMut(usize, usize));
22}
23
24impl<T: Columnation> HeapSize for ColumnationStack<T> {
25    fn heap_size(&self, callback: impl FnMut(usize, usize)) {
26        ColumnationStack::heap_size(self, callback)
27    }
28}
29
30impl<C: Columnar> HeapSize for Coltainer<C> {
31    fn heap_size(&self, mut callback: impl FnMut(usize, usize)) {
32        // Columnar containers expose their contents as byte slices, one per column, and each
33        // non-empty column is one `Vec` allocation, so the callback count is right. They do
34        // not expose spare capacity, so each slice reports its length as both size and
35        // capacity, and the capacity is a lower bound.
36        for (_align, bytes) in self.container.borrow().as_bytes() {
37            callback(bytes.len(), bytes.len());
38        }
39    }
40}