console_subscriber/aggregator/
shrink.rs

1use std::{
2    any::type_name,
3    collections::hash_map::{HashMap, RandomState},
4    hash::{BuildHasher, Hash},
5    mem,
6    ops::{Deref, DerefMut},
7};
8
9#[derive(Debug, Clone)]
10pub(crate) struct ShrinkMap<K, V, S = RandomState> {
11    map: HashMap<K, V, S>,
12    shrink: Shrink,
13}
14
15#[derive(Debug, Clone)]
16pub(crate) struct ShrinkVec<T> {
17    vec: Vec<T>,
18    shrink: Shrink,
19}
20
21#[derive(Debug, Clone)]
22pub(crate) struct Shrink {
23    shrink_every: usize,
24    since_shrink: usize,
25    min_bytes: usize,
26}
27
28// === impl ShrinkMap ===
29
30impl<K, V> ShrinkMap<K, V>
31where
32    K: Hash + Eq,
33{
34    pub(crate) fn new() -> Self {
35        Self {
36            map: HashMap::new(),
37            shrink: Shrink::default(),
38        }
39    }
40}
41
42impl<K, V, S> ShrinkMap<K, V, S>
43where
44    K: Hash + Eq,
45    S: BuildHasher,
46{
47    pub(crate) fn try_shrink(&mut self) {
48        self.shrink.try_shrink_map(&mut self.map)
49    }
50
51    pub(crate) fn retain_and_shrink(&mut self, f: impl FnMut(&K, &mut V) -> bool) {
52        let len0 = self.len();
53
54        self.retain(f);
55
56        if self.len() < len0 {
57            tracing::debug!(
58                len = self.len(),
59                dropped = len0.saturating_sub(self.len()),
60                data.key = %type_name::<K>(),
61                data.val = %type_name::<V>(),
62                "dropped unused entries"
63            );
64            self.try_shrink();
65        }
66    }
67}
68
69impl<K, V, S> Deref for ShrinkMap<K, V, S> {
70    type Target = HashMap<K, V, S>;
71    fn deref(&self) -> &Self::Target {
72        &self.map
73    }
74}
75
76impl<K, V, S> DerefMut for ShrinkMap<K, V, S> {
77    fn deref_mut(&mut self) -> &mut Self::Target {
78        &mut self.map
79    }
80}
81
82impl<K, V> Default for ShrinkMap<K, V>
83where
84    K: Hash + Eq,
85{
86    fn default() -> Self {
87        Self::new()
88    }
89}
90
91// === impl ShrinkVec ===
92
93impl<T> ShrinkVec<T> {
94    pub(crate) fn new() -> Self {
95        Self {
96            vec: Vec::new(),
97            shrink: Shrink::default(),
98        }
99    }
100
101    pub(crate) fn try_shrink(&mut self) {
102        self.shrink.try_shrink_vec(&mut self.vec)
103    }
104
105    pub(crate) fn retain_and_shrink(&mut self, f: impl FnMut(&T) -> bool) {
106        let len0 = self.len();
107
108        self.retain(f);
109
110        if self.len() < len0 {
111            tracing::debug!(
112                len = self.len(),
113                dropped = len0.saturating_sub(self.len()),
114                data = %type_name::<T>(),
115                "dropped unused data"
116            );
117            self.try_shrink();
118        }
119    }
120}
121
122impl<T> Deref for ShrinkVec<T> {
123    type Target = Vec<T>;
124    fn deref(&self) -> &Self::Target {
125        &self.vec
126    }
127}
128
129impl<T> DerefMut for ShrinkVec<T> {
130    fn deref_mut(&mut self) -> &mut Self::Target {
131        &mut self.vec
132    }
133}
134
135impl<T> Default for ShrinkVec<T> {
136    fn default() -> Self {
137        Self::new()
138    }
139}
140
141// === impl Shrink ===
142
143impl Shrink {
144    /// Shrinking every 60 flushes should be roughly every minute.
145    pub(crate) const DEFAULT_SHRINK_INTERVAL: usize = 60;
146
147    /// Don't bother if we'd free less than 4KB of memory.
148    // TODO(eliza): this number was chosen totally arbitrarily; it's the minimum
149    // page size on x86.
150    pub(crate) const DEFAULT_MIN_SIZE_BYTES: usize = 1024 * 4;
151
152    pub(crate) fn try_shrink_map<K, V, S>(&mut self, map: &mut HashMap<K, V, S>)
153    where
154        K: Hash + Eq,
155        S: BuildHasher,
156    {
157        if self.should_shrink::<(K, V)>(map.capacity(), map.len()) {
158            map.shrink_to_fit();
159        }
160    }
161
162    pub(crate) fn try_shrink_vec<T>(&mut self, vec: &mut Vec<T>) {
163        if self.should_shrink::<T>(vec.capacity(), vec.len()) {
164            vec.shrink_to_fit();
165        }
166    }
167
168    /// Returns `true` if we should shrink with a capacity of `capacity` Ts and
169    /// an actual length of `len` Ts.
170    fn should_shrink<T>(&mut self, capacity: usize, len: usize) -> bool {
171        // Has the required interval elapsed since the last shrink?
172        self.since_shrink = self.since_shrink.saturating_add(1);
173        if self.since_shrink < self.shrink_every {
174            tracing::trace!(
175                self.since_shrink,
176                self.shrink_every,
177                capacity_bytes = capacity * mem::size_of::<T>(),
178                used_bytes = len * mem::size_of::<T>(),
179                data = %type_name::<T>(),
180                "should_shrink: shrink interval has not elapsed"
181            );
182            return false;
183        }
184
185        // Okay, would we free up at least `min_bytes` by shrinking?
186        let capacity_bytes = capacity * mem::size_of::<T>();
187        let used_bytes = len * mem::size_of::<T>();
188        let diff = capacity_bytes.saturating_sub(used_bytes);
189        if diff < self.min_bytes {
190            tracing::trace!(
191                self.since_shrink,
192                self.shrink_every,
193                self.min_bytes,
194                freed_bytes = diff,
195                capacity_bytes,
196                used_bytes,
197                data = %type_name::<T>(),
198                "should_shrink: would not free sufficient bytes"
199            );
200            return false;
201        }
202
203        // Reset the clock! time to shrink!
204        self.since_shrink = 0;
205        tracing::debug!(
206            self.since_shrink,
207            self.shrink_every,
208            self.min_bytes,
209            freed_bytes = diff,
210            capacity_bytes,
211            used_bytes,
212            data = %type_name::<T>(),
213            "should_shrink: shrinking!"
214        );
215        true
216    }
217}
218
219impl Default for Shrink {
220    fn default() -> Self {
221        Self {
222            since_shrink: 0,
223            shrink_every: Self::DEFAULT_SHRINK_INTERVAL,
224            min_bytes: Self::DEFAULT_MIN_SIZE_BYTES,
225        }
226    }
227}