Support lgalloc for columnar (#31230)

Adds support for allocating memory for columnar's aligned type from
lgalloc. Adds a flag `enable_lgalloc_columnar` to control the behavior.

### Checklist

- [ ] This PR has adequate test coverage / QA involvement has been duly
considered. ([trigger-ci for additional test/nightly
runs](https://trigger-ci.dev.materialize.com/))
- [ ] This PR has an associated up-to-date [design
doc](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/README.md),
is a design doc
([template](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/00000000_template.md)),
or is sufficiently small to not require a design.
  <!-- Reference the design in the description. -->
- [ ] If this PR evolves [an existing `$T ⇔ Proto$T`
mapping](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/command-and-response-binary-encoding.md)
(possibly in a backwards-incompatible way), then it is tagged with a
`T-proto` label.
- [ ] If this PR will require changes to cloud orchestration or tests,
there is a companion cloud PR to account for those changes that is
tagged with the release-blocker label
([example](MaterializeInc/cloud#5021)).
<!-- Ask in #team-cloud on Slack if you need help preparing the cloud
PR. -->
- [ ] If this PR includes major [user-facing behavior
changes](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/guide-changes.md#what-changes-require-a-release-note),
I have pinged the relevant PM to schedule a changelog post.

Signed-off-by: Moritz Hoffmann <[email protected]>

Author: antiguru

Cargo.lock

@@ -57,6 +57,13 @@ pub const ENABLE_COLUMNATION_LGALLOC: Config<bool> = Config::new(
     "Enable allocating regions from lgalloc.",
 );
 
+/// Enable lgalloc for columnar.
+pub const ENABLE_COLUMNAR_LGALLOC: Config<bool> = Config::new(
+    "enable_columnar_lgalloc",
+    true,
+    "Enable allocating aligned regions in columnar from lgalloc.",
+);
+
 /// Enable lgalloc's eager memory return/reclamation feature.
 pub const ENABLE_LGALLOC_EAGER_RECLAMATION: Config<bool> = Config::new(
     "enable_lgalloc_eager_reclamation",
@@ -188,6 +195,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&ENABLE_LGALLOC)
         .add(&ENABLE_COLUMNATION_LGALLOC)
         .add(&ENABLE_LGALLOC_EAGER_RECLAMATION)
+        .add(&ENABLE_COLUMNAR_LGALLOC)
         .add(&COMPUTE_SERVER_MAINTENANCE_INTERVAL)
         .add(&DATAFLOW_MAX_INFLIGHT_BYTES)
         .add(&DATAFLOW_MAX_INFLIGHT_BYTES_CC)

src/compute-types/src/dyncfgs.rs

@@ -303,6 +303,9 @@ impl ComputeState {
             std::sync::atomic::Ordering::Relaxed,
         );
 
+        let enable_columnar_lgalloc = ENABLE_COLUMNAR_LGALLOC.get(config);
+        mz_timely_util::containers::set_enable_columnar_lgalloc(enable_columnar_lgalloc);
+
         // Remember the maintenance interval locally to avoid reading it from the config set on
         // every server iteration.
         self.server_maintenance_interval = COMPUTE_SERVER_MAINTENANCE_INTERVAL.get(config);

src/compute/src/compute_state.rs

@@ -28,6 +28,7 @@ rust_library(
         "assert",
         "async",
         "async-trait",
+        "bytemuck",
         "bytes",
         "bytes_",
         "capture",
@@ -110,6 +111,7 @@ rust_test(
         "assert",
         "async",
         "async-trait",
+        "bytemuck",
         "bytes",
         "bytes_",
         "capture",

src/ore/BUILD.bazel

@@ -19,6 +19,7 @@ anyhow = { version = "1.0.95", optional = true }
 # The only exception to the above is the `either` crate, which itself has zero
 # dependencies and is widely considered to be basically part of the stdlib.
 async-trait = { version = "0.1.83", optional = true }
+bytemuck = { version = "1.21.0", optional = true }
 bytes = { version = "1.3.0", optional = true }
 chrono = { version = "0.4.39", default-features = false, features = [
   "std",
@@ -119,7 +120,7 @@ async = [
 bytes_ = ["bytes", "compact_bytes", "smallvec", "smallvec/const_generics", "region", "tracing_"]
 network = ["async", "bytes", "smallvec", "tonic", "tracing"]
 process = ["libc"]
-region = ["lgalloc"]
+region = ["dep:lgalloc", "dep:bytemuck"]
 tracing_ = [
   "anyhow",
   "tracing",

src/ore/Cargo.toml

@@ -170,6 +170,7 @@ impl<T> LgAllocRegion<T> {
 /// NOTE: We plan to deprecate this type soon. Users should switch to different types or the raw
 /// `lgalloc` API instead.
 #[derive(Debug)]
+#[must_use]
 pub enum Region<T> {
     /// A possibly empty heap-allocated region, represented as a vector.
     Heap(Vec<T>),
@@ -218,14 +219,12 @@ impl<T> Default for Region<T> {
 impl<T> Region<T> {
     /// Create a new empty region.
     #[inline]
-    #[must_use]
     pub fn new_empty() -> Region<T> {
         Region::Heap(Vec::new())
     }
 
     /// Create a new heap-allocated region of a specific capacity.
     #[inline]
-    #[must_use]
     pub fn new_heap(capacity: usize) -> Region<T> {
         Region::Heap(Vec::with_capacity(capacity))
     }
@@ -255,7 +254,6 @@ impl<T> Region<T> {
     /// but can be larger if the implementation requires it.
     ///
     /// Returns a [`Region::MMap`] if possible, and falls back to [`Region::Heap`] otherwise.
-    #[must_use]
     pub fn new_auto(capacity: usize) -> Region<T> {
         if ENABLE_LGALLOC_REGION.load(std::sync::atomic::Ordering::Relaxed) {
             match Region::new_mmap(capacity) {
@@ -340,6 +338,8 @@ impl<T> Region<T> {
     /// Otherwise, the vector representation could try to reallocate the underlying memory
     /// using the global allocator, which would cause problems because the memory might not
     /// have originated from it. This is undefined behavior.
+    ///
+    /// Private because it is too dangerous to expose to the public.
     #[inline]
     unsafe fn as_mut_vec(&mut self) -> &mut Vec<T> {
         match self {
@@ -349,6 +349,50 @@ impl<T> Region<T> {
     }
 }
 
+impl<T: bytemuck::AnyBitPattern> Region<T> {
+    /// Create a new file-based mapped region of a specific capacity, initialized to 0. The
+    /// capacity of the returned region can be larger than requested to accommodate page sizes.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the memory allocation fails.
+    #[inline(always)]
+    pub fn new_mmap_zeroed(capacity: usize) -> Result<Self, lgalloc::AllocError> {
+        let (ptr, capacity, handle) = lgalloc::allocate::<T>(capacity)?;
+        // SAFETY: `allocate` returns a valid memory block, and `T` supports a null-bit pattern.
+        unsafe { ptr.as_ptr().write_bytes(0, capacity) }
+        // SAFETY: `ptr` points to suitable memory.
+        // It is UB to call `from_raw_parts` with a pointer not allocated from the global
+        // allocator, but we accept this here because we promise never to reallocate the vector.
+        let inner =
+            ManuallyDrop::new(unsafe { Vec::from_raw_parts(ptr.as_ptr(), capacity, capacity) });
+        let handle = Some(handle);
+        Ok(Self::MMap(MMapRegion { inner, handle }))
+    }
+
+    /// Allocate a zeroed region on the heap.
+    #[inline(always)]
+    pub fn new_heap_zeroed(capacity: usize) -> Self {
+        Self::Heap(vec![T::zeroed(); capacity])
+    }
+
+    /// Construct a new region with the specified capacity, initialized to 0.
+    pub fn new_auto_zeroed(capacity: usize) -> Self {
+        if ENABLE_LGALLOC_REGION.load(std::sync::atomic::Ordering::Relaxed) {
+            match Region::new_mmap_zeroed(capacity) {
+                Ok(r) => return r,
+                Err(lgalloc::AllocError::Disabled)
+                | Err(lgalloc::AllocError::InvalidSizeClass(_)) => {}
+                Err(e) => {
+                    eprintln!("lgalloc error: {e}, falling back to heap");
+                }
+            }
+        }
+        // Fall-through
+        Self::new_heap_zeroed(capacity)
+    }
+}
+
 impl<T: Clone> Region<T> {
     /// Extend the region from a slice.
     ///

src/ore/src/region.rs

@@ -24,14 +24,49 @@ use timely::container::columnation::TimelyStack;
 
 pub mod stack;
 
+pub(crate) use alloc::alloc_aligned_zeroed;
+pub use alloc::{enable_columnar_lgalloc, set_enable_columnar_lgalloc};
+pub use builder::ColumnBuilder;
 pub use container::Column;
+pub use provided_builder::ProvidedBuilder;
+
+mod alloc {
+    use mz_ore::region::Region;
+
+    /// Allocate a region of memory with a capacity of at least `len` that is properly aligned
+    /// and zeroed. The memory in Regions is always aligned to its content type.
+    #[inline]
+    pub(crate) fn alloc_aligned_zeroed<T: bytemuck::AnyBitPattern>(len: usize) -> Region<T> {
+        if enable_columnar_lgalloc() {
+            Region::new_auto_zeroed(len)
+        } else {
+            Region::new_heap_zeroed(len)
+        }
+    }
+
+    thread_local! {
+        static ENABLE_COLUMNAR_LGALLOC: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
+    }
+
+    /// Returns `true` if columnar allocations should come from lgalloc.
+    #[inline]
+    pub fn enable_columnar_lgalloc() -> bool {
+        ENABLE_COLUMNAR_LGALLOC.get()
+    }
+
+    /// Set whether columnar allocations should come from lgalloc. Applies to future allocations.
+    pub fn set_enable_columnar_lgalloc(enabled: bool) {
+        ENABLE_COLUMNAR_LGALLOC.set(enabled);
+    }
+}
 
 mod container {
     use columnar::bytes::{EncodeDecode, Sequence};
     use columnar::common::IterOwn;
     use columnar::Columnar;
     use columnar::Container as _;
     use columnar::{Clear, FromBytes, Index, Len};
+    use mz_ore::region::Region;
     use timely::bytes::arc::Bytes;
     use timely::container::PushInto;
     use timely::dataflow::channels::ContainerBytes;
@@ -51,7 +86,7 @@ mod container {
         ///
         /// Reasons could include misalignment, cloning of data, or wanting
         /// to release the `Bytes` as a scarce resource.
-        Align(Box<[u64]>),
+        Align(Region<u64>),
     }
 
     impl<C: Columnar> Column<C> {
@@ -90,11 +125,15 @@ mod container {
                 Column::Typed(t) => Column::Typed(t.clone()),
                 Column::Bytes(b) => {
                     assert_eq!(b.len() % 8, 0);
-                    let mut alloc: Vec<u64> = vec![0; b.len() / 8];
+                    let mut alloc: Region<u64> = super::alloc_aligned_zeroed(b.len() / 8);
                     bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&b[..]);
-                    Self::Align(alloc.into())
+                    Self::Align(alloc)
+                }
+                Column::Align(a) => {
+                    let mut alloc = super::alloc_aligned_zeroed(a.len());
+                    alloc.extend_from_slice(&a[..]);
+                    Column::Align(alloc)
                 }
-                Column::Align(a) => Column::Align(a.clone()),
             }
         }
     }
@@ -111,8 +150,7 @@ mod container {
         fn clear(&mut self) {
             match self {
                 Column::Typed(t) => t.clear(),
-                Column::Bytes(_) => *self = Column::Typed(Default::default()),
-                Column::Align(_) => *self = Column::Typed(Default::default()),
+                Column::Bytes(_) | Column::Align(_) => *self = Column::Typed(Default::default()),
             }
         }
 
@@ -159,9 +197,9 @@ mod container {
                 Self::Bytes(bytes)
             } else {
                 // We failed to cast the slice, so we'll reallocate.
-                let mut alloc: Vec<u64> = vec![0; bytes.len() / 8];
+                let mut alloc: Region<u64> = super::alloc_aligned_zeroed(bytes.len() / 8);
                 bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&bytes[..]);
-                Self::Align(alloc.into())
+                Self::Align(alloc)
             }
         }
 
@@ -183,7 +221,6 @@ mod container {
     }
 }
 
-pub use builder::ColumnBuilder;
 mod builder {
     use std::collections::VecDeque;
 
@@ -192,7 +229,7 @@ mod builder {
     use timely::container::PushInto;
     use timely::container::{ContainerBuilder, LengthPreservingContainerBuilder};
 
-    use super::Column;
+    use crate::containers::Column;
 
     /// A container builder for `Column<C>`.
     pub struct ColumnBuilder<C: Columnar> {
@@ -219,11 +256,24 @@ mod builder {
             let words = Sequence::length_in_words(&self.current.borrow());
             let round = (words + ((1 << 18) - 1)) & !((1 << 18) - 1);
             if round - words < round / 10 {
-                let mut alloc = Vec::with_capacity(round);
-                Sequence::encode(&mut alloc, &self.current.borrow());
-                self.pending
-                    .push_back(Column::Align(alloc.into_boxed_slice()));
-                self.current.clear();
+                /// Move the contents from `current` to an aligned allocation, and push it to `pending`.
+                /// The contents must fit in `round` words (u64).
+                #[cold]
+                fn outlined_align<C>(
+                    current: &mut C::Container,
+                    round: usize,
+                    pending: &mut VecDeque<Column<C>>,
+                ) where
+                    C: Columnar,
+                {
+                    let mut alloc = super::alloc_aligned_zeroed(round);
+                    let writer = std::io::Cursor::new(bytemuck::cast_slice_mut(&mut alloc[..]));
+                    Sequence::write(writer, &current.borrow()).unwrap();
+                    pending.push_back(Column::Align(alloc));
+                    current.clear();
+                }
+
+                outlined_align(&mut self.current, round, &mut self.pending);
             }
         }
     }
@@ -390,8 +440,6 @@ pub mod batcher {
     }
 }
 
-pub use provided_builder::ProvidedBuilder;
-
 mod provided_builder {
     use timely::container::ContainerBuilder;
     use timely::Container;