mz_timely_util/
containers.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Reusable containers.
17
18pub mod stack;
19
20pub(crate) use alloc::alloc_aligned_zeroed;
21pub use alloc::{enable_columnar_lgalloc, set_enable_columnar_lgalloc};
22pub use provided_builder::ProvidedBuilder;
23
24mod alloc {
25    use mz_ore::region::Region;
26
27    /// Allocate a region of memory with a capacity of at least `len` that is properly aligned
28    /// and zeroed. The memory in Regions is always aligned to its content type.
29    #[inline]
30    pub(crate) fn alloc_aligned_zeroed<T: bytemuck::AnyBitPattern>(len: usize) -> Region<T> {
31        if enable_columnar_lgalloc() {
32            Region::new_auto_zeroed(len)
33        } else {
34            Region::new_heap_zeroed(len)
35        }
36    }
37
38    thread_local! {
39        static ENABLE_COLUMNAR_LGALLOC: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
40    }
41
42    /// Returns `true` if columnar allocations should come from lgalloc.
43    #[inline]
44    pub fn enable_columnar_lgalloc() -> bool {
45        ENABLE_COLUMNAR_LGALLOC.get()
46    }
47
48    /// Set whether columnar allocations should come from lgalloc. Applies to future allocations.
49    pub fn set_enable_columnar_lgalloc(enabled: bool) {
50        ENABLE_COLUMNAR_LGALLOC.set(enabled);
51    }
52}
53
54mod provided_builder {
55    use timely::Container;
56    use timely::container::ContainerBuilder;
57
58    /// A container builder that doesn't support pushing elements, and is only suitable for pushing
59    /// whole containers at Timely sessions. See [`give_container`] for more information.
60    ///
61    ///  [`give_container`]: timely::dataflow::channels::pushers::buffer::Session::give_container
62    pub struct ProvidedBuilder<C> {
63        _marker: std::marker::PhantomData<C>,
64    }
65
66    impl<C> Default for ProvidedBuilder<C> {
67        fn default() -> Self {
68            Self {
69                _marker: std::marker::PhantomData,
70            }
71        }
72    }
73
74    impl<C: Container + Clone + 'static> ContainerBuilder for ProvidedBuilder<C> {
75        type Container = C;
76
77        #[inline(always)]
78        fn extract(&mut self) -> Option<&mut Self::Container> {
79            None
80        }
81
82        #[inline(always)]
83        fn finish(&mut self) -> Option<&mut Self::Container> {
84            None
85        }
86    }
87}
mz_timely_util/containers.rs

mz_timely_util/
containers.rs
