fix: detoast borrowed FlatArray arguments

Author: eeeebbbbrrrr

pgrx-sql-entity-graph/src/pg_extern/mod.rs

@@ -468,8 +468,8 @@ impl PgExtern {
                             let call_flow = <#ret_ty as ::pgrx::callconv::RetAbi>::check_and_prepare(fcinfo);
                             let result = match call_flow {
                                 ::pgrx::callconv::CallCx::WrappedFn(mcx) => {
+                                    let mut #args_ident = unsafe { fcinfo.args_in(mcx) };
                                     let mut mcx = ::pgrx::PgMemoryContexts::For(mcx);
-                                    let #args_ident = &mut fcinfo.args();
                                     let call_result = mcx.switch_to(|_| {
                                         #(#arg_fetches)*
                                         #func_name( #(#arg_pats),* )

pgrx-unit-tests/src/tests/array_borrowed.rs

@@ -95,6 +95,11 @@ fn borrow_get_arr_nelems(arr: &FlatArray<'_, i32>) -> libc::c_int {
     arr.nelems() as _
 }
 
+#[pg_extern]
+fn borrow_iter_array<'a>(arr: &'a FlatArray<'a, i32>) -> SetOfIterator<'a, i32> {
+    SetOfIterator::new(arr.iter_non_null().copied())
+}
+
 #[pg_extern]
 fn borrow_get_arr_data_ptr_nth_elem(arr: &FlatArray<'_, i32>, elem: i32) -> Option<i32> {
     arr.get(elem as usize).unwrap().into_option().copied()
@@ -181,7 +186,7 @@ mod tests {
     use pgrx::datum::DatumWithOid;
     use pgrx::memcx;
     use pgrx::prelude::*;
-    use pgrx::{IntoDatum, Json};
+    use pgrx::{IntoDatum, Json, PgMemoryContexts, direct_pg_extern_function_call};
     use serde_json::json;
 
     #[pg_test]
@@ -330,6 +335,67 @@ mod tests {
         assert_eq!(len, Ok(Some(5)));
     }
 
+    #[pg_test]
+    fn borrow_test_toasted_flat_array() -> Result<(), pgrx::spi::Error> {
+        Spi::run("CREATE TEMP TABLE flat_array_toast_test (arr integer[])")?;
+        Spi::run("ALTER TABLE flat_array_toast_test ALTER COLUMN arr SET STORAGE EXTERNAL")?;
+        Spi::run(
+            "INSERT INTO flat_array_toast_test \
+             SELECT array_agg(i) FROM generate_series(1, 2500) i",
+        )?;
+
+        let result = Spi::get_two::<i32, i32>(
+            "SELECT borrow_get_arr_nelems(arr), borrow_sum_array(arr) \
+             FROM flat_array_toast_test",
+        );
+        assert_eq!(result, Ok((Some(2500), Some(3_126_250))));
+
+        let iterated = Spi::get_two::<i64, i64>(
+            "SELECT count(*), sum(value) \
+             FROM flat_array_toast_test, LATERAL borrow_iter_array(arr) value",
+        );
+        assert_eq!(iterated, Ok((Some(2500), Some(3_126_250))));
+
+        Spi::connect(|client| {
+            let table =
+                client.select("SELECT arr FROM flat_array_toast_test", Some(1), &[])?.first();
+            let datum = table.get_datum_by_ordinal(1)?.expect("array was null");
+            assert!(unsafe { pgrx::varlena::varatt_is_1b_e(datum.cast_mut_ptr()) });
+
+            unsafe {
+                PgMemoryContexts::Transient {
+                    parent: PgMemoryContexts::CurrentMemoryContext.value(),
+                    name: "toasted FlatArray cleanup test",
+                    min_context_size: 8 * 1024,
+                    initial_block_size: 8 * 1024,
+                    max_block_size: 8 * 1024,
+                }
+                .switch_to(|context| {
+                    let call = || {
+                        direct_pg_extern_function_call::<i32>(
+                            super::borrow_get_arr_nelems_wrapper,
+                            &[Some(datum)],
+                        )
+                    };
+
+                    assert_eq!(call(), Some(2500));
+                    let warmed = pg_sys::MemoryContextMemAllocated(context.value(), true);
+                    for _ in 0..64 {
+                        assert_eq!(call(), Some(2500));
+                    }
+                    let after = pg_sys::MemoryContextMemAllocated(context.value(), true);
+
+                    assert!(
+                        after <= warmed + 64 * 1024,
+                        "detoasted arguments accumulated {} bytes",
+                        after - warmed
+                    );
+                });
+            }
+            Ok(())
+        })
+    }
+
     #[pg_test]
     fn borrow_test_get_arr_data_ptr_nth_elem() {
         let nth =

pgrx/src/array/flat_array.rs

@@ -316,6 +316,7 @@ fn alloc_zeroed_head(
 
 unsafe impl<T: ?Sized> BorrowDatum for FlatArray<'_, T> {
     const PASS: layout::PassBy = layout::PassBy::Ref;
+
     unsafe fn point_from(ptr: ptr::NonNull<u8>) -> ptr::NonNull<Self> {
         unsafe {
             let len =
@@ -325,6 +326,22 @@ unsafe impl<T: ?Sized> BorrowDatum for FlatArray<'_, T> {
             )
         }
     }
+
+    unsafe fn borrow_arg_unchecked<'dat>(
+        ptr: ptr::NonNull<u8>,
+        register: impl FnOnce(ptr::NonNull<u8>),
+    ) -> &'dat Self {
+        // SAFETY: The caller guarantees `ptr` points to a valid PostgreSQL array Datum.
+        let detoasted = unsafe { pg_sys::pg_detoast_datum(ptr.as_ptr().cast()) };
+        let detoasted = ptr::NonNull::new(detoasted).expect("pg_detoast_datum returned null");
+        if detoasted.cast() != ptr {
+            register(detoasted.cast());
+        }
+
+        // SAFETY: `pg_detoast_datum` returns an aligned, contiguous, initialized ArrayType, and a
+        // fresh allocation is registered above to live until the function result has been boxed.
+        unsafe { Self::point_from(detoasted.cast()).as_ref() }
+    }
 }
 
 /// `T[]` in Postgres