Allow explicit handling of allocation errors (pgcentralfoundation#2226)

Right now, we have several fallible constructors for Postgres types.
When this is an allocation error, we don't allow handling it.
As I expect there to be more such fallible constructors in the future,
and transient memory contexts can be arbitrarily limited in size,
I think we should expose this problem in otherwise-fallible code.
We can retain some infallible constructors elsewhere, though.
If nothing else, this can allow better diagnostics via `track_caller`.

Resolves pgcentralfoundation#2225

Author: workingjubilee

pgrx/src/array/flat_array.rs

@@ -1,7 +1,7 @@
 #![deny(missing_docs)]
 use crate::datum::{Array, BorrowDatum, Datum};
 use crate::layout::{Align, Layout};
-use crate::memcx::MemCx;
+use crate::memcx::{MemCx, OutOfMemory};
 use crate::nullable::Nullable;
 use crate::palloc::PBox;
 use crate::pgrx_sql_entity_graph::metadata::{
@@ -77,6 +77,8 @@ pub enum ArrayAllocError {
     TooManyElems,
     /// One or more dimensions are zero
     ZeroLenDim,
+    /// There is insufficient memory in this context
+    OutOfMemory,
 }
 
 const MAX_ALLOC_SIZE: usize = 0x3fffffff;
@@ -104,7 +106,7 @@ where
 
         let ndims = N;
         if N == 0 {
-            return Ok(FlatArray::new_empty(memcx));
+            return FlatArray::new_empty(memcx);
         }
         const { assert!(N <= MAX_DIMS) };
 
@@ -151,7 +153,7 @@ where
         let elemtype = <T as Scalar>::OID;
         let tail_size = size - base_size;
 
-        let ptr = alloc_zeroed_head(memcx, tail_size, ndims as i32, dataoffset, elemtype);
+        let ptr = alloc_zeroed_head(memcx, tail_size, ndims as i32, dataoffset, elemtype)?;
         // SAFETY: we allocated enough for our dimensions and lbounds
         unsafe {
             ptr.byte_add(base_size).cast().write(dim_ints);
@@ -164,11 +166,13 @@ where
     }
 
     /// Allocate a 0-dimension array
-    pub fn new_empty<'cx>(memcx: &MemCx<'cx>) -> PBox<'cx, FlatArray<'cx, T>> {
+    pub fn new_empty<'cx>(
+        memcx: &MemCx<'cx>,
+    ) -> Result<PBox<'cx, FlatArray<'cx, T>>, ArrayAllocError> {
         let nbytes = mem::size_of::<pg_sys::ArrayType>();
-        let ptr = alloc_zeroed_head(memcx, 0, 0, 0, <T as Scalar>::OID);
+        let ptr = alloc_zeroed_head(memcx, 0, 0, 0, <T as Scalar>::OID)?;
         // SAFETY: it's valid, if 0-dimensional
-        unsafe { PBox::from_raw_in(FlatArray::cast_tailed(ptr), memcx) }
+        Ok(unsafe { PBox::from_raw_in(FlatArray::cast_tailed(ptr), memcx) })
     }
 
     /// Allocate an array sized to fit a slice and copy it
@@ -179,7 +183,7 @@ where
         memcx: &MemCx<'cx>,
     ) -> Result<PBox<'cx, FlatArray<'cx, T>>, ArrayAllocError> {
         match data.len() {
-            0 => Ok(FlatArray::new_empty(memcx)),
+            0 => FlatArray::new_empty(memcx),
             len => {
                 let mut array = FlatArray::new_zeroed_in([len], false, memcx)?;
                 array.as_non_null_slice_mut().unwrap().copy_from_slice(data);
@@ -291,9 +295,9 @@ fn alloc_zeroed_head(
     ndim: ffi::c_int,
     dataoffset: i32,
     elemtype: pg_sys::Oid,
-) -> ptr::NonNull<[u8]> {
+) -> Result<ptr::NonNull<[u8]>, ArrayAllocError> {
     let nbytes = size_of::<pg_sys::ArrayType>() + tail_size;
-    let ptr = memcx.alloc_zeroed_bytes(nbytes);
+    let ptr = memcx.alloc_zeroed_bytes(nbytes)?;
     // SAFETY: We allocated at least the ArrayType's prefix of bytes
     unsafe {
         // COMPAT: write ArrayType so fields must be initialized even if ArrayType changes
@@ -305,7 +309,7 @@ fn alloc_zeroed_head(
             elemtype,
         })
     }
-    ptr::NonNull::slice_from_raw_parts(ptr, tail_size)
+    Ok(ptr::NonNull::slice_from_raw_parts(ptr, tail_size))
 }
 
 unsafe impl<T: ?Sized> BorrowDatum for FlatArray<'_, T> {
@@ -408,3 +412,9 @@ where
 
 impl<'arr, T> ExactSizeIterator for ArrayIter<'arr, T> where T: ?Sized + Element {}
 impl<'arr, T> FusedIterator for ArrayIter<'arr, T> where T: ?Sized + Element {}
+
+impl From<OutOfMemory> for ArrayAllocError {
+    fn from(value: OutOfMemory) -> Self {
+        ArrayAllocError::OutOfMemory
+    }
+}

pgrx/src/list/flat_list.rs

@@ -112,7 +112,8 @@ impl<'cx, T: Enlist> List<'cx, T> {
                 // No silly reasoning, simply allocate ~2 cache lines for a list
                 let list_size = 128;
                 unsafe {
-                    let list: *mut pg_sys::List = mcx.alloc_bytes(list_size).cast().as_ptr();
+                    let list: *mut pg_sys::List =
+                        mcx.alloc_bytes(list_size).unwrap().cast().as_ptr();
                     assert!(list.is_non_null());
                     (*list).type_ = T::LIST_TAG;
                     (*list).max_length = ((list_size - mem::size_of::<pg_sys::List>())

pgrx/src/memcx.rs

@@ -36,18 +36,18 @@ impl<'mcx> MemCx<'mcx> {
 
     /// Allocate a raw byte buffer `size` bytes in length
     /// and returns a pointer to the new allocation.
-    pub fn alloc_bytes(&self, size: usize) -> NonNull<u8> {
-        let ptr = unsafe { pg_sys::MemoryContextAlloc(self.ptr.as_ptr(), size).cast() };
-        // SAFETY: `pg_sys::MemoryContextAlloc` is not allowed to return NULL, it must error.
-        unsafe { NonNull::new_unchecked(ptr) }
+    pub fn alloc_bytes(&self, size: usize) -> Result<NonNull<u8>, OutOfMemory> {
+        let flags = (pg_sys::MCXT_ALLOC_NO_OOM) as i32;
+        let ptr = unsafe { pg_sys::MemoryContextAllocExtended(self.ptr.as_ptr(), size, flags) };
+        NonNull::new(ptr.cast()).ok_or(OutOfMemory::new())
     }
 
     /// Allocate a raw byte buffer `size` bytes in length
     /// and returns a pointer to the new allocation.
-    pub fn alloc_zeroed_bytes(&self, size: usize) -> NonNull<u8> {
-        let ptr = unsafe { pg_sys::MemoryContextAllocZero(self.ptr.as_ptr(), size).cast() };
-        // SAFETY: `pg_sys::MemoryContextAlloc` is not allowed to return NULL, it must error.
-        unsafe { NonNull::new_unchecked(ptr) }
+    pub fn alloc_zeroed_bytes(&self, size: usize) -> Result<NonNull<u8>, OutOfMemory> {
+        let flags = (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_ZERO) as i32;
+        let ptr = unsafe { pg_sys::MemoryContextAllocExtended(self.ptr.as_ptr(), size, flags) };
+        NonNull::new(ptr.cast()).ok_or(OutOfMemory::new())
     }
 
     /// Stores the current memory context, switches to *this* memory context,
@@ -68,6 +68,16 @@ impl<'mcx> MemCx<'mcx> {
     }
 }
 
+#[derive(Debug)]
+pub struct OutOfMemory {
+    _reserve: (),
+}
+impl OutOfMemory {
+    pub fn new() -> OutOfMemory {
+        OutOfMemory { _reserve: () }
+    }
+}
+
 /// Acquire the current context and operate inside it.
 pub fn current_context<'curr, F, T>(f: F) -> T
 where

pgrx/src/palloc/pbox.rs

@@ -1,7 +1,7 @@
 use crate::callconv::{BoxRet, FcInfo};
 use crate::datum::{BorrowDatum, Datum};
 use crate::layout::PassBy;
-use crate::memcx::MemCx;
+use crate::memcx::{MemCx, OutOfMemory};
 use crate::pg_sys;
 use core::marker::PhantomData;
 use core::ops::{Deref, DerefMut};
@@ -33,13 +33,18 @@ impl<'mcx, T: ?Sized> PBox<'mcx, T> {
 }
 
 impl<'mcx, T: Sized> PBox<'mcx, T> {
+    #[track_caller]
     pub fn new_in(val: T, memcx: &MemCx<'mcx>) -> Self {
+        PBox::try_new_in(val, memcx).unwrap()
+    }
+
+    pub fn try_new_in(val: T, memcx: &MemCx<'mcx>) -> Result<Self, OutOfMemory> {
         const { assert!(align_of::<T>() <= size_of::<pg_sys::Datum>()) };
-        let ptr = memcx.alloc_bytes(size_of::<T>()).cast();
+        let ptr = memcx.alloc_bytes(size_of::<T>())?.cast();
         // SAFETY: We were guaranteed an appropriately sized allocation to write to,
         // and we have asserted our alignment maximum was upheld
         unsafe { ptr.write(val) };
-        PBox { ptr, _cx: PhantomData }
+        Ok(PBox { ptr, _cx: PhantomData })
     }
 }