AnyNumeric is no longer backed by Postgres-allocated memory

pgrx/src/datum/numeric.rs

@@ -14,7 +14,7 @@ use std::iter::Sum;
 
 use crate::numeric_support::convert::{from_primitive_helper, FromPrimitiveFunc};
 pub use crate::numeric_support::error::Error;
-use crate::{direct_function_call, pg_sys, varsize, PgMemoryContexts};
+use crate::{direct_function_call, pg_sys};
 
 /// A wrapper around the Postgres SQL `NUMERIC(P, S)` type.  Its `Precision` and `Scale` values
 /// are known at compile-time to assist with scale conversions and general type safety.
@@ -40,30 +40,12 @@ pub struct Numeric<const P: u32, const S: u32>(pub(crate) AnyNumeric);
 /// to represent any Rust primitive value from `i128::MIN` to `u128::MAX` and anything in between.
 ///
 /// Generally, this is the type you'll want to use as function arguments.
+///
+#[derive(Debug, Clone)]
 pub struct AnyNumeric {
-    pub(crate) inner: pg_sys::Numeric,
-    pub(crate) need_pfree: bool,
-}
-
-impl Clone for AnyNumeric {
-    /// Performs a deep clone of this [`AnyNumeric`] into the [`pg_sys::CurrentMemoryContext`].
-    fn clone(&self) -> Self {
-        unsafe {
-            let copy = PgMemoryContexts::CurrentMemoryContext
-                .copy_ptr_into(self.inner, varsize(self.inner.cast()));
-            AnyNumeric { inner: copy, need_pfree: true }
-        }
-    }
-}
-
-impl Drop for AnyNumeric {
-    fn drop(&mut self) {
-        if self.need_pfree {
-            unsafe {
-                pg_sys::pfree(self.inner.cast());
-            }
-        }
-    }
+    // we represent a NUMERIC as opaque bytes -- we never inspect these ourselves, only a pointer
+    // to them (cast as a [`Datum`]) are passed to Postgres
+    pub(super) inner: Box<[u8]>,
 }
 
 impl Display for AnyNumeric {
@@ -76,12 +58,6 @@ impl Display for AnyNumeric {
     }
 }
 
