egglog_core_relations/pool/
mod.rs

1//! Utilities for pooling object allocations.
2
3use std::{
4    cell::{Cell, RefCell},
5    fmt,
6    hash::{Hash, Hasher},
7    mem::{self, ManuallyDrop},
8    ops::{Deref, DerefMut},
9    ptr,
10    rc::Rc,
11};
12
13use crate::{
14    AtomId,
15    free_join::execute::TrieNode,
16    numeric_id::{DenseIdMap, IdVec},
17};
18use fixedbitset::FixedBitSet;
19use hashbrown::HashTable;
20
21use crate::{
22    ColumnId, RowId,
23    action::{Instr, PredictedVals},
24    common::{HashMap, HashSet, IndexMap, IndexSet, ShardId, Value},
25    free_join::execute::SortedColumnIndex,
26    hash_index::{BufferedSubset, TableEntry},
27    offsets::SortedOffsetVector,
28    table::TableEntry as SwTableEntry,
29    table_spec::Constraint,
30};
31
32#[cfg(test)]
33mod tests;
34
35/// A trait for types whose allocations can be reused.
36pub trait Clear: Default {
37    /// Clear the object.
38    ///
39    /// The end result must be equivalent to `Self::default()`.
40    fn clear(&mut self);
41    /// Indicate whether or not this object should be reused.
42    fn reuse(&self) -> bool {
43        true
44    }
45    /// A rough approximation for the in-memory overhead of this object.
46    fn bytes(&self) -> usize;
47}
48
49impl<T> Clear for Vec<T> {
50    fn clear(&mut self) {
51        self.clear()
52    }
53    fn reuse(&self) -> bool {
54        self.capacity() > 256
55    }
56    fn bytes(&self) -> usize {
57        self.capacity() * mem::size_of::<T>()
58    }
59}
60
61impl<T: Clear> Clear for Rc<T> {
62    fn clear(&mut self) {
63        Rc::get_mut(self).unwrap().clear()
64    }
65    fn reuse(&self) -> bool {
66        Rc::strong_count(self) == 1 && Rc::weak_count(self) == 0
67    }
68    fn bytes(&self) -> usize {
69        mem::size_of::<T>()
70    }
71}
72
73impl<T: Clear> Clone for Pooled<Rc<T>>
74where
75    Rc<T>: InPoolSet<PoolSet>,
76{
77    fn clone(&self) -> Self {
78        Pooled {
79            data: self.data.clone(),
80        }
81    }
82}
83
84impl<T> Clear for HashSet<T> {
85    fn clear(&mut self) {
86        self.clear()
87    }
88    fn reuse(&self) -> bool {
89        self.capacity() > 0
90    }
91    fn bytes(&self) -> usize {
92        self.capacity() * mem::size_of::<T>()
93    }
94}
95
96impl<T> Clear for HashTable<T> {
97    fn clear(&mut self) {
98        self.clear()
99    }
100    fn reuse(&self) -> bool {
101        self.capacity() > 0
102    }
103    fn bytes(&self) -> usize {
104        self.capacity() * mem::size_of::<T>()
105    }
106}
107
108impl<K, V> Clear for HashMap<K, V> {
109    fn clear(&mut self) {
110        self.clear()
111    }
112    fn reuse(&self) -> bool {
113        self.capacity() > 0
114    }
115    fn bytes(&self) -> usize {
116        self.capacity() * mem::size_of::<(K, V)>()
117    }
118}
119
120impl<K, V> Clear for IndexMap<K, V> {
121    fn clear(&mut self) {
122        self.clear()
123    }
124    fn reuse(&self) -> bool {
125        self.capacity() > 0
126    }
127    fn bytes(&self) -> usize {
128        self.capacity() * (mem::size_of::<u64>() + mem::size_of::<(K, V)>())
129    }
130}
131
132impl<T> Clear for IndexSet<T> {
133    fn clear(&mut self) {
134        self.clear()
135    }
136    fn reuse(&self) -> bool {
137        self.capacity() > 0
138    }
139    fn bytes(&self) -> usize {
140        self.capacity() * (mem::size_of::<u64>() + mem::size_of::<T>())
141    }
142}
143
144impl Clear for FixedBitSet {
145    fn clear(&mut self) {
146        self.clone_from(&Default::default());
147    }
148    fn reuse(&self) -> bool {
149        !self.is_empty()
150    }
151    fn bytes(&self) -> usize {
152        self.len() / 8
153    }
154}
155
156impl<K, V> Clear for IdVec<K, V> {
157    fn clear(&mut self) {
158        self.clear()
159    }
160    fn reuse(&self) -> bool {
161        self.capacity() > 0
162    }
163    fn bytes(&self) -> usize {
164        self.capacity() * mem::size_of::<V>()
165    }
166}
167
168struct PoolState<T> {
169    data: Vec<T>,
170    bytes: usize,
171    limit: usize,
172}
173
174impl<T: Clear> PoolState<T> {
175    fn new(limit: usize) -> Self {
176        PoolState {
177            data: Vec::new(),
178            bytes: 0,
179            limit,
180        }
181    }
182
183    fn push(&mut self, mut item: T) {
184        if !item.reuse() {
185            return;
186        }
187        if self.bytes + item.bytes() > self.limit {
188            return;
189        }
190        item.clear();
191        self.bytes += item.bytes();
192        self.data.push(item);
193    }
194
195    fn pop(&mut self) -> T {
196        if let Some(got) = self.data.pop() {
197            self.bytes -= got.bytes();
198            got
199        } else {
200            Default::default()
201        }
202    }
203
204    fn clear_and_shrink(&mut self) {
205        self.data.clear();
206        self.bytes = 0;
207        self.data.shrink_to_fit();
208    }
209}
210
211/// A shared pool of objects.
212pub struct Pool<T> {
213    data: Rc<RefCell<PoolState<T>>>,
214}
215
216impl<T> Clone for Pool<T> {
217    fn clone(&self) -> Self {
218        Pool {
219            data: self.data.clone(),
220        }
221    }
222}
223
224impl<T: Clear> Default for Pool<T> {
225    fn default() -> Self {
226        Pool {
227            data: Rc::new(RefCell::new(PoolState::new(usize::MAX))),
228        }
229    }
230}
231
232impl<T: Clear + InPoolSet<PoolSet>> Pool<T> {
233    pub(crate) fn new(limit: usize) -> Pool<T> {
234        Pool {
235            data: Rc::new(RefCell::new(PoolState::new(limit))),
236        }
237    }
238    /// Get an empty value of type `T`, potentially reused from the pool.
239    pub(crate) fn get(&self) -> Pooled<T> {
240        let empty = self.data.borrow_mut().pop();
241
242        Pooled {
243            data: ManuallyDrop::new(empty),
244        }
245    }
246
247    /// Clear the contents of the pool and release any memory associated with it.
248    pub(crate) fn clear(&self) {
249        let mut data_mut = self.data.borrow_mut();
250        data_mut.clear_and_shrink();
251    }
252}
253
254/// An owned value of type `T` that can be returned to a memory pool when it is
255/// no longer used.
256pub struct Pooled<T: Clear + InPoolSet<PoolSet>> {
257    data: ManuallyDrop<T>,
258}
259
260impl<T: Clear + InPoolSet<PoolSet>> Default for Pooled<T> {
261    fn default() -> Self {
262        with_pool_set(|ps| ps.get::<T>())
263    }
264}
265
266impl<T: Clear + fmt::Debug + InPoolSet<PoolSet>> fmt::Debug for Pooled<T> {
267    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
268        let data: &T = &self.data;
269        data.fmt(f)
270    }
271}
272impl<T: Clear + PartialEq + InPoolSet<PoolSet>> PartialEq for Pooled<T> {
273    fn eq(&self, other: &Self) -> bool {
274        // This form rid of a spuriou clippy warning about unconditional recursion.
275        <T as PartialEq>::eq(&self.data, &other.data)
276    }
277}
278
279impl<T: Clear + InPoolSet<PoolSet> + Eq> Eq for Pooled<T> {}
280
281impl<T: Clear + Hash + InPoolSet<PoolSet>> Hash for Pooled<T> {
282    fn hash<H: Hasher>(&self, state: &mut H) {
283        self.data.hash(state)
284    }
285}
286
287impl<T: Clear + InPoolSet<PoolSet> + 'static> Pooled<T> {
288    /// Clear the contents the wrapped object. If the object cannot be reused,
289    /// attempt to fetch another value from the pool.
290    ///
291    /// This method can be used in concert with `relinquish` to provide a
292    /// `clear` operation that hands data back to the pool, and then grabs it
293    /// back again if it needs to be reused.
294    ///
295    /// This pattern is likely only suitable for "temporary" buffers.
296    pub(crate) fn refresh(this: &mut Pooled<T>) {
297        this.data.clear();
298        if this.data.reuse() {
299            return;
300        }
301        let pool = with_pool_set(|ps| ps.get_pool::<T>());
302        let mut other = pool.data.borrow_mut().pop();
303        if !other.reuse() {
304            return;
305        }
306        let slot: &mut T = &mut this.data;
307        mem::swap(slot, &mut other);
308    }
309
310    pub(crate) fn into_inner(this: Pooled<T>) -> T {
311        // SAFETY: ownership of `this.data` is transferred to the caller. We
312        // will not drop `this` or use it again.
313        let inner = unsafe { ptr::read(&this.data) };
314        mem::forget(this);
315        ManuallyDrop::into_inner(inner)
316    }
317
318    pub(crate) fn new(data: T) -> Pooled<T> {
319        Pooled {
320            data: ManuallyDrop::new(data),
321        }
322    }
323}
324
325impl<T: Clear + Clone + InPoolSet<PoolSet>> Pooled<T> {
326    pub(crate) fn cloned(this: &Pooled<T>) -> Pooled<T> {
327        let mut res = with_pool_set(|ps| ps.get::<T>());
328        res.clone_from(this);
329        res
330    }
331}
332
333impl<T: Clear + InPoolSet<PoolSet>> Drop for Pooled<T> {
334    fn drop(&mut self) {
335        let reuse = self.data.reuse();
336        if !reuse {
337            // SAFETY: we own `self.data` and being in the drop method means no
338            // one else will access it.
339            unsafe { ManuallyDrop::drop(&mut self.data) };
340            return;
341        }
342        self.data.clear();
343        let t: &T = &self.data;
344        // SAFETY: ownership of `self.data` is transferred to the pool
345        with_pool_set(|ps| {
346            T::with_pool(ps, |pool| {
347                pool.data.borrow_mut().push(unsafe { ptr::read(t) })
348            })
349        });
350    }
351}
352
353impl<T: Clear + InPoolSet<PoolSet>> Deref for Pooled<T> {
354    type Target = T;
355
356    fn deref(&self) -> &T {
357        &self.data
358    }
359}
360
361impl<T: Clear + InPoolSet<PoolSet>> DerefMut for Pooled<T> {
362    fn deref_mut(&mut self) -> &mut T {
363        &mut self.data
364    }
365}
366
367/// Helper trait for allowing the trait resolution system to infer the correct
368/// pool type during allocation.
369pub trait InPoolSet<PoolSet>
370where
371    Self: Sized + Clear,
372{
373    fn with_pool<R>(pool_set: &PoolSet, f: impl FnOnce(&Pool<Self>) -> R) -> R;
374}
375
376macro_rules! pool_set {
377    ($vis:vis $name:ident { $($ident:ident : $ty:ty [ $bytes:expr ],)* }) => {
378        $vis struct $name {
379            $(
380                $ident: Pool<$ty>,
381            )*
382        }
383
384        impl Default for $name {
385            fn default() -> Self {
386                $name {
387                $(
388                    $ident: Pool::new($bytes),
389                )*
390                }
391            }
392        }
393
394        impl $name {
395            $vis fn get_pool<T: InPoolSet<Self>>(&self) -> Pool<T> {
396                T::with_pool(self, Pool::clone)
397            }
398
399            $vis fn get<T: InPoolSet<Self> + Default>(&self) -> Pooled<T> {
400                T::with_pool(self, |pool| pool.get())
401            }
402            $vis fn clear(&self) {
403                $( self.$ident.clear(); )*
404            }
405        }
406
407        $(
408            impl InPoolSet<$name> for $ty {
409                fn with_pool<R>(pool_set: &$name, f: impl FnOnce(&Pool<Self>) -> R) -> R {
410                    f(&pool_set.$ident)
411                }
412            }
413        )*
414    }
415}
416
417// The main thread-local memory pool used for reusing allocations. The syntax is:
418//
419// <name> : <type> [ <bytes> ],
420//
421// Where `name` is not used for anything, `type` feeds into the `InPoolSet` machinery and allows
422// anything of that type to be allocated using `with_pool_set`, and `bytes` is a per-type limit on
423// the total bytes that can be buffered in a single (per-thread) memory pool.
424
425pool_set! {
426    pub PoolSet {
427        vec_vals: Vec<Value> [ 1 << 25 ],
428        vec_cell_vals: Vec<Cell<Value>> [ 1 << 25 ],
429        // TODO: work on scaffolding/DI/etc. so that we can share allocations
430        // between vec_vals and shared_vals.
431        rows: Vec<RowId> [ 1 << 25 ],
432        offset_vec: SortedOffsetVector [ 1 << 20 ],
433        column_index: IndexMap<Value, BufferedSubset> [ 1 << 20 ],
434        constraints: Vec<Constraint> [ 1 << 20 ],
435        bitsets: FixedBitSet [ 1 << 20 ],
436        instrs: Vec<Instr> [ 1 << 20 ],
437        tuple_indexes: HashTable<TableEntry<BufferedSubset>> [ 1 << 20 ],
438        staged_outputs: HashTable<SwTableEntry> [ 1 << 25 ],
439        predicted_vals: PredictedVals [ 1 << 20 ],
440        shard_hist: DenseIdMap<ShardId, usize> [ 1 << 20 ],
441        instr_indexes: Vec<u32> [ 1 << 20 ],
442        cached_subsets: IdVec<ColumnId, std::sync::OnceLock<std::sync::Arc<SortedColumnIndex>>> [ 4 << 20 ],
443        intersected_on: DenseIdMap<AtomId, i64> [ 1 << 20 ],
444
445        cached_child: IdVec<ColumnId, std::sync::RwLock<HashMap<Value, (std::sync::Arc<TrieNode>, Box<[Constraint]>)>>> [ 1 << 20 ],
446    }
447}
448
449/// Run `f` on the thread-local [`PoolSet`].
450pub(crate) fn with_pool_set<R>(f: impl FnOnce(&PoolSet) -> R) -> R {
451    POOL_SET.with(|pool_set| f(pool_set))
452}
453
454thread_local! {
455    /// A thread-local pool set. All pooled allocations land back in the local thread.
456    ///
457    /// We don't drop this PoolSet because it does not contain any resources
458    /// that need to be released, other than memory (which will be reclaimed
459    /// when the process exits, right after drop runs).
460    ///
461    /// For large egraphs, this be a big runtime win. The main egglog binary
462    /// avoids dropping the egraph for the same reason.
463    static POOL_SET: ManuallyDrop<PoolSet> = Default::default();
464}
egglog_core_relations/pool/mod.rs

egglog_core_relations/pool/
mod.rs
