Implement memory limits in Global allocator

crates/rune-alloc/Cargo.toml

@@ -21,6 +21,7 @@ alloc = []
 [dependencies]
 serde = { version = "1.0", optional = true }
 ahash = { version = "0.8.3", default-features = false }
+pin-project = "1.1.0"
 
 [dev-dependencies]
 rand = { version = "0.8.5", features = ["small_rng"] }

crates/rune-alloc/src/alloc/allocator.rs

@@ -3,9 +3,7 @@
 use core::alloc::Layout;
 use core::fmt;
 
-use crate::ptr::{self, invalid_mut, NonNull};
-
-use ::rust_alloc::alloc::{alloc, alloc_zeroed, dealloc};
+use crate::ptr::{self, NonNull};
 
 /// Error raised while allocating.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -299,8 +297,7 @@ where
 
     #[inline]
     unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
-        // SAFETY: the safety contract must be upheld by the caller
-        unsafe { (**self).deallocate(ptr, layout) }
+        (**self).deallocate(ptr, layout)
     }
 
     #[inline]
@@ -310,8 +307,7 @@ where
         old_layout: Layout,
         new_layout: Layout,
     ) -> Result<NonNull<[u8]>, AllocError> {
-        // SAFETY: the safety contract must be upheld by the caller
-        unsafe { (**self).grow(ptr, old_layout, new_layout) }
+        (**self).grow(ptr, old_layout, new_layout)
     }
 
     #[inline]
@@ -321,63 +317,6 @@ where
         old_layout: Layout,
         new_layout: Layout,
     ) -> Result<NonNull<[u8]>, AllocError> {
-        // SAFETY: the safety contract must be upheld by the caller
-        unsafe { (**self).shrink(ptr, old_layout, new_layout) }
-    }
-}
-
-/// The default global allocator.
-#[derive(Default, Debug, Clone)]
-pub struct Global;
-
-impl Global {
-    #[inline]
-    fn alloc_impl(&self, layout: Layout, zeroed: bool) -> Result<NonNull<[u8]>, AllocError> {
-        /// Creates a `NonNull` that is dangling, but well-aligned for this Layout.
-        ///
-        /// Note that the pointer value may potentially represent a valid pointer, which
-        /// means this must not be used as a "not yet initialized" sentinel value. Types
-        /// that lazily allocate must track initialization by some other means.
-        pub(crate) const fn dangling<T>(layout: &Layout) -> NonNull<T> {
-            unsafe { NonNull::new_unchecked(invalid_mut::<T>(layout.align())) }
-        }
-
-        match layout.size() {
-            0 => Ok(NonNull::slice_from_raw_parts(dangling(&layout), 0)),
-            // SAFETY: `layout` is non-zero in size,
-            size => unsafe {
-                let raw_ptr = if zeroed {
-                    alloc_zeroed(layout)
-                } else {
-                    alloc(layout)
-                };
-
-                let Some(ptr) = NonNull::new(raw_ptr) else {
-                    return Err(AllocError { layout });
-                };
-
-                Ok(NonNull::slice_from_raw_parts(ptr, size))
-            },
-        }
-    }
-}
-
-unsafe impl Allocator for Global {
-    #[inline]
-    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
-        self.alloc_impl(layout, false)
-    }
-
-    #[inline]
-    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
-        self.alloc_impl(layout, true)
-    }
-
-    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
-        if layout.size() != 0 {
-            // SAFETY: `layout` is non-zero in size,
-            // other conditions must be upheld by the caller
-            unsafe { dealloc(ptr.as_ptr(), layout) }
-        }
+        (**self).shrink(ptr, old_layout, new_layout)
     }
 }

crates/rune-alloc/src/alloc/boxed.rs

@@ -32,10 +32,10 @@ impl<T> Box<T, Global> {
     /// ```
     /// use rune_alloc::Box;
     ///
-    /// let five = Box::new(5)?;
+    /// let five = Box::try_new(5)?;
     /// # Ok::<_, rune_alloc::AllocError>(())
     /// ```
-    pub fn new(value: T) -> Result<Self, AllocError> {
+    pub fn try_new(value: T) -> Result<Self, AllocError> {
         Self::try_new_in(value, Global)
     }
 }
@@ -148,7 +148,7 @@ impl<T: ?Sized, A: Allocator> Box<T, A> {
     /// # fn main() -> Result<(), rune_alloc::Error> {
     /// use rune_alloc::Box;
     ///
-    /// let x = Box::new(41)?;
+    /// let x = Box::try_new(41)?;
     /// let static_ref: &'static mut usize = Box::leak(x);
     /// *static_ref += 1;
     /// assert_eq!(*static_ref, 42);

crates/rune-alloc/src/alloc/global.rs

@@ -0,0 +1,72 @@
+use core::alloc::Layout;
+
+use crate::alloc::{AllocError, Allocator};
+use crate::ptr::{invalid_mut, NonNull};
+
+use ::rust_alloc::alloc::{alloc, alloc_zeroed, dealloc};
+
+/// Creates a `NonNull` that is dangling, but well-aligned for this Layout.
+///
+/// Note that the pointer value may potentially represent a valid pointer, which
+/// means this must not be used as a "not yet initialized" sentinel value. Types
+/// that lazily allocate must track initialization by some other means.
+pub(crate) const fn dangling<T>(layout: &Layout) -> NonNull<T> {
+    unsafe { NonNull::new_unchecked(invalid_mut::<T>(layout.align())) }
+}
+
+/// The default global allocator for Rune.
+///
+/// This supports enforcing thread-local memory limits through the [`limit`]
+/// module.
+///
+/// [`limit`]: crate::limit
+#[derive(Default, Debug, Clone, Copy)]
+pub struct Global;
+
+impl Global {
+    #[inline]
+    fn alloc_impl(&self, layout: Layout, zeroed: bool) -> Result<NonNull<[u8]>, AllocError> {
+        if !crate::limit::take(layout.size()) {
+            return Err(AllocError { layout });
+        }
+
+        match layout.size() {
+            0 => Ok(NonNull::slice_from_raw_parts(dangling(&layout), 0)),
+            // SAFETY: `layout` is non-zero in size,
+            size => unsafe {
+                let raw_ptr = if zeroed {
+                    alloc_zeroed(layout)
+                } else {
+                    alloc(layout)
+                };
+
+                let Some(ptr) = NonNull::new(raw_ptr) else {
+                    return Err(AllocError { layout });
+                };
+
+                Ok(NonNull::slice_from_raw_parts(ptr, size))
+            },
+        }
+    }
+}
+
+unsafe impl Allocator for Global {
+    #[inline]
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        self.alloc_impl(layout, false)
+    }
+
+    #[inline]
+    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        self.alloc_impl(layout, true)
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        if layout.size() != 0 {
+            // SAFETY: `layout` is non-zero in size,
+            // other conditions must be upheld by the caller
+            dealloc(ptr.as_ptr(), layout);
+            crate::limit::release(layout.size());
+        }
+    }
+}

