Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/common/src/utils/proxy.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
//! [`VecAllocExt`] and [`RawTableAllocExt`] to help tracking of memory allocations
19
20
use hashbrown::raw::{Bucket, RawTable};
21
22
/// Extension trait for [`Vec`] to account for allocations.
23
pub trait VecAllocExt {
24
    /// Item type.
25
    type T;
26
27
    /// [Push](Vec::push) new element to vector and increase
28
    /// `accounting` by any newly allocated bytes.
29
    ///
30
    /// Note that allocation counts  capacity, not size
31
    ///
32
    /// # Example:
33
    /// ```
34
    /// # use datafusion_common::utils::proxy::VecAllocExt;
35
    /// // use allocated to incrementally track how much memory is allocated in the vec
36
    /// let mut allocated = 0;
37
    /// let mut vec = Vec::new();
38
    /// // Push data into the vec and the accounting will be updated to reflect
39
    /// // memory allocation
40
    /// vec.push_accounted(1, &mut allocated);
41
    /// assert_eq!(allocated, 16); // space for 4 u32s
42
    /// vec.push_accounted(1, &mut allocated);
43
    /// assert_eq!(allocated, 16); // no new allocation needed
44
    ///
45
    /// // push more data into the vec
46
    /// for _ in 0..10 { vec.push_accounted(1, &mut allocated); }
47
    /// assert_eq!(allocated, 64); // underlying vec has space for 10 u32s
48
    /// assert_eq!(vec.allocated_size(), 64);
49
    /// ```
50
    /// # Example with other allocations:
51
    /// ```
52
    /// # use datafusion_common::utils::proxy::VecAllocExt;
53
    /// // You can use the same allocated size to track memory allocated by
54
    /// // another source. For example
55
    /// let mut allocated = 27;
56
    /// let mut vec = Vec::new();
57
    /// vec.push_accounted(1, &mut allocated); // allocates 16 bytes for vec
58
    /// assert_eq!(allocated, 43); // 16 bytes for vec, 27 bytes for other
59
    /// ```
60
    fn push_accounted(&mut self, x: Self::T, accounting: &mut usize);
61
62
    /// Return the amount of memory allocated by this Vec to store elements
63
    /// (`size_of<T> * capacity`).
64
    ///
65
    /// Note this calculation is not recursive, and does not include any heap
66
    /// allocations contained within the Vec's elements. Does not include the
67
    /// size of `self`
68
    ///
69
    /// # Example:
70
    /// ```
71
    /// # use datafusion_common::utils::proxy::VecAllocExt;
72
    /// let mut vec = Vec::new();
73
    /// // Push data into the vec and the accounting will be updated to reflect
74
    /// // memory allocation
75
    /// vec.push(1);
76
    /// assert_eq!(vec.allocated_size(), 16); // space for 4 u32s
77
    /// vec.push(1);
78
    /// assert_eq!(vec.allocated_size(), 16); // no new allocation needed
79
    ///
80
    /// // push more data into the vec
81
    /// for _ in 0..10 { vec.push(1); }
82
    /// assert_eq!(vec.allocated_size(), 64); // space for 64 now
83
    /// ```
84
    fn allocated_size(&self) -> usize;
85
}
86
87
impl<T> VecAllocExt for Vec<T> {
88
    type T = T;
89
90
0
    fn push_accounted(&mut self, x: Self::T, accounting: &mut usize) {
91
0
        let prev_capacty = self.capacity();
92
0
        self.push(x);
93
0
        let new_capacity = self.capacity();
94
0
        if new_capacity > prev_capacty {
95
0
            // capacity changed, so we allocated more
96
0
            let bump_size = (new_capacity - prev_capacty) * std::mem::size_of::<T>();
97
0
            // Note multiplication should never overflow because `push` would
98
0
            // have panic'd first, but the checked_add could potentially
99
0
            // overflow since accounting could be tracking additional values, and
100
0
            // could be greater than what is stored in the Vec
101
0
            *accounting = (*accounting).checked_add(bump_size).expect("overflow");
102
0
        }
103
0
    }
104
834
    fn allocated_size(&self) -> usize {
105
834
        std::mem::size_of::<T>() * self.capacity()
106
834
    }
107
}
108
109
/// Extension trait for hash browns [`RawTable`] to account for allocations.
110
pub trait RawTableAllocExt {
111
    /// Item type.
112
    type T;
113
114
    /// [Insert](RawTable::insert) new element into table and increase
115
    /// `accounting` by any newly allocated bytes.
116
    ///
117
    /// Returns the bucket where the element was inserted.
118
    /// Note that allocation counts capacity, not size.
119
    ///
120
    /// # Example:
121
    /// ```
122
    /// # use datafusion_common::utils::proxy::RawTableAllocExt;
123
    /// # use hashbrown::raw::RawTable;
124
    /// let mut table = RawTable::new();
125
    /// let mut allocated = 0;
126
    /// let hash_fn = |x: &u32| (*x as u64) % 1000;
127
    /// // pretend 0x3117 is the hash value for 1
128
    /// table.insert_accounted(1, hash_fn, &mut allocated);
129
    /// assert_eq!(allocated, 64);
130
    ///
131
    /// // insert more values
132
    /// for i in 0..100 { table.insert_accounted(i, hash_fn, &mut allocated); }
133
    /// assert_eq!(allocated, 400);
134
    /// ```
135
    fn insert_accounted(
136
        &mut self,
137
        x: Self::T,
138
        hasher: impl Fn(&Self::T) -> u64,
139
        accounting: &mut usize,
140
    ) -> Bucket<Self::T>;
141
}
142
143
impl<T> RawTableAllocExt for RawTable<T> {
144
    type T = T;
145
146
165
    fn insert_accounted(
147
165
        &mut self,
148
165
        x: Self::T,
149
165
        hasher: impl Fn(&Self::T) -> u64,
150
165
        accounting: &mut usize,
151
165
    ) -> Bucket<Self::T> {
152
165
        let hash = hasher(&x);
153
165
154
165
        match self.try_insert_no_grow(hash, x) {
155
151
            Ok(bucket) => bucket,
156
14
            Err(x) => {
157
14
                // need to request more memory
158
14
159
14
                let bump_elements = self.capacity().max(16);
160
14
                let bump_size = bump_elements * std::mem::size_of::<T>();
161
14
                *accounting = (*accounting).checked_add(bump_size).expect("overflow");
162
14
163
14
                self.reserve(bump_elements, hasher);
164
14
165
14
                // still need to insert the element since first try failed
166
14
                // Note: cannot use `.expect` here because `T` may not implement `Debug`
167
14
                match self.try_insert_no_grow(hash, x) {
168
14
                    Ok(bucket) => bucket,
169
0
                    Err(_) => panic!("just grew the container"),
170
                }
171
            }
172
        }
173
165
    }
174
}