diff --git a/pgrx-sql-entity-graph/src/pg_extern/mod.rs b/pgrx-sql-entity-graph/src/pg_extern/mod.rs index 822846c2f..5c0ce77c0 100644 --- a/pgrx-sql-entity-graph/src/pg_extern/mod.rs +++ b/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),* ) diff --git a/pgrx-unit-tests/src/tests/array_borrowed.rs b/pgrx-unit-tests/src/tests/array_borrowed.rs index d8dbee83e..05586918b 100644 --- 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 { 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::( + "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::( + "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::( + 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 index 781a9e260..cd7782eb1 100644 --- a/pgrx/src/array/flat_array.rs +++ b/pgrx/src/array/flat_array.rs @@ -316,6 +316,7 @@ fn alloc_zeroed_head( unsafe impl BorrowDatum for FlatArray<'_, T> { const PASS: layout::PassBy = layout::PassBy::Ref; + unsafe fn point_from(ptr: ptr::NonNull) -> ptr::NonNull { unsafe { let len = @@ -325,6 +326,22 @@ unsafe impl BorrowDatum for FlatArray<'_, T> { ) } } + + unsafe fn borrow_arg_unchecked<'dat>( + ptr: ptr::NonNull, + register: impl FnOnce(ptr::NonNull), + ) -> &'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 index df19a86a8..4b0416a18 100644 --- 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); - -impl ReadOnly { - unsafe fn refer_to(&self) -> &T { - &self.0 - } -} - -unsafe impl Sync for ReadOnly {} - -static VIRTUAL_ARGUMENT: ReadOnly = - 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) + })) + }, } } } @@ -332,7 +323,7 @@ pub unsafe trait RetAbi: Sized { } fn check_and_prepare(fcinfo: &mut FcInfo<'_>) -> CallCx { - unsafe { Self::check_fcinfo_and_prepare(fcinfo.0) } + unsafe { Self::check_fcinfo_and_prepare(fcinfo.ptr) } } /// answer what kind and how many returns happen from this type @@ -340,7 +331,7 @@ pub unsafe trait RetAbi: Sized { // move into the function context and obtain a Datum unsafe fn box_ret_in<'fcx>(fcinfo: &mut FcInfo<'fcx>, ret: Self::Ret) -> Datum<'fcx> { - let fcinfo = fcinfo.0; + let fcinfo = fcinfo.ptr; unsafe { mem::transmute(Self::box_ret_in_fcinfo(fcinfo, ret)) } } @@ -369,7 +360,7 @@ pub unsafe trait RetAbi: Sized { } fn ret_from_fcx(fcinfo: &mut FcInfo<'_>) -> Self::Ret { - let fcinfo = fcinfo.0; + let fcinfo = fcinfo.ptr; unsafe { Self::ret_from_fcinfo_fcx(fcinfo) } } @@ -599,8 +590,27 @@ 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> { + ptr: pgrx_pg_sys::FunctionCallInfo, + detoasted_arg_allocations: Vec>, + entry_memory_context: pg_sys::MemoryContext, + argument_memory_context: pg_sys::MemoryContext, + _marker: PhantomData<&'fcx mut FcInfoData>, +} + +impl Drop for FcInfo<'_> { + fn drop(&mut self) { + for allocation in self.detoasted_arg_allocations.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. @@ -615,12 +625,40 @@ impl<'fcx> FcInfo<'fcx> { /// /// # Safety /// - /// This function is unsafe as we cannot ensure the `fcinfo` argument is a valid - /// [`pg_sys::FunctionCallInfo`] pointer. This is your responsibility. + /// - `fcinfo` must be a valid [`pg_sys::FunctionCallInfo`] pointer and must not be wrapped by + /// another live `FcInfo` for `'fcx`. + /// - PostgreSQL must be initialized, and its current memory context must not be reset or deleted + /// while the function call is active. #[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) + let entry_memory_context = NonNull::new(unsafe { pg_sys::CurrentMemoryContext }) + .expect("CurrentMemoryContext must be non-null") + .as_ptr(); + Self { + ptr: fcinfo, + detoasted_arg_allocations: Vec::new(), + entry_memory_context, + argument_memory_context: entry_memory_context, + _marker: PhantomData, + } + } + + fn register_detoasted_arg(&mut self, ptr: NonNull) { + // 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. + assert_eq!( + unsafe { pg_sys::CurrentMemoryContext }, + self.argument_memory_context, + "detoasted argument was allocated outside its selected memory context" + ); + if self.argument_memory_context != self.entry_memory_context { + return; + } + + self.detoasted_arg_allocations.push(ptr); } /// Retrieve the arguments to this function call as a slice of [`pgrx_pg_sys::NullableDatum`] #[inline] @@ -628,8 +666,8 @@ impl<'fcx> FcInfo<'fcx> { // Null pointer check already performed on immutable pointer // at construction time. unsafe { - let arg_len = (*self.0).nargs; - let args_ptr: *const pg_sys::NullableDatum = ptr::addr_of!((*self.0).args).cast(); + let arg_len = (*self.ptr).nargs; + let args_ptr: *const pg_sys::NullableDatum = ptr::addr_of!((*self.ptr).args).cast(); // A valid FcInfoWrapper constructed from a valid FuntionCallInfo should always have // at least nargs elements of NullableDatum. std::slice::from_raw_parts(args_ptr, arg_len as _) @@ -641,7 +679,7 @@ impl<'fcx> FcInfo<'fcx> { /// that type is sufficient. #[inline] pub unsafe fn as_mut_ptr(&self) -> pg_sys::FunctionCallInfo { - self.0 + self.ptr } /// Accessor for the "is null" flag @@ -650,7 +688,7 @@ impl<'fcx> FcInfo<'fcx> { /// If this flag is set to "false", then the resulting return must be a valid [`Datum`] for /// the function call's result type. pub unsafe fn set_return_is_null(&mut self) -> &mut bool { - unsafe { &mut (*self.0).isnull } + unsafe { &mut (*self.ptr).isnull } } /// Modifies the function call's return to be null @@ -721,7 +759,7 @@ impl<'fcx> FcInfo<'fcx> { #[inline] pub fn get_collation(&self) -> Option { // SAFETY: see FcInfo::from_ptr - let fcinfo = unsafe { self.0.as_mut() }.unwrap(); + let fcinfo = unsafe { self.ptr.as_mut() }.unwrap(); (fcinfo.fncollation.to_u32() != 0).then_some(fcinfo.fncollation) } @@ -732,11 +770,11 @@ impl<'fcx> FcInfo<'fcx> { // SAFETY: see FcInfo::from_ptr unsafe { // bool::then() is lazy-evaluated, then_some is not. - (num < ((*self.0).nargs as usize)).then( + (num < ((*self.ptr).nargs as usize)).then( #[inline] || { pg_sys::get_fn_expr_argtype( - self.0.as_ref().unwrap().flinfo, + self.ptr.as_ref().unwrap().flinfo, num as std::os::raw::c_int, ) }, @@ -754,7 +792,7 @@ impl<'fcx> FcInfo<'fcx> { // to construct a FcInfo. If that constraint is maintained, this should // be safe. unsafe { - let mut flinfo = NonNull::new((*self.0).flinfo).unwrap(); + let mut flinfo = NonNull::new((*self.ptr).flinfo).unwrap(); if flinfo.as_ref().fn_extra.is_null() { flinfo.as_mut().fn_extra = PgMemoryContexts::For(flinfo.as_ref().fn_mcxt) .leak_and_drop_on_delete(default()) @@ -771,17 +809,17 @@ impl<'fcx> FcInfo<'fcx> { // Safety: User must supply a valid fcinfo to from_ptr() in order // to construct a FcInfo. If that constraint is maintained, this should // be safe. - unsafe { !(*(*self.0).flinfo).fn_extra.is_null() } + unsafe { !(*(*self.ptr).flinfo).fn_extra.is_null() } } /// Thin wrapper around [`pg_sys::init_MultiFuncCall`], made necessary /// because this structure's FunctionCallInfo is a private field. /// - /// This should initialize `self.0.flinfo.fn_extra` + /// This should initialize `self.ptr.flinfo.fn_extra` #[inline] pub unsafe fn init_multi_func_call(&mut self) -> &'fcx mut pg_sys::FuncCallContext { unsafe { - let fcx: *mut pg_sys::FuncCallContext = pg_sys::init_MultiFuncCall(self.0); + let fcx: *mut pg_sys::FuncCallContext = pg_sys::init_MultiFuncCall(self.ptr); debug_assert!(!fcx.is_null()); &mut *fcx } @@ -790,18 +828,18 @@ impl<'fcx> FcInfo<'fcx> { /// Equivalent to "per_MultiFuncCall" with no FFI cost, and a lifetime /// constraint. /// - /// Safety: Assumes `self.0.flinfo.fn_extra` is non-null + /// Safety: Assumes `self.ptr.flinfo.fn_extra` is non-null /// i.e. [`FcInfo::srf_is_initialized()`] would be `true`. #[inline] pub(crate) unsafe fn deref_fcx(&mut self) -> &'fcx mut pg_sys::FuncCallContext { unsafe { - let fcx: *mut pg_sys::FuncCallContext = (*(*self.0).flinfo).fn_extra.cast(); + let fcx: *mut pg_sys::FuncCallContext = (*(*self.ptr).flinfo).fn_extra.cast(); debug_assert!(!fcx.is_null()); &mut *fcx } } - /// Safety: Assumes `self.0.flinfo.fn_extra` is non-null + /// Safety: Assumes `self.ptr.flinfo.fn_extra` is non-null /// i.e. [`FcInfo::srf_is_initialized()`] would be `true`. #[inline] pub unsafe fn srf_return_next(&mut self) { @@ -811,12 +849,12 @@ impl<'fcx> FcInfo<'fcx> { } } - /// Safety: Assumes `self.0.flinfo.fn_extra` is non-null + /// Safety: Assumes `self.ptr.flinfo.fn_extra` is non-null /// i.e. [`FcInfo::srf_is_initialized()`] would be `true`. #[inline] pub unsafe fn srf_return_done(&mut self) { unsafe { - pg_sys::end_MultiFuncCall(self.0, self.deref_fcx()); + pg_sys::end_MultiFuncCall(self.ptr, self.deref_fcx()); self.get_result_info().set_is_done(pg_sys::ExprDoneCond::ExprEndResult); } } @@ -826,18 +864,42 @@ impl<'fcx> FcInfo<'fcx> { #[inline] pub unsafe fn get_result_info(&self) -> ReturnSetInfoWrapper<'fcx> { unsafe { - ReturnSetInfoWrapper::from_ptr((*self.0).resultinfo as *mut pg_sys::ReturnSetInfo) + ReturnSetInfoWrapper::from_ptr((*self.ptr).resultinfo as *mut pg_sys::ReturnSetInfo) } } + /// 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 mut self) -> Args<'arg, 'fcx> { + let entry_memory_context = self.entry_memory_context; + unsafe { self.args_in(entry_memory_context) } + } + + /// Create a decoder whose temporary argument allocations belong to `memory_context`. + /// + /// # Safety + /// - `memory_context` must be valid and must not be reset or deleted while any borrow returned + /// by the decoder is live. + /// - Arguments must be decoded while `memory_context` is PostgreSQL's current memory context. + /// - If this is not the call's entry memory context, it must own and release temporary argument + /// allocations after all returned borrows expire. + #[doc(hidden)] #[inline] - pub fn args<'arg>(&'arg self) -> Args<'arg, 'fcx> { - Args { iter: self.raw_args().iter().enumerate(), fcinfo: self } + pub unsafe fn args_in<'arg>( + &'arg mut self, + memory_context: pg_sys::MemoryContext, + ) -> Args<'arg, 'fcx> { + let _nullptr_check = NonNull::new(memory_context).expect("memory context must be non-null"); + self.argument_memory_context = memory_context; + 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 +942,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>, - fcinfo: &'a FcInfo<'fcx>, -} - -impl<'a, 'fcx> Iterator for Args<'a, 'fcx> { - type Item = Arg<'a, 'fcx>; - - fn next(&mut self) -> Option { - 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> { + 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 +979,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 +992,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 index ff475b1a9..9ce696b17 100644 --- a/pgrx/src/datum/borrow.rs +++ b/pgrx/src/datum/borrow.rs @@ -63,6 +63,31 @@ pub unsafe trait BorrowDatum { unsafe fn borrow_unchecked<'dat>(ptr: ptr::NonNull) -> &'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. + /// + /// # Safety + /// - `ptr` must point to an initialized function argument representation that is valid for + /// `Self` and [`BorrowDatum::PASS`]. + /// - The pointer ultimately passed to [`BorrowDatum::point_from`] must satisfy that method's + /// safety requirements for `Self`. + /// - The returned reference must remain valid for `'dat`. If it refers to a temporary copy, + /// the implementation must pass that copy to `register` before returning. + /// - An allocation passed to `register` must be a fresh, live PostgreSQL allocation that may be + /// released with `pg_sys::pfree`. Passing it transfers ownership, so the implementation must + /// not free it or register the caller-owned `ptr`. + #[doc(hidden)] + unsafe fn borrow_arg_unchecked<'dat>( + ptr: ptr::NonNull, + _register: impl FnOnce(ptr::NonNull), + ) -> &'dat Self { + unsafe { BorrowDatum::borrow_unchecked(ptr) } + } } /// From a pointer to a Datum, obtain a pointer to T's bytes