crates/rune-alloc/src/alloc/mod.rs

@@ -5,9 +5,12 @@ pub(crate) mod raw_vec;
 pub use self::borrow::TryToOwned;
 pub(crate) mod borrow;
 
-pub use self::allocator::{AllocError, Allocator, Global};
+pub use self::allocator::{AllocError, Allocator};
 pub(crate) mod allocator;
 
+pub use self::global::Global;
+pub(crate) mod global;
+
 pub use self::boxed::Box;
 pub mod boxed;
 

crates/rune-alloc/src/alloc/slice.rs

@@ -4,17 +4,8 @@ use crate::alloc::{Allocator, Box, Error, Global, TryClone, TryToOwned, Vec};
 ///
 /// The resulting vector can be converted back into a box via
 /// `Vec<T>`'s `into_boxed_slice` method.
-///
-/// # Examples
-///
-/// ```
-/// let s: Box<[i32]> = Box::new([10, 40, 30]);
-/// let x = s.into_vec();
-/// // `s` cannot be used anymore because it has been converted into `x`.
-///
-/// assert_eq!(x, vec![10, 40, 30]);
-/// ```
 #[inline]
+#[doc(hidden)]
 pub fn into_vec<T, A: Allocator>(this: Box<[T], A>) -> Vec<T, A> {
     // N.B., see the `hack` module in this file for more details.
     hack::into_vec(this)

crates/rune-alloc/src/alloc/string.rs

@@ -1736,7 +1736,7 @@ impl<A: Allocator> FromUtf8Error<A> {
     }
 }
 
-impl Borrow<str> for String {
+impl<A: Allocator> Borrow<str> for String<A> {
     #[inline]
     fn borrow(&self) -> &str {
         &self[..]

crates/rune-alloc/src/alloc/vec/mod.rs

@@ -84,6 +84,7 @@ use self::is_zero::IsZero;
 #[cfg(rune_nightly)]
 mod is_zero;
 
+#[doc(hidden)]
 pub use crate::alloc::slice::into_vec;
 
 use core::borrow::Borrow;
@@ -2899,16 +2900,24 @@ impl<T, A: Allocator, const N: usize> TryFrom<Vec<T, A>> for [T; N] {
 }
 
 impl<T, A: Allocator> From<Box<[T], A>> for Vec<T, A> {
-    /// Convert a boxed slice into a vector by transferring ownership of
-    /// the existing heap allocation.
+    /// Convert a boxed slice into a vector by transferring ownership of the
+    /// existing heap allocation.
     ///
     /// # Examples
     ///
     /// ```
-    /// use rune_alloc::{Vec, Box};
+    /// use rune_alloc::{Box, Vec};
+    /// use rune_alloc::try_vec;
+    ///
+    /// let s: Box<[i32]> = Box::try_from([10, 40, 30])?;
+    /// let x: Vec<i32> = Vec::from(s);
+    ///
+    /// assert_eq!(x, [10, 40, 30]);
+    ///
+    /// let s: Box<[i32]> = try_vec![10, 40, 30].try_into_boxed_slice()?;
+    /// let x: Vec<i32> = Vec::from(s);
     ///
-    /// let b: Box<[i32]> = rune_alloc::try_vec![1, 2, 3].try_into_boxed_slice()?;
-    /// assert_eq!(Vec::from(b), rune_alloc::try_vec![1, 2, 3]);
+    /// assert_eq!(x, [10, 40, 30]);
     /// # Ok::<_, rune_alloc::Error>(())
     /// ```
     fn from(s: Box<[T], A>) -> Self {

crates/rune-alloc/src/lib.rs

@@ -94,6 +94,8 @@ pub mod prelude {
     };
 }
 
+pub mod limit;
+
 #[cfg(test)]
 mod testing;