-impl Debug for AnyNumeric {
-    fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
-        write!(fmt, "AnyNumeric(inner: {self}, need_pfree:{})", self.need_pfree)
-    }
-}
-
 impl<const P: u32, const S: u32> Display for Numeric<P, S> {
     fn fmt(&self, fmt: &mut Formatter<'_>) -> Result<(), fmt::Error> {
         write!(fmt, "{}", self.0)
@@ -142,7 +118,7 @@ impl AnyNumeric {
 
     /// Is this [`AnyNumeric`] not-a-number?
     pub fn is_nan(&self) -> bool {
-        unsafe { pg_sys::numeric_is_nan(self.inner) }
+        unsafe { pg_sys::numeric_is_nan(self.as_ptr() as *mut _) }
     }
 
     /// The absolute value of this [`AnyNumeric`]
@@ -191,7 +167,7 @@ impl AnyNumeric {
     /// compare equal.
     pub fn normalize(&self) -> &str {
         unsafe {
-            let s = pg_sys::numeric_normalize(self.inner.cast());
+            let s = pg_sys::numeric_normalize(self.as_ptr() as *mut _);
             let cstr = CStr::from_ptr(s);
             let normalized = cstr.to_str().unwrap();
             normalized
@@ -217,12 +193,12 @@ impl AnyNumeric {
 
     #[inline]
     pub(crate) fn as_datum(&self) -> Option<pg_sys::Datum> {
-        Some(pg_sys::Datum::from(self.inner))
+        Some(pg_sys::Datum::from(self.inner.as_ptr()))
     }
 
     #[inline]
-    pub(crate) fn copy(&self) -> AnyNumeric {
-        AnyNumeric { inner: self.inner, need_pfree: false }
+    fn as_ptr(&self) -> *const pg_sys::NumericData {
+        self.inner.as_ptr().cast()
     }
 }
 

pgrx/src/datum/numeric_support/convert.rs

@@ -14,7 +14,8 @@ use pgrx_pg_sys::PgTryBuilder;
 use crate::numeric::make_typmod;
 use crate::numeric_support::error::Error;
 use crate::{
-    direct_function_call, direct_function_call_as_datum, pg_sys, AnyNumeric, IntoDatum, Numeric,
+    direct_function_call, direct_function_call_as_datum, pg_sys, AnyNumeric, FromDatum, IntoDatum,
+    Numeric,
 };
 
 pub use super::convert_anynumeric::*;
@@ -86,9 +87,15 @@ pub(crate) fn from_primitive_helper<I: IntoDatum, const P: u32, const S: u32>(
     };
 
     PgTryBuilder::new(|| {
-        let datum = materialize_numeric_datum();
-        // we asked Postgres to create this Numeric datum for us, so it'll need to be freed at some point
-        Ok(Numeric(AnyNumeric { inner: datum.cast_mut_ptr(), need_pfree: true }))
+        unsafe {
+            let datum = materialize_numeric_datum();
+            let anynumeric = AnyNumeric::from_datum(datum, false).unwrap();
+
+            // SAFETY:  We asked Postgres to create a new NUMERIC instance, so it now needs to be freed
+            // after we've copied it into Rust memory
+            pg_sys::pfree(datum.cast_mut_ptr());
+            Ok(Numeric(anynumeric))
+        }
     })
     .catch_when(PgSqlErrorCode::ERRCODE_INVALID_TEXT_REPRESENTATION, |e| {
         if let CaughtError::PostgresError(ref ereport) = e {

pgrx/src/datum/numeric_support/convert_numeric.rs

@@ -22,7 +22,7 @@ impl<const P: u32, const S: u32> TryFrom<AnyNumeric> for Numeric<P, S> {
 
     #[inline]
     fn try_from(value: AnyNumeric) -> Result<Self, Self::Error> {
-        from_primitive_helper::<_, P, S>(value.copy(), FromPrimitiveFunc::Numeric)
+        from_primitive_helper::<_, P, S>(value.clone(), FromPrimitiveFunc::Numeric)
     }
 }
 

pgrx/src/datum/numeric_support/datum.rs

@@ -7,7 +7,7 @@
 //LICENSE All rights reserved.
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
-use crate::{pg_sys, AnyNumeric, FromDatum, IntoDatum, Numeric, PgMemoryContexts};
+use crate::{pg_sys, varsize_any, AnyNumeric, FromDatum, IntoDatum, Numeric};
 
 impl FromDatum for AnyNumeric {
     #[inline]
@@ -22,42 +22,48 @@ impl FromDatum for AnyNumeric {
         if is_null {
             None
         } else {
-            let numeric = pg_sys::pg_detoast_datum(datum.cast_mut_ptr()) as pg_sys::Numeric;
+            // Going back to Postgres v9.1, `pg_sys::NumericData` is really a "varlena" in disguise.
+            //
+            // We want to copy it out of the Postgres-allocated memory it's in and into something
+            // managed by Rust.
+            //
+            // First we'll detoast it and then copy the backing varlena bytes into a `Vec<u8>`.
+            // This is what we'll use later if we need to convert back into a postgres-allocated Datum
+            // or to provide a view over the bytes as a Datum.  The latter is what most of the AnyNumeric
+            // support functions use, via the `as_datum()` function.
+
+            // detoast
+            let numeric = pg_sys::pg_detoast_datum(datum.cast_mut_ptr());
             let is_copy = !std::ptr::eq(
-                numeric.cast::<pg_sys::NumericData>(),
-                datum.cast_mut_ptr::<pg_sys::NumericData>(),
+                numeric.cast::<pg_sys::varlena>(),
+                datum.cast_mut_ptr::<pg_sys::varlena>(),
             );
-            Some(AnyNumeric { inner: numeric, need_pfree: is_copy })
-        }
-    }
 
-    unsafe fn from_datum_in_memory_context(
-        mut memory_context: PgMemoryContexts,
-        datum: pg_sys::Datum,
-        is_null: bool,
-        _typoid: pg_sys::Oid,
-    ) -> Option<Self>
-    where
-        Self: Sized,
-    {
-        if is_null {
-            None
-        } else {
-            memory_context.switch_to(|_| {
-                // copy the Datum into this MemoryContext and then create the AnyNumeric over that
-                let copy = pg_sys::pg_detoast_datum_copy(datum.cast_mut_ptr());
-                Some(AnyNumeric { inner: copy.cast(), need_pfree: true })
-            })
+            // copy us into a rust-owned/allocated Box<[u8]>
+            let size = varsize_any(numeric);
+            let slice = std::slice::from_raw_parts(numeric.cast::<u8>(), size);
+            let boxed: Box<[u8]> = slice.into();
+
+            // free the copy detoast might have made
+            if is_copy {
+                pg_sys::pfree(numeric.cast());
+            }
+
+            Some(AnyNumeric { inner: boxed })
         }
     }
 }
 
 impl IntoDatum for AnyNumeric {
     #[inline]
-    fn into_datum(mut self) -> Option<pg_sys::Datum> {
-        // we're giving it to Postgres so we don't want our drop impl to free the inner pointer
-        self.need_pfree = false;
-        Some(pg_sys::Datum::from(self.inner))
+    fn into_datum(self) -> Option<pg_sys::Datum> {
+        unsafe {
+            let size = self.inner.len();
+            let src = self.inner.as_ptr();
+            let dest = pg_sys::palloc(size).cast();
+            std::ptr::copy_nonoverlapping(src, dest, size);
+            Some(pg_sys::Datum::from(dest))
+        }
     }
 
     #[inline]

pgrx/src/datum/numeric_support/mod.rs

@@ -7,7 +7,7 @@
 //LICENSE All rights reserved.
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
-use crate::{direct_function_call_as_datum, pg_sys, AnyNumeric};
+use crate::{direct_function_call_as_datum, pg_sys, AnyNumeric, FromDatum};
 
 pub mod cmp;
 pub mod convert;
@@ -29,7 +29,11 @@ pub(super) fn call_numeric_func(
     unsafe {
         // SAFETY: this call to direct_function_call_as_datum will never return None
         let numeric_datum = direct_function_call_as_datum(func, args).unwrap_unchecked();
-        // we asked Postgres to create this numeric so it'll need to be freed eventually
-        AnyNumeric { inner: numeric_datum.cast_mut_ptr(), need_pfree: true }
+
+        // we asked Postgres to create this numeric so it'll need to be freed after we copy it
+        let anynumeric = AnyNumeric::from_datum(numeric_datum, false);
+        pg_sys::pfree(numeric_datum.cast_mut_ptr());
+
+        anynumeric.unwrap()
     }
 }

pgrx/src/datum/numeric_support/ops.rs

@@ -123,7 +123,7 @@ macro_rules! anynumeric_assign_op_from_primitive {
         impl $opname<$ty> for AnyNumeric {
             #[inline]
             fn $trait_fname(&mut self, rhs: $ty) {
-                *self = self.copy() $op AnyNumeric::from(rhs);
+                *self = self.clone() $op AnyNumeric::from(rhs);
             }
         }
     }
@@ -202,14 +202,14 @@ macro_rules! anynumeric_assign_op_from_float {
                 // these versions of Postgres could produce an error when unwrapping a try_from(float)
                 #[cfg(any(feature = "pg11", feature = "pg12", feature = "pg13"))]
                 {
-                    *self = self.copy() $op AnyNumeric::try_from(rhs).unwrap();
+                    *self = self.clone() $op AnyNumeric::try_from(rhs).unwrap();
                 }
 
                 // these versions won't, so we use .unwrap_unchecked()
                 #[cfg(any(feature = "pg14", feature = "pg15", feature = "pg16"))]
                 {
                     unsafe {
-                        *self = self.copy() $op AnyNumeric::try_from(rhs).unwrap_unchecked();
+                        *self = self.clone() $op AnyNumeric::try_from(rhs).unwrap_unchecked();
                     }
                 }
             }

pgrx/src/varlena.rs

@@ -295,14 +295,6 @@ pub unsafe fn vardata_any(ptr: *const pg_sys::varlena) -> *const std::os::raw::c
     }
 }
 
-/// ## Safety
-///
-/// This function is unsafe because it blindly dereferences the varlena pointer argument
-#[inline]
-pub unsafe fn varlena_size(t: *const pg_sys::varlena) -> usize {
-    std::mem::size_of_val(&(*t).vl_len_) + varsize_any_exhdr(t)
-}
-
 /// Convert a Postgres `varlena *` (or `text *`) into a Rust `&str`.
 ///
 /// ## Safety