fix: detoast borrowed FlatArray arguments

eeeebbbbrrrr · eeeebbbbrrrr · commit e60a8e469214 · 2026-07-13T18:26:57.000Z
diff --git a/pgrx-sql-entity-graph/src/pg_extern/mod.rs b/pgrx-sql-entity-graph/src/pg_extern/mod.rs
@@ -469,7 +469,7 @@ impl PgExtern {
                             let result = match call_flow {
                                 ::pgrx::callconv::CallCx::WrappedFn(mcx) => {
                                     let mut mcx = ::pgrx::PgMemoryContexts::For(mcx);
-                                    let #args_ident = &mut fcinfo.args();
+                                    let mut #args_ident = fcinfo.args();
                                     let call_result = mcx.switch_to(|_| {
                                         #(#arg_fetches)*
                                         #func_name( #(#arg_pats),* )
diff --git a/pgrx-unit-tests/src/tests/array_borrowed.rs b/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 =
diff --git a/pgrx/src/array/flat_array.rs b/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
diff --git a/pgrx/src/callconv.rs b/pgrx/src/callconv.rs
@@ -27,23 +27,10 @@ use crate::rel::PgRelation;
 use crate::{PgBox, PgMemoryContexts};
 
 use core::marker::PhantomData;
-use core::{iter, mem, ptr, slice};
+use core::{mem, ptr};
 use std::ffi::{CStr, CString};
 use std::ptr::NonNull;
 
-struct ReadOnly<T>(T);
-
-impl<T> ReadOnly<T> {
-    unsafe fn refer_to(&self) -> &T {
-        &self.0
-    }
-}
-
-unsafe impl<T> Sync for ReadOnly<T> {}
-
-static VIRTUAL_ARGUMENT: ReadOnly<pg_sys::NullableDatum> =
-    ReadOnly(pg_sys::NullableDatum { value: pg_sys::Datum::null(), isnull: true });
-
 /// How to pass a value from Postgres to Rust
 ///
 /// This bound is necessary to distinguish things which can be passed into a `#[pg_extern] fn`.
@@ -288,7 +275,7 @@ where
         };
         unsafe {
             let ptr = ptr::NonNull::new_unchecked(ptr);
-            T::borrow_unchecked(ptr)
+            T::borrow_arg_unchecked(ptr, |ptr| arg.0.register_detoasted_arg(ptr))
         }
     }
 
@@ -299,7 +286,11 @@ where
         });
         match (arg.is_null(), ptr) {
             (true, _) | (false, None) => Nullable::Null,
-            (false, Some(ptr)) => unsafe { Nullable::Valid(T::borrow_unchecked(ptr)) },
+            (false, Some(ptr)) => unsafe {
+                Nullable::Valid(T::borrow_arg_unchecked(ptr, |ptr| {
+                    arg.0.register_detoasted_arg(ptr)
+                }))
+            },
         }
     }
 }
@@ -599,8 +590,26 @@ where
 
 type FcInfoData = pg_sys::FunctionCallInfoBaseData;
 
-#[derive(Clone)]
-pub struct FcInfo<'fcx>(pgrx_pg_sys::FunctionCallInfo, PhantomData<&'fcx mut FcInfoData>);
+/// A uniquely owned PostgreSQL function call frame.
+///
+/// Argument decoding may register temporary detoast allocations here, so this type is intentionally
+/// not cloneable.
+pub struct FcInfo<'fcx>(
+    pgrx_pg_sys::FunctionCallInfo,
+    PhantomData<&'fcx mut FcInfoData>,
+    Vec<NonNull<u8>>,
+    pg_sys::MemoryContext,
+);
+
+impl Drop for FcInfo<'_> {
+    fn drop(&mut self) {
+        for allocation in self.2.drain(..) {
+            // SAFETY: `BorrowDatum::borrow_arg_unchecked` requires registered pointers to be
+            // PostgreSQL allocations that can be released after the function result is boxed.
+            unsafe { pg_sys::pfree(allocation.as_ptr().cast()) }
+        }
+    }
+}
 
 // when talking about this, there's the lifetime for setreturningfunction, and then there's the current context's lifetime.
 // Potentially <'srf, 'curr, 'ret: 'curr + 'srf> -> <'ret> but don't start with that.
@@ -620,7 +629,19 @@ impl<'fcx> FcInfo<'fcx> {
     #[inline]
     pub unsafe fn from_ptr(fcinfo: pg_sys::FunctionCallInfo) -> FcInfo<'fcx> {
         let _nullptr_check = NonNull::new(fcinfo).expect("fcinfo pointer must be non-null");
-        Self(fcinfo, PhantomData)
+        Self(fcinfo, PhantomData, Vec::new(), unsafe { pg_sys::CurrentMemoryContext })
+    }
+
+    fn register_detoasted_arg(&mut self, ptr: NonNull<u8>) {
+        // Set-returning functions unbox their arguments in a multi-call context so an iterator
+        // may retain borrows across calls. That context owns its detoast allocations and releases
+        // them when the SRF finishes. Ordinary functions unbox in the entry context, where we can
+        // release fresh detoast copies as soon as their result has been boxed.
+        if unsafe { pg_sys::CurrentMemoryContext } != self.3 {
+            return;
+        }
+
+        self.2.push(ptr);
     }
     /// Retrieve the arguments to this function call as a slice of [`pgrx_pg_sys::NullableDatum`]
     #[inline]
@@ -830,14 +851,18 @@ impl<'fcx> FcInfo<'fcx> {
         }
     }
 
+    /// Create a decoder for this function call's arguments.
+    ///
+    /// The decoder borrows the call frame exclusively so temporary argument allocations have one
+    /// unambiguous owner.
     #[inline]
-    pub fn args<'arg>(&'arg self) -> Args<'arg, 'fcx> {
-        Args { iter: self.raw_args().iter().enumerate(), fcinfo: self }
+    pub fn args<'arg>(&'arg mut self) -> Args<'arg, 'fcx> {
+        Args { next: 0, fcinfo: self }
     }
 }
 
 // TODO: rebadge this as AnyElement
-pub struct Arg<'a, 'fcx>(&'a FcInfo<'fcx>, usize, &'a pg_sys::NullableDatum);
+pub struct Arg<'a, 'fcx>(&'a mut FcInfo<'fcx>, usize, pg_sys::NullableDatum);
 
 impl<'a, 'fcx> Arg<'a, 'fcx> {
     /// # Performance note
@@ -880,23 +905,30 @@ impl<'a, 'fcx> Arg<'a, 'fcx> {
     }
 }
 
+/// A lending-style decoder for PostgreSQL function arguments.
+///
+/// This is deliberately not an [`Iterator`]: each decoded [`Arg`] temporarily borrows the call
+/// frame's detoast cleanup state.
 pub struct Args<'a, 'fcx> {
-    iter: iter::Enumerate<slice::Iter<'a, pg_sys::NullableDatum>>,
-    fcinfo: &'a FcInfo<'fcx>,
-}
-
-impl<'a, 'fcx> Iterator for Args<'a, 'fcx> {
-    type Item = Arg<'a, 'fcx>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.iter.next().map(|(a, b)| Arg(self.fcinfo, a, b))
-    }
+    next: usize,
+    fcinfo: &'a mut FcInfo<'fcx>,
 }
 
 impl<'a, 'fcx> Args<'a, 'fcx> {
     // generate an argument for use by virtual args
-    fn synthesize_virtual_arg(&self) -> Arg<'a, 'fcx> {
-        Arg(self.fcinfo, usize::MAX, unsafe { VIRTUAL_ARGUMENT.refer_to() })
+    fn synthesize_virtual_arg(&mut self) -> Arg<'_, 'fcx> {
+        Arg(
+            self.fcinfo,
+            usize::MAX,
+            pg_sys::NullableDatum { value: pg_sys::Datum::null(), isnull: true },
+        )
+    }
+
+    fn next_raw(&mut self) -> Option<Arg<'_, 'fcx>> {
+        let index = self.next;
+        let datum = self.fcinfo.raw_args().get(index).copied()?;
+        self.next += 1;
+        Some(Arg(self.fcinfo, index, datum))
     }
 
     /// # Safety
@@ -910,7 +942,7 @@ impl<'a, 'fcx> Args<'a, 'fcx> {
             unsafe { Some(T::unbox_arg_unchecked(self.synthesize_virtual_arg())) }
         } else {
             // SAFETY: caller upholds
-            unsafe { self.next().map(|next| T::unbox_arg_unchecked(next)) }
+            unsafe { self.next_raw().map(|next| T::unbox_arg_unchecked(next)) }
         }
     }
 
@@ -923,7 +955,7 @@ impl<'a, 'fcx> Args<'a, 'fcx> {
             unsafe { Some(T::unbox_nullable_arg(self.synthesize_virtual_arg())) }
         } else {
             // SAFETY: caller upholds
-            unsafe { self.next().map(|next| T::unbox_nullable_arg(next)) }
+            unsafe { self.next_raw().map(|next| T::unbox_nullable_arg(next)) }
         }
     }
 }
diff --git a/pgrx/src/datum/borrow.rs b/pgrx/src/datum/borrow.rs
@@ -63,6 +63,20 @@ pub unsafe trait BorrowDatum {
     unsafe fn borrow_unchecked<'dat>(ptr: ptr::NonNull<u8>) -> &'dat Self {
         unsafe { BorrowDatum::point_from(ptr).as_ref() }
     }
+
+    /// Prepare and borrow a function argument, registering any temporary allocation for cleanup.
+    ///
+    /// This is used by pgrx's function-call machinery. Implementors that must copy an argument
+    /// before it can be borrowed should pass that allocation to `register` before returning the
+    /// borrow. The allocation must have been made by PostgreSQL and ownership is transferred to
+    /// `register`, which keeps it alive for the function's return ABI and releases it afterward.
+    #[doc(hidden)]
+    unsafe fn borrow_arg_unchecked<'dat>(
+        ptr: ptr::NonNull<u8>,
+        _register: impl FnOnce(ptr::NonNull<u8>),
+    ) -> &'dat Self {
+        unsafe { BorrowDatum::borrow_unchecked(ptr) }
+    }
 }
 
 /// From a pointer to a Datum, obtain a pointer to T's bytes
