From f840dfffd3508c63317d51e242735003717230c6 Mon Sep 17 00:00:00 2001 From: danielshih Date: Mon, 15 Jun 2026 05:50:06 +0000 Subject: [PATCH 1/3] feat(memcx): safe slice + varlena allocation parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Layered API for sized + aligned allocation. Uses MemoryContextAllocAligned on PG >= 16, falls back to MemoryContextAllocExtended on older versions. feat(pbox): add PBox<[MaybeUninit]>::new_uninit_slice_in Lifetime-bound uninitialized slice allocation, mirroring alloc::boxed::Box::new_uninit_slice. Also adds as_ptr/as_mut_ptr on PBox and len/is_empty on PBox<[T]>. feat(pbox): add new_zeroed_slice_in and assume_init Adds zero-initialized uninit slice allocation and a safe transmute from PBox<[MaybeUninit]> to PBox<[T]>. The new test is intentionally incomplete pending Task 4's Deref<[T]> impl — confirmed by leaving a single, expected E0599 error on .iter() resolution. feat(pbox): Deref/DerefMut for PBox<[T]> Specific impls for slice elements; coexist with the BorrowDatum-bounded blanket impl since [T] does not impl BorrowDatum. Closes the test gap left by Task 3 — new_zeroed_slice_in_is_zero now compiles and passes. feat(pbox): add PBox::<[T]>::from_slice_in for Copy types Fast-path constructor: allocate uninit, copy_nonoverlapping, assume_init. feat(pbox): add from_iter_in with RAII panic-safety guard Iterator-driven slice constructor. Requires ExactSizeIterator so the allocation is sized upfront. On panic during iteration, an RAII Guard drops only the initialized prefix exactly once; underlying memory stays in the MemCx as designed (PBox has no Drop). Iterator under-yield is treated as an OOM-shaped error. feat(examples): add pbox_slice example crate Demonstrates MemCx::alloc_varlena returning PBox directly to Postgres as bytea, with no intermediate Vec. Includes pg_test cases for length, byte-pattern, and zero-length boundaries. feat(memcx): add alloc_huge_layout + guard alloc_layout at 1 GiB Two-part change: 1. New MemCx::alloc_huge_layout and alloc_huge_layout_zeroed methods that pass MCXT_ALLOC_HUGE alongside NO_OOM. These permit allocations up to MaxAllocHugeSize (~half the address space) for genuinely huge Rust-side scratch buffers. Mirrored at the PBox layer via new_huge_uninit_slice_in / new_huge_zeroed_slice_in. 2. alloc_layout / alloc_layout_zeroed now reject sizes above MaxAllocSize (1 GiB - 1) up front with OutOfMemory, instead of forwarding to palloc and triggering ereport/longjmp out of the Rust frame. This closes a latent unsoundness where the documented "panic-free" path could in fact longjmp on > 1 GiB requests. Varlena values are deliberately not given a "_huge_" variant: the varlena 4-byte header encodes a 30-bit size, so > 1 GiB varlenas are invalid by Postgres design regardless of how they were allocated. --- Cargo.lock | 8 + pgrx-examples/pbox_slice/.gitignore | 7 + pgrx-examples/pbox_slice/Cargo.toml | 35 +++ pgrx-examples/pbox_slice/README.md | 25 ++ pgrx-examples/pbox_slice/pbox_slice.control | 5 + pgrx-examples/pbox_slice/src/lib.rs | 168 ++++++++++++ pgrx-unit-tests/src/tests/mod.rs | 2 + pgrx-unit-tests/src/tests/pbox_slice_tests.rs | 253 ++++++++++++++++++ .../src/tests/varlena_buf_tests.rs | 60 +++++ pgrx/src/datum/mod.rs | 2 + pgrx/src/datum/varlena_buf.rs | 196 ++++++++++++++ pgrx/src/memcx.rs | 120 +++++++++ pgrx/src/palloc/pbox.rs | 246 ++++++++++++++++- pgrx/src/varlena.rs | 15 +- 14 files changed, 1125 insertions(+), 17 deletions(-) create mode 100644 pgrx-examples/pbox_slice/.gitignore create mode 100644 pgrx-examples/pbox_slice/Cargo.toml create mode 100644 pgrx-examples/pbox_slice/README.md create mode 100644 pgrx-examples/pbox_slice/pbox_slice.control create mode 100644 pgrx-examples/pbox_slice/src/lib.rs create mode 100644 pgrx-unit-tests/src/tests/pbox_slice_tests.rs create mode 100644 pgrx-unit-tests/src/tests/varlena_buf_tests.rs create mode 100644 pgrx/src/datum/varlena_buf.rs diff --git a/Cargo.lock b/Cargo.lock index ceda6800e..4493ba502 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2149,6 +2149,14 @@ dependencies = [ "hmac", ] +[[package]] +name = "pbox_slice" +version = "0.0.0" +dependencies = [ + "pgrx", + "pgrx-tests", +] + [[package]] name = "pem-rfc7468" version = "1.0.0" diff --git a/pgrx-examples/pbox_slice/.gitignore b/pgrx-examples/pbox_slice/.gitignore new file mode 100644 index 000000000..d4de4f351 --- /dev/null +++ b/pgrx-examples/pbox_slice/.gitignore @@ -0,0 +1,7 @@ +.DS_Store +.idea/ +/target +*.iml +**/*.rs.bk +Cargo.lock +sql/pbox_slice-*.sql diff --git a/pgrx-examples/pbox_slice/Cargo.toml b/pgrx-examples/pbox_slice/Cargo.toml new file mode 100644 index 000000000..109edf354 --- /dev/null +++ b/pgrx-examples/pbox_slice/Cargo.toml @@ -0,0 +1,35 @@ +#LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +#LICENSE +#LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +#LICENSE +#LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +#LICENSE +#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. + +[package] +name = "pbox_slice" +version = "0.0.0" +edition.workspace = true +rust-version.workspace = true +publish = false + +[lib] +crate-type = ["cdylib"] + +[features] +default = ["pg13"] +pg13 = ["pgrx/pg13", "pgrx-tests/pg13"] +pg14 = ["pgrx/pg14", "pgrx-tests/pg14"] +pg15 = ["pgrx/pg15", "pgrx-tests/pg15"] +pg16 = ["pgrx/pg16", "pgrx-tests/pg16"] +pg17 = ["pgrx/pg17", "pgrx-tests/pg17"] +pg18 = ["pgrx/pg18", "pgrx-tests/pg18"] +pg_test = [] + +[dependencies] +pgrx = { path = "../../pgrx" } + +[dev-dependencies] +pgrx-tests = { path = "../../pgrx-tests" } diff --git a/pgrx-examples/pbox_slice/README.md b/pgrx-examples/pbox_slice/README.md new file mode 100644 index 000000000..1a6c89acf --- /dev/null +++ b/pgrx-examples/pbox_slice/README.md @@ -0,0 +1,25 @@ +# pbox_slice — `MemCx` slice + varlena demos + +Demonstrates the safe, lifetime-bound allocation APIs added in support of +[issue #2217](https://github.com/pgcentralfoundation/pgrx/issues/2217). + +All four `#[pg_extern]` functions return a `PBox` directly to +Postgres — no intermediate `Vec` allocated or copied. + +| Function | Demonstrates | +|----------|--------------| +| `build_bytea(int)` | Bare-bones varlena allocation + payload write | +| `xor_bytea(bytea, bytea)` | Common bytea-in / bytea-out transformation | +| `pack_i32_array(int[])` | Exact-size varlena alloc + structured binary serialization | +| `hash_chain(bytea, int)` | Pairing `PBox<[T]>` scratch buffer with `VarlenaBuf` output | + +The slice-returning analogue (`PBox<[T]>` as a SQL array) is deferred to a +follow-up; it requires `BoxRet` + `SqlTranslatable` impls that bridge `[T]` +to Postgres's `ArrayType` representation. + +## Build & test + +```bash +cd pgrx-examples/pbox_slice +cargo pgrx test pg18 +``` diff --git a/pgrx-examples/pbox_slice/pbox_slice.control b/pgrx-examples/pbox_slice/pbox_slice.control new file mode 100644 index 000000000..cb11d1aab --- /dev/null +++ b/pgrx-examples/pbox_slice/pbox_slice.control @@ -0,0 +1,5 @@ +comment = 'pbox_slice: demonstrates MemCx slice + varlena APIs' +default_version = '@CARGO_VERSION@' +module_pathname = '$libdir/pbox_slice' +relocatable = false +superuser = false diff --git a/pgrx-examples/pbox_slice/src/lib.rs b/pgrx-examples/pbox_slice/src/lib.rs new file mode 100644 index 000000000..06c2f9007 --- /dev/null +++ b/pgrx-examples/pbox_slice/src/lib.rs @@ -0,0 +1,168 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//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. + +//! `pbox_slice` — demonstrates `MemCx::alloc_varlena` returning a +//! `PBox` directly to Postgres (no intermediate `Vec`), +//! plus the `PBox<[T]>` slice constructors used as in-context scratch buffers. + +use pgrx::datum::varlena_buf::{RawVarlena, VarlenaBuf}; +use pgrx::memcx::MemCx; +use pgrx::palloc::PBox; +use pgrx::prelude::*; + +pgrx::pg_module_magic!(name, version); + +/// Build a bytea of the requested length, filled with a `(i % 256)` pattern, allocated directly in the caller's `MemCx`. Returning `PBox` hands the varlena to Postgres without an extra Rust-side `Vec`. +#[pg_extern] +fn build_bytea<'mcx>(len: i32, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let n = len.max(0) as usize; + let mut buf: VarlenaBuf<'_> = cx.alloc_varlena(n).expect("OOM"); + for (i, b) in buf.payload_mut().iter_mut().enumerate() { + *b = (i & 0xff) as u8; + } + buf.into_pbox() +} + +/// Bytewise XOR of two byteas, truncated to the shorter length. Output is +/// allocated once at the exact final size — no intermediate `Vec`. +#[pg_extern] +fn xor_bytea<'mcx>(a: &[u8], b: &[u8], cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let len = a.len().min(b.len()); + let mut buf = cx.alloc_varlena(len).expect("OOM"); + let dst = buf.payload_mut(); + for i in 0..len { + dst[i] = a[i] ^ b[i]; + } + buf.into_pbox() +} + +/// Pack a SQL `int[]` into a little-endian binary bytea (4 bytes per element). +/// NULLs are encoded as zero. Demonstrates exact-size varlena allocation and in-place structured writes for binary-protocol-style serialization. +#[pg_extern] +fn pack_i32_array<'mcx>(arr: Array<'_, i32>, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let n = arr.len(); + let mut buf = cx.alloc_varlena(n * 4).expect("OOM"); + let dst = buf.payload_mut(); + for (i, v) in arr.iter().enumerate() { + let bytes = v.unwrap_or(0).to_le_bytes(); + dst[i * 4..(i + 1) * 4].copy_from_slice(&bytes); + } + buf.into_pbox() +} + +/// Apply `rounds` of a byte-rotating transformation to `data`. Uses a `PBox<[u8]>` scratch buffer allocated in `cx` (not a Rust `Vec`), then copies the final state into a `VarlenaBuf` for return. Demonstrates pairing the slice and varlena APIs in one function. +#[pg_extern] +fn hash_chain<'mcx>(data: &[u8], rounds: i32, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let rounds = rounds.max(0) as u32; + let mut scratch: PBox<[u8]> = PBox::from_slice_in(data, cx).expect("OOM"); + for _ in 0..rounds { + for b in scratch.iter_mut() { + *b = b.wrapping_add(1); + } + } + let mut out = cx.alloc_varlena(scratch.len()).expect("OOM"); + out.payload_mut().copy_from_slice(&scratch); + out.into_pbox() +} + +#[cfg(any(test, feature = "pg_test"))] +#[pg_schema] +mod tests { + use pgrx::prelude::*; + + #[pg_test] + fn test_build_bytea_length() { + let r = Spi::get_one::("SELECT length(build_bytea(1024))"); + assert_eq!(r, Ok(Some(1024))); + } + + #[pg_test] + fn test_build_bytea_pattern() { + // First 4 bytes of build_bytea(4) should be 0x00, 0x01, 0x02, 0x03. + let r = Spi::get_one::>("SELECT build_bytea(4)::bytea"); + assert_eq!(r, Ok(Some(vec![0u8, 1, 2, 3]))); + } + + #[pg_test] + fn test_build_bytea_zero_length() { + let r = Spi::get_one::("SELECT length(build_bytea(0))"); + assert_eq!(r, Ok(Some(0))); + } + + #[pg_test] + fn test_xor_bytea_basic() { + // 0xFF XOR 0x0F = 0xF0; 0xAA XOR 0x55 = 0xFF. + let r = Spi::get_one::>(r"SELECT xor_bytea('\xffaa'::bytea, '\x0f55'::bytea)"); + assert_eq!(r, Ok(Some(vec![0xf0u8, 0xff]))); + } + + #[pg_test] + fn test_xor_bytea_truncates_to_shorter() { + // 3-byte input XOR 5-byte input yields 3-byte output. + let r = Spi::get_one::( + r"SELECT length(xor_bytea('\x010203'::bytea, '\x0405060708'::bytea))", + ); + assert_eq!(r, Ok(Some(3))); + } + + #[pg_test] + fn test_xor_bytea_self_is_zero() { + let r = + Spi::get_one::>(r"SELECT xor_bytea('\xdeadbeef'::bytea, '\xdeadbeef'::bytea)"); + assert_eq!(r, Ok(Some(vec![0u8, 0, 0, 0]))); + } + + #[pg_test] + fn test_pack_i32_array_length() { + let r = Spi::get_one::("SELECT length(pack_i32_array(ARRAY[1,2,3,4]::int[]))"); + assert_eq!(r, Ok(Some(16))); // 4 elems * 4 bytes + } + + #[pg_test] + fn test_pack_i32_array_little_endian() { + // 1_i32 in LE = 01 00 00 00; 256_i32 in LE = 00 01 00 00. + let r = Spi::get_one::>("SELECT pack_i32_array(ARRAY[1, 256]::int[])::bytea"); + assert_eq!(r, Ok(Some(vec![1u8, 0, 0, 0, 0, 1, 0, 0]))); + } + + #[pg_test] + fn test_pack_i32_array_null_as_zero() { + let r = Spi::get_one::>("SELECT pack_i32_array(ARRAY[NULL, 1]::int[])::bytea"); + assert_eq!(r, Ok(Some(vec![0u8, 0, 0, 0, 1, 0, 0, 0]))); + } + + #[pg_test] + fn test_hash_chain_zero_rounds_is_identity() { + let r = Spi::get_one::>(r"SELECT hash_chain('\x010203'::bytea, 0)"); + assert_eq!(r, Ok(Some(vec![1u8, 2, 3]))); + } + + #[pg_test] + fn test_hash_chain_adds_rounds() { + // Each round increments every byte. 5 rounds on [0x00] = [0x05]. + let r = Spi::get_one::>(r"SELECT hash_chain('\x00'::bytea, 5)"); + assert_eq!(r, Ok(Some(vec![5u8]))); + } + + #[pg_test] + fn test_hash_chain_wraps() { + // 256 rounds on any byte returns the same byte (wrap_add). + let r = Spi::get_one::>(r"SELECT hash_chain('\xab'::bytea, 256)"); + assert_eq!(r, Ok(Some(vec![0xabu8]))); + } +} + +#[cfg(test)] +pub mod pg_test { + pub fn setup(_options: Vec<&str>) {} + pub fn postgresql_conf_options() -> Vec<&'static str> { + vec![] + } +} diff --git a/pgrx-unit-tests/src/tests/mod.rs b/pgrx-unit-tests/src/tests/mod.rs index a2da900f5..17c35bbe8 100644 --- a/pgrx-unit-tests/src/tests/mod.rs +++ b/pgrx-unit-tests/src/tests/mod.rs @@ -44,6 +44,7 @@ mod memcxt_tests; mod name_tests; mod numeric_tests; mod oid_tests; +mod pbox_slice_tests; mod pg_cast_tests; mod pg_extern_tests; mod pg_guard_tests; @@ -70,6 +71,7 @@ mod trigger_tests; mod type_ident_tests; mod uuid_tests; mod variadic_tests; +mod varlena_buf_tests; mod xact_callback_tests; mod xid64_tests; mod zero_datum_edge_cases; diff --git a/pgrx-unit-tests/src/tests/pbox_slice_tests.rs b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs new file mode 100644 index 000000000..59c0b9d8e --- /dev/null +++ b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs @@ -0,0 +1,253 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//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. + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use core::alloc::Layout; + use core::mem::MaybeUninit; + use pgrx::PgMemoryContexts; + use pgrx::memcx; + use pgrx::palloc::PBox; + use pgrx::prelude::*; + + /// Run `f` inside a freshly-created child memory context, then delete the + /// child and assert the parent's `MemoryContextMemAllocated` returned to + /// its baseline. If anything `f` allocated escaped the child context, + /// the parent's byte count grew and we fail. + /// + /// We use `PgMemoryContexts::new` (which wraps `AllocSetContextCreateExtended`) + /// and rely on its `Drop` impl to call `MemoryContextDelete`, so the child + /// context is freed even on unwind. The `recurse=false` argument to + /// `MemoryContextMemAllocated` measures only the parent's own blocks, + /// so child allocations don't pollute the baseline. + /// + /// Caveats: + /// - This harness is **not** unwind-safe in the strict sense: if `f` + /// panics, the `MemoryContextSwitchTo(prev)` call never runs and + /// `CurrentMemoryContext` is left pointing at the now-deleted child + /// until the outer `pg_test` harness catches the panic and restores + /// the test's outer context. Don't use this harness in code paths + /// where the post-panic state matters before that catch. + /// - `MemoryContextMemAllocated` reports the AllocSet's *block* count, + /// not byte-precise usage. Allocations that fit inside the initial + /// 8 KiB block (`ALLOCSET_DEFAULT_INITSIZE`) won't move the counter, + /// so a hypothetical leak smaller than that is invisible to this + /// harness. Tests that need leak detection should allocate well past + /// one block to force a fresh AllocSet block to be claimed. + fn with_scoped_memcx(f: F) + where + F: for<'mcx> FnOnce(&pgrx::memcx::MemCx<'mcx>), + { + unsafe { + let parent = pg_sys::CurrentMemoryContext; + let before = pg_sys::MemoryContextMemAllocated(parent, false); + + // Build the child as an Owned context; its Drop will MemoryContextDelete. + let child = PgMemoryContexts::new("pbox_slice_test_scope"); + // current_context() reads pg_sys::CurrentMemoryContext, so switch + // first, then restore. We deliberately do NOT use child.switch_to() + // here because its closure receives `&mut PgMemoryContexts`, not a + // `&MemCx<'_>`; the manual switch lets us reach memcx::current_context. + let prev = pg_sys::MemoryContextSwitchTo(child.value()); + memcx::current_context(|cx| f(cx)); + pg_sys::MemoryContextSwitchTo(prev); + + // Drop the owned child explicitly so MemoryContextDelete runs + // before we sample `after`. + drop(child); + + let after = pg_sys::MemoryContextMemAllocated(parent, false); + assert_eq!( + before, + after, + "leak: parent context grew by {} bytes after child deletion", + after.saturating_sub(before), + ); + } + } + + #[pg_test] + fn alloc_layout_returns_aligned_nonnull() { + memcx::current_context(|cx| { + let layout = Layout::from_size_align(128, 8).unwrap(); + let p = cx.alloc_layout(layout).expect("alloc"); + assert_eq!(p.as_ptr() as usize % 8, 0); + }); + } + + #[pg_test] + fn alloc_layout_zeroed_is_zero() { + memcx::current_context(|cx| { + let layout = Layout::from_size_align(64, 8).unwrap(); + let p = cx.alloc_layout_zeroed(layout).expect("alloc"); + let s = unsafe { core::slice::from_raw_parts(p.as_ptr(), 64) }; + assert!(s.iter().all(|&b| b == 0)); + }); + } + + #[pg_test] + fn new_uninit_slice_in_len_and_align() { + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = PBox::new_uninit_slice_in(16, cx).expect("alloc"); + assert_eq!(s.len(), 16); + assert_eq!(s.as_ptr() as *const u32 as usize % core::mem::align_of::(), 0); + }); + } + + #[pg_test] + fn new_zeroed_slice_in_is_zero() { + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = PBox::new_zeroed_slice_in(32, cx).expect("alloc"); + // SAFETY: zero is a valid bit pattern for u8. + let init: PBox<[u8]> = unsafe { s.assume_init() }; + assert_eq!(init.len(), 32); + assert!(init.iter().all(|&b| b == 0)); + }); + } + + #[pg_test] + fn from_slice_in_roundtrip() { + memcx::current_context(|cx| { + let src = [1u32, 2, 3, 4, 5]; + let dst: PBox<[u32]> = PBox::from_slice_in(&src, cx).expect("alloc"); + assert_eq!(&*dst, &src); + }); + } + + #[pg_test] + fn from_iter_in_happy_path() { + memcx::current_context(|cx| { + let v: Vec = (0..10).collect(); + let dst: PBox<[i32]> = PBox::from_iter_in(v.iter().copied(), cx).expect("alloc"); + assert_eq!(&*dst, v.as_slice()); + }); + } + + #[pg_test] + fn from_iter_in_drops_prefix_on_panic() { + use std::sync::atomic::{AtomicUsize, Ordering}; + static DROPS: AtomicUsize = AtomicUsize::new(0); + struct D(#[allow(dead_code)] usize); + impl Drop for D { + fn drop(&mut self) { + DROPS.fetch_add(1, Ordering::SeqCst); + } + } + + struct PanicIter { + i: usize, + max: usize, + panic_at: usize, + } + impl Iterator for PanicIter { + type Item = D; + fn next(&mut self) -> Option { + if self.i == self.panic_at { + panic!("boom") + } + if self.i >= self.max { + return None; + } + let v = D(self.i); + self.i += 1; + Some(v) + } + } + impl ExactSizeIterator for PanicIter { + fn len(&self) -> usize { + self.max - self.i + } + } + + DROPS.store(0, Ordering::SeqCst); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + memcx::current_context(|cx| { + let _ = PBox::<[D]>::from_iter_in(PanicIter { i: 0, max: 10, panic_at: 3 }, cx); + }); + })); + assert!(result.is_err()); + // Initialized prefix (0,1,2) must drop exactly once each. + assert_eq!(DROPS.load(Ordering::SeqCst), 3); + } + + #[pg_test] + fn zero_length_slice_is_safe() { + memcx::current_context(|cx| { + let s: PBox<[i64]> = PBox::from_slice_in(&[], cx).expect("alloc"); + assert_eq!(s.len(), 0); + assert!(s.is_empty()); + // Deref to an empty slice must not UB. + #[allow(clippy::explicit_auto_deref)] + let r: &[i64] = &*s; + assert_eq!(r.len(), 0); + }); + } + + #[pg_test] + fn oom_returns_err() { + memcx::current_context(|cx| { + // Request a length so large that `Layout::array` overflows + // (length * size > isize::MAX). PBox maps that to OutOfMemory + // before ever reaching palloc, so no ereport can occur. + let r = PBox::<[core::mem::MaybeUninit]>::new_uninit_slice_in(usize::MAX, cx); + assert!(r.is_err(), "expected OOM, got Ok"); + }); + } + + #[pg_test] + fn no_leak_after_context_delete() { + with_scoped_memcx(|cx| { + // Allocate a slice large enough to force the AllocSet to grab + // a fresh block (well past ALLOCSET_DEFAULT_INITSIZE = 8 KB), + // so a hypothetical leak would visibly grow the parent's + // MemoryContextMemAllocated. Intentionally do NOT pfree — + // the harness verifies that deleting the child reclaims it. + let _s: PBox<[u64]> = PBox::from_slice_in(&[1u64; 131_072], cx).expect("alloc"); + }); + } + + #[pg_test] + fn alloc_layout_rejects_over_max_alloc_size() { + memcx::current_context(|cx| { + // MaxAllocSize = 1 GiB - 1. Request 1 GiB + 1 byte through the + // non-huge path: must be rejected up front as OutOfMemory + // (not via PG ereport longjmp). + let layout = Layout::from_size_align((1usize << 30) + 1, 8).unwrap(); + assert!(cx.alloc_layout(layout).is_err()); + assert!(cx.alloc_layout_zeroed(layout).is_err()); + }); + } + + #[pg_test] + fn alloc_huge_layout_accepts_over_max_alloc_size() { + memcx::current_context(|cx| { + // 1.5 GiB — above MaxAllocSize, well below MaxAllocHugeSize. + // Succeeds on a machine with sufficient RAM; on a tight test + // environment we accept OutOfMemory as a valid response. + // The critical property: no longjmp, no panic — Rust gets + // back a Result either way. + let layout = Layout::from_size_align(3usize << 29, 8).unwrap(); + let _ = cx.alloc_huge_layout(layout); + }); + } + + #[pg_test] + fn alloc_huge_layout_zeroed_is_zero() { + memcx::current_context(|cx| { + // Small probe — verifies the ZERO flag still flows through + // when HUGE is set. No need to actually allocate huge here. + let layout = Layout::from_size_align(64, 8).unwrap(); + let p = cx.alloc_huge_layout_zeroed(layout).expect("alloc"); + let s = unsafe { core::slice::from_raw_parts(p.as_ptr(), 64) }; + assert!(s.iter().all(|&b| b == 0)); + }); + } +} diff --git a/pgrx-unit-tests/src/tests/varlena_buf_tests.rs b/pgrx-unit-tests/src/tests/varlena_buf_tests.rs new file mode 100644 index 000000000..918172720 --- /dev/null +++ b/pgrx-unit-tests/src/tests/varlena_buf_tests.rs @@ -0,0 +1,60 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//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. + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use pgrx::datum::varlena_buf::VarlenaBuf; + use pgrx::memcx; + use pgrx::pg_sys; + use pgrx::prelude::*; + + #[pg_test] + fn alloc_varlena_header_is_valid() { + memcx::current_context(|cx| { + let buf: VarlenaBuf = cx.alloc_varlena(64).expect("alloc"); + assert_eq!(buf.payload().len(), 64); + // payload should be zeroed + assert!(buf.payload().iter().all(|&b| b == 0)); + // total_size = header + payload + assert_eq!(buf.total_size(), 64 + pg_sys::VARHDRSZ); + // hand off to PG and check via varsize_any + let raw = buf.into_raw(); + unsafe { + assert_eq!(pgrx::varlena::varsize_any(raw), 64 + pg_sys::VARHDRSZ,); + } + }); + } + + #[pg_test] + fn write_then_read_payload() { + memcx::current_context(|cx| { + let mut buf: VarlenaBuf = cx.alloc_varlena(16).expect("alloc"); + for (i, b) in buf.payload_mut().iter_mut().enumerate() { + *b = i as u8; + } + for (i, &b) in buf.payload().iter().enumerate() { + assert_eq!(b as usize, i); + } + }); + } + + #[pg_test] + fn inspect_varlena_roundtrip() { + memcx::current_context(|cx| { + let mut buf: VarlenaBuf = cx.alloc_varlena(11).expect("alloc"); + buf.payload_mut().copy_from_slice(b"hello world"); + let ptr = buf.into_raw(); + // SAFETY: `ptr` came from alloc_varlena above, still in `cx`, header valid. + let view = unsafe { cx.inspect_varlena(ptr) }; + assert_eq!(view.payload(), b"hello world"); + }); + } +} diff --git a/pgrx/src/datum/mod.rs b/pgrx/src/datum/mod.rs index 7dd3f38fa..bccdf6d4e 100644 --- a/pgrx/src/datum/mod.rs +++ b/pgrx/src/datum/mod.rs @@ -32,6 +32,7 @@ mod tuples; mod unbox; mod uuid; mod varlena; +pub mod varlena_buf; pub use self::uuid::*; pub use crate::datetime::support as datetime_support; @@ -50,6 +51,7 @@ pub use numeric::{AnyNumeric, Numeric}; pub use range::*; pub use unbox::*; pub use varlena::*; +pub use varlena_buf::{RawVarlena, VarlenaBuf}; use crate::memcx::MemCx; use crate::pg_sys; diff --git a/pgrx/src/datum/varlena_buf.rs b/pgrx/src/datum/varlena_buf.rs new file mode 100644 index 000000000..313d15c42 --- /dev/null +++ b/pgrx/src/datum/varlena_buf.rs @@ -0,0 +1,196 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//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. + +//! Safe, lifetime-bound varlena allocation. +//! +//! [`MemCx::alloc_varlena`] returns a [`VarlenaBuf<'mcx>`] whose 4-byte header +//! is set to a valid total size before the wrapper exists. Writing into the +//! payload via [`VarlenaBuf::payload_mut`] cannot resize, so the header +//! invariant is preserved. +//! +//! # Examples +//! +//! Building a bytea and handing it to Postgres: +//! +//! ```no_run +//! use pgrx::memcx; +//! use pgrx::pg_sys; +//! +//! memcx::current_context(|cx| { +//! let mut buf = cx.alloc_varlena(11).expect("OOM"); +//! buf.payload_mut().copy_from_slice(b"hello world"); +//! let ptr: *mut pg_sys::varlena = buf.into_raw(); +//! // `ptr` can be returned from a #[pg_extern] as a varlena Datum. +//! // The allocation lives in `cx` until the memory context resets. +//! let _ = ptr; +//! }); +//! ``` +//! +//! See: `docs/superpowers/specs/2026-06-15-memcx-slice-varlena-parity-design.md` + +use crate::memcx::MemCx; +use crate::palloc::PBox; +use crate::pg_sys; +use core::ptr::NonNull; + +/// A `varlena` allocation with a valid 4-byte header, maintained as a type invariant. +/// +/// `RawVarlena` is a DST whose in-memory representation is the full varlena byte +/// sequence (header + payload). Safely obtainable only through +/// [`MemCx::alloc_varlena`] (which sets the header) or [`MemCx::inspect_varlena`] +/// (caller-asserted). +/// +/// # Layout +/// +/// `#[repr(transparent)]` over `[u8]`. The header is at byte offset 0; payload +/// starts at offset `VARHDRSZ` (= 4). The total length matches the header's +/// stored size. +#[repr(transparent)] +pub struct RawVarlena { + bytes: [u8], +} + +impl RawVarlena { + /// Total size in bytes (header + payload), as recorded in the header. + pub fn total_size(&self) -> usize { + // SAFETY: type invariant — header at offset 0 is a valid 4-byte varlena header. + unsafe { crate::varlena::varsize_any(self.as_ptr()) } + } + + /// Read-only access to the payload bytes (excludes the header). + pub fn payload(&self) -> &[u8] { + let total = self.total_size(); + let hdr = pg_sys::VARHDRSZ; + let payload_len = total.checked_sub(hdr).expect("corrupt varlena: total < VARHDRSZ"); + // SAFETY: payload starts at offset VARHDRSZ; `checked_sub` above rules out + // an underflowed length on a corrupt header (release-mode UB guard). + unsafe { + core::slice::from_raw_parts((self as *const Self as *const u8).add(hdr), payload_len) + } + } + + /// Mutable access to the payload bytes. Cannot be used to grow or shrink: + /// the header (and therefore `total_size`) is unchanged. + pub fn payload_mut(&mut self) -> &mut [u8] { + let total = self.total_size(); + let hdr = pg_sys::VARHDRSZ; + let payload_len = total.checked_sub(hdr).expect("corrupt varlena: total < VARHDRSZ"); + // SAFETY: payload starts at offset VARHDRSZ; `checked_sub` above rules out + // an underflowed length on a corrupt header. We hold `&mut self`, so the + // borrow is exclusive. + unsafe { + core::slice::from_raw_parts_mut((self as *mut Self as *mut u8).add(hdr), payload_len) + } + } + + /// Raw `*const pg_sys::varlena` pointer. The pointee remains valid for + /// the lifetime of `&self`. + pub fn as_ptr(&self) -> *const pg_sys::varlena { + self as *const Self as *const pg_sys::varlena + } + + /// Raw `*mut pg_sys::varlena` pointer. + pub fn as_mut_ptr(&mut self) -> *mut pg_sys::varlena { + self as *mut Self as *mut pg_sys::varlena + } +} + +/// A varlena buffer: a `PBox`-owned `RawVarlena`, scoped to a memory context. +pub struct VarlenaBuf<'mcx> { + inner: PBox<'mcx, RawVarlena>, +} + +impl<'mcx> VarlenaBuf<'mcx> { + /// # Safety + /// `data_ptr` must point to a valid varlena allocation in `cx` whose header has + /// already been set, and the total byte length must equal `total`. + pub(crate) unsafe fn from_raw(data_ptr: NonNull, total: usize, cx: &MemCx<'mcx>) -> Self { + // Build a *mut [u8] fat pointer of the right length, then cast to + // *mut RawVarlena (sound: RawVarlena is #[repr(transparent)] over [u8]). + let slice_ptr: *mut [u8] = core::ptr::slice_from_raw_parts_mut(data_ptr.as_ptr(), total); + let raw_ptr: *mut RawVarlena = slice_ptr as *mut RawVarlena; + // SAFETY: data_ptr was non-null, so the fat pointer's data pointer is non-null. + let nn = NonNull::new_unchecked(raw_ptr); + let inner = PBox::from_raw_in(nn, cx); + VarlenaBuf { inner } + } + + /// Read-only payload access. + pub fn payload(&self) -> &[u8] { + // SAFETY: `inner.ptr` is a valid fat pointer to a RawVarlena whose header + // is valid by construction. + unsafe { (&*self.inner.as_ptr()).payload() } + } + + /// Mutable payload access. + pub fn payload_mut(&mut self) -> &mut [u8] { + // SAFETY: `inner` exclusively owns the allocation; we hold `&mut self`. + unsafe { (&mut *self.inner.as_mut_ptr()).payload_mut() } + } + + /// Total size (header + payload) as recorded in the header. + pub fn total_size(&self) -> usize { + // SAFETY: same as `payload`. + unsafe { (&*self.inner.as_ptr()).total_size() } + } + + /// Convert into the underlying `PBox`. Ownership of the + /// memory-context-bound allocation is preserved. + pub fn into_pbox(self) -> PBox<'mcx, RawVarlena> { + self.inner + } + + /// Leak the `PBox` and return a raw `*mut pg_sys::varlena` for handoff + /// to Postgres (e.g., as a function return). The allocation remains in + /// `cx` and is reclaimed at context reset. + pub fn into_raw(self) -> *mut pg_sys::varlena { + let mut me = core::mem::ManuallyDrop::new(self); + me.inner.as_mut_ptr() as *mut pg_sys::varlena + } +} + +// --- Returning `PBox` directly from a `#[pg_extern]` ----------- +// +// `RawVarlena` deliberately does NOT implement `BorrowDatum` (it has no +// fixed in-Datum representation: a varlena is always pass-by-reference and +// its size is header-driven). The blanket `BoxRet`/`SqlTranslatable` impls +// for `PBox` in `palloc/pbox.rs` therefore do not cover it, +// and we add specific impls here. + +use crate::callconv::{BoxRet, FcInfo}; +use crate::datum::Datum; +use pgrx_sql_entity_graph::metadata::{ + ArgumentError, ReturnsError, ReturnsRef, SqlMappingRef, SqlTranslatable, +}; + +/// Returning a `PBox<'mcx, RawVarlena>` from a `#[pg_extern]` hands the +/// underlying varlena pointer to Postgres unchanged. No copy: the +/// allocation stays in the caller's `MemCx` and is reclaimed on context +/// reset. +unsafe impl<'mcx> BoxRet for PBox<'mcx, RawVarlena> { + unsafe fn box_into<'fcx>(self, fcinfo: &mut FcInfo<'fcx>) -> Datum<'fcx> { + // Leak the PBox; ownership of the allocation transfers to Postgres. + let mut me = core::mem::ManuallyDrop::new(self); + let ptr = me.as_mut_ptr() as *mut pg_sys::varlena; + // SAFETY: `ptr` points at a valid varlena allocation in the caller's + // memory context; Postgres owns it from here on. + unsafe { fcinfo.return_raw_datum(pg_sys::Datum::from(ptr)) } + } +} + +/// `PBox<'mcx, RawVarlena>` maps to SQL `bytea`. +unsafe impl<'mcx> SqlTranslatable for PBox<'mcx, RawVarlena> { + const TYPE_IDENT: &'static str = "bytea"; + const TYPE_ORIGIN: pgrx_sql_entity_graph::metadata::TypeOrigin = + pgrx_sql_entity_graph::metadata::TypeOrigin::External; + const ARGUMENT_SQL: Result = Ok(SqlMappingRef::literal("bytea")); + const RETURN_SQL: Result = + Ok(ReturnsRef::One(SqlMappingRef::literal("bytea"))); +} diff --git a/pgrx/src/memcx.rs b/pgrx/src/memcx.rs index d2fd5177f..0e6c83c66 100644 --- a/pgrx/src/memcx.rs +++ b/pgrx/src/memcx.rs @@ -13,8 +13,15 @@ use pgrx_sql_entity_graph::metadata::{ use crate::callconv::{Arg, ArgAbi}; use crate::nullable::Nullable; use crate::pg_sys; +use core::alloc::Layout; use core::{marker::PhantomData, ptr::NonNull}; +/// Postgres's `MaxAllocSize` (1 GiB - 1). `bindgen` does not expose it +/// (see `pgrx/src/array.rs:219`), so we mirror the definition here. Any +/// `palloc` call exceeding this size requires `MCXT_ALLOC_HUGE`; without +/// that flag, Postgres `ereport`s — which longjmps out of Rust frames. +const MAX_ALLOC_SIZE: usize = 0x3fff_ffff; + /// A borrowed memory context. #[repr(transparent)] pub struct MemCx<'mcx> { @@ -50,6 +57,83 @@ impl<'mcx> MemCx<'mcx> { NonNull::new(ptr.cast()).ok_or(OutOfMemory::new()) } + /// Allocate a buffer matching `layout`. + /// + /// Sizes above Postgres's `MaxAllocSize` (1 GiB - 1) are rejected up + /// front with `OutOfMemory`. Without this guard, the underlying palloc + /// call would `ereport` and longjmp out of the Rust frame. For genuine + /// huge allocations, use [`MemCx::alloc_huge_layout`] instead. + /// + /// On PG >= 16 this uses `MemoryContextAllocAligned` for arbitrary power-of-two alignment. On earlier versions this falls back to size-only `MemoryContextAllocExtended`; callers requiring alignment greater than palloc's `MAXIMUM_ALIGNOF` (typically 8) must enforce that constraint at compile time at the type layer. + pub fn alloc_layout(&self, layout: Layout) -> Result, OutOfMemory> { + if layout.size() > MAX_ALLOC_SIZE { + return Err(OutOfMemory::new()); + } + self.alloc_layout_with_flags(layout, pg_sys::MCXT_ALLOC_NO_OOM as i32) + } + + /// Like [`MemCx::alloc_layout`] but zeroes the allocation. + pub fn alloc_layout_zeroed(&self, layout: Layout) -> Result, OutOfMemory> { + if layout.size() > MAX_ALLOC_SIZE { + return Err(OutOfMemory::new()); + } + self.alloc_layout_with_flags( + layout, + (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_ZERO) as i32, + ) + } + + /// Like [`MemCx::alloc_layout`] but permits sizes above `MaxAllocSize` + /// (1 GiB - 1), up to `MaxAllocHugeSize` (~half the address space). + /// + /// Use only when the allocation is genuinely known to be huge. The 1 GiB + /// cap on [`MemCx::alloc_layout`] is a deliberate sanity check that + /// catches sizing bugs early; opting into `_huge_` should be a + /// considered choice, not a routine workaround. + /// + /// Note: varlena values cannot exceed 1 GiB by design (the 4-byte + /// header encodes 30-bit size). Huge allocations are only useful for + /// Rust-side scratch buffers, not for values returned to Postgres. + pub fn alloc_huge_layout(&self, layout: Layout) -> Result, OutOfMemory> { + self.alloc_layout_with_flags( + layout, + (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_HUGE) as i32, + ) + } + + /// Like [`MemCx::alloc_huge_layout`] but zeroes the allocation. + pub fn alloc_huge_layout_zeroed(&self, layout: Layout) -> Result, OutOfMemory> { + self.alloc_layout_with_flags( + layout, + (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_HUGE | pg_sys::MCXT_ALLOC_ZERO) as i32, + ) + } + + #[inline] + fn alloc_layout_with_flags( + &self, + layout: Layout, + flags: i32, + ) -> Result, OutOfMemory> { + let ptr = unsafe { + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + { + pg_sys::MemoryContextAllocAligned( + self.ptr.as_ptr(), + layout.size(), + layout.align(), + flags, + ) + } + #[cfg(not(any(feature = "pg16", feature = "pg17", feature = "pg18")))] + { + let _ = layout.align(); + pg_sys::MemoryContextAllocExtended(self.ptr.as_ptr(), layout.size(), flags) + } + }; + NonNull::new(ptr.cast()).ok_or_else(OutOfMemory::new) + } + /// Stores the current memory context, switches to *this* memory context, /// and executes the closure `f`. /// Once `f` completes, the previous current memory context is restored. @@ -66,6 +150,42 @@ impl<'mcx> MemCx<'mcx> { pg_sys::MemoryContextSwitchTo(remembered); res } + + /// Allocate a varlena with a valid 4-byte header and `payload_len` zeroed payload bytes. The returned [`VarlenaBuf`] maintains the header as a type invariant. + /// + /// [`VarlenaBuf`]: crate::datum::varlena_buf::VarlenaBuf + pub fn alloc_varlena( + &self, + payload_len: usize, + ) -> Result, OutOfMemory> { + let total = (pg_sys::VARHDRSZ).checked_add(payload_len).ok_or_else(OutOfMemory::new)?; + let layout = core::alloc::Layout::from_size_align(total, core::mem::align_of::()) + .map_err(|_| OutOfMemory::new())?; + let raw = self.alloc_layout_zeroed(layout)?; + // Set the 4-byte header. The total_size encoded in the header is `total` (header bytes + payload bytes), per Postgres convention for 4-byte headers. + // SAFETY: `raw` points at `total` zeroed bytes, large enough for a 4-byte header. + unsafe { + crate::varlena::set_varsize_4b(raw.as_ptr() as *mut pg_sys::varlena, total as i32); + } + // SAFETY: `raw` is non-null, sized at `total`, header is now valid. + Ok(unsafe { crate::datum::varlena_buf::VarlenaBuf::from_raw(raw, total, self) }) + } + + /// Inspect an existing varlena pointer as a `&RawVarlena`. + /// + /// # Safety + /// - `ptr` must point to a valid, non-TOASTed varlena + /// - the header at `ptr` must have its size correctly set + /// - the returned reference must not outlive the underlying allocation + pub unsafe fn inspect_varlena<'a>( + &self, + ptr: *const pg_sys::varlena, + ) -> &'a crate::datum::varlena_buf::RawVarlena { + let total = crate::varlena::varsize_any(ptr); + let slice_ptr: *const [u8] = core::ptr::slice_from_raw_parts(ptr as *const u8, total); + let raw_ptr: *const crate::datum::varlena_buf::RawVarlena = slice_ptr as *const _; + &*raw_ptr + } } #[derive(Debug)] diff --git a/pgrx/src/palloc/pbox.rs b/pgrx/src/palloc/pbox.rs index 69e181d67..e0d239a0c 100644 --- a/pgrx/src/palloc/pbox.rs +++ b/pgrx/src/palloc/pbox.rs @@ -1,9 +1,48 @@ +//! Lifetime-bound boxed allocations in a Postgres memory context. +//! +//! [`PBox<'mcx, T>`] is the safe, lifetime-bound analogue of [`alloc::boxed::Box`] +//! for memory allocated inside a [`MemCx<'mcx>`]. Unlike `Box`, the value is +//! reclaimed when the memory context is reset — not when `PBox` is dropped. +//! +//! # Examples +//! +//! Allocating a zero-initialized slice and reading it back: +//! +//! ```no_run +//! use pgrx::memcx; +//! use pgrx::palloc::PBox; +//! use core::mem::MaybeUninit; +//! +//! memcx::current_context(|cx| { +//! let s = PBox::<[MaybeUninit]>::new_zeroed_slice_in(8, cx) +//! .expect("OOM"); +//! // SAFETY: zero is a valid bit pattern for u32. +//! let s: PBox<[u32]> = unsafe { s.assume_init() }; +//! assert_eq!(s.len(), 8); +//! assert!(s.iter().all(|&n| n == 0)); +//! }); +//! ``` +//! +//! Copying from an existing slice: +//! +//! ```no_run +//! use pgrx::memcx; +//! use pgrx::palloc::PBox; +//! +//! memcx::current_context(|cx| { +//! let src = [1u32, 2, 3, 4]; +//! let dst: PBox<[u32]> = PBox::from_slice_in(&src, cx).expect("OOM"); +//! assert_eq!(&*dst, &src); +//! }); +//! ``` + use crate::callconv::{BoxRet, FcInfo}; use crate::datum::{BorrowDatum, Datum}; use crate::layout::PassBy; use crate::memcx::{MemCx, OutOfMemory}; use crate::pg_sys; use core::marker::PhantomData; +use core::mem::MaybeUninit; use core::ops::{Deref, DerefMut}; use core::ptr::NonNull; @@ -15,6 +54,19 @@ use pgrx_sql_entity_graph::metadata::{ [stdbox]: alloc::boxed::Box + +Lifetime escape is statically rejected. The following snippet must fail to type-check because the `MemCx`-bound lifetime cannot outlive the context it was borrowed from: + +```compile_fail +use pgrx::memcx::MemCx; +use pgrx::palloc::PBox; + +fn smuggle<'long, 'short>(cx: &MemCx<'short>) -> PBox<'long, [u32]> { + PBox::<[core::mem::MaybeUninit]>::new_uninit_slice_in(4, cx) + .map(|p| unsafe { p.assume_init() }) + .unwrap() +} +``` */ #[repr(transparent)] pub struct PBox<'mcx, T: ?Sized> { @@ -24,7 +76,7 @@ pub struct PBox<'mcx, T: ?Sized> { impl<'mcx, T: ?Sized> PBox<'mcx, T> { // # Safety - // The same constraints as [`Box::from_raw`], AND + // The same constraints as [`Box::from_raw`] // - you assert the pointer was allocated in the `MemCx` // - you assert the pointer may be freed by `pfree` pub unsafe fn from_raw_in(ptr: NonNull, _cx: &MemCx<'mcx>) -> PBox<'mcx, T> { @@ -41,13 +93,179 @@ impl<'mcx, T: Sized> PBox<'mcx, T> { pub fn try_new_in(val: T, memcx: &MemCx<'mcx>) -> Result { const { assert!(align_of::() <= size_of::()) }; let ptr = memcx.alloc_bytes(size_of::())?.cast(); - // SAFETY: We were guaranteed an appropriately sized allocation to write to, - // and we have asserted our alignment maximum was upheld + // SAFETY: We were guaranteed an appropriately sized allocation to write to, and we have asserted our alignment maximum was upheld unsafe { ptr.write(val) }; Ok(PBox { ptr, _cx: PhantomData }) } } +impl<'mcx, T: ?Sized> PBox<'mcx, T> { + /// Raw pointer to the boxed value. Does not dereference. + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() as *const T + } + /// Raw mutable pointer to the boxed value. Does not dereference. + pub fn as_mut_ptr(&mut self) -> *mut T { + self.ptr.as_ptr() + } +} + +impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { + /// Allocate a slice of `len` uninitialized `T` inside `cx`. + /// + /// Mirrors [`alloc::boxed::Box::new_uninit_slice`] but the resulting slice is bound to the lifetime of `cx`. + pub fn new_uninit_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_layout(layout)?; + // SAFETY: `raw` is non-null, sized by `layout`; MaybeUninit needs no initialization, so building the fat pointer over `len` elements is sound. + let slice_ptr: *mut [MaybeUninit] = + core::ptr::slice_from_raw_parts_mut(raw.as_ptr().cast::>(), len); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } + + /// Allocate a slice of `len` zero-initialized `T` (kept as `MaybeUninit`) inside `cx`. Caller can transmute via [`assume_init`] when zero is a valid bit pattern for `T`. + pub fn new_zeroed_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_layout_zeroed(layout)?; + // SAFETY: `raw` is non-null, sized by `layout`; MaybeUninit needs no init. + let slice_ptr: *mut [core::mem::MaybeUninit] = core::ptr::slice_from_raw_parts_mut( + raw.as_ptr().cast::>(), + len, + ); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } + + /// # Safety + /// Caller asserts every element of the slice has been initialized. + pub unsafe fn assume_init(self) -> PBox<'mcx, [T]> { + let p = self.ptr.as_ptr() as *mut [T]; + // SAFETY: caller upholds the init invariant; the data pointer was already non-null (it came from PBox::new_uninit_slice_in / new_zeroed_slice_in). + let ptr = unsafe { NonNull::new_unchecked(p) }; + // Avoid running any (non-existent) Drop on `self`. + core::mem::forget(self); + PBox { ptr, _cx: PhantomData } + } + + /// Like [`new_uninit_slice_in`] but permits allocations above Postgres's 1 GiB `MaxAllocSize`. Only use for genuinely huge Rust-side scratch buffers — values returned to Postgres still cap at 1 GiB regardless of how they were allocated. + /// + /// [`new_uninit_slice_in`]: PBox::new_uninit_slice_in + pub fn new_huge_uninit_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_huge_layout(layout)?; + // SAFETY: same as new_uninit_slice_in. + let slice_ptr: *mut [MaybeUninit] = + core::ptr::slice_from_raw_parts_mut(raw.as_ptr().cast::>(), len); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } + + /// Zeroed counterpart of [`new_huge_uninit_slice_in`]. + pub fn new_huge_zeroed_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_huge_layout_zeroed(layout)?; + // SAFETY: same as new_zeroed_slice_in. + let slice_ptr: *mut [MaybeUninit] = + core::ptr::slice_from_raw_parts_mut(raw.as_ptr().cast::>(), len); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } +} + +impl<'mcx, T> PBox<'mcx, [T]> { + /// Number of elements in the slice. + pub fn len(&self) -> usize { + self.ptr.len() + } + + /// `true` if the slice has zero elements. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +impl<'mcx, T: Copy> PBox<'mcx, [T]> { + /// Allocate a slice of `src.len()` elements in `cx` and copy `src` into it. + pub fn from_slice_in(src: &[T], cx: &MemCx<'mcx>) -> Result { + let mut uninit = PBox::<[MaybeUninit]>::new_uninit_slice_in(src.len(), cx)?; + // SAFETY: `uninit.len() == src.len()` by construction. T: Copy, so a bytewise copy fully initializes every element with no Drop concerns. + unsafe { + let dst = (uninit.as_mut_ptr() as *mut MaybeUninit).cast::(); + core::ptr::copy_nonoverlapping(src.as_ptr(), dst, src.len()); + Ok(uninit.assume_init()) + } + } +} + +impl<'mcx, T> PBox<'mcx, [T]> { + /// Allocate a slice of `iter.len()` elements in `cx` and write each element produced by `iter`. Requires `ExactSizeIterator` so the allocation size is known upfront with no resizing. + /// + /// On panic during iteration, the initialized prefix is dropped exactly once. The underlying memory remains in `cx` until the context resets. + /// + /// # Errors + /// + /// - `Err(OutOfMemory)` if the initial allocation fails. + /// - `Err(OutOfMemory)` if `iter` yields **fewer** items than its + /// `ExactSizeIterator::len()` advertised. This is technically a contract violation by the iterator rather than a true OOM, but + /// reusing `OutOfMemory` avoids introducing a new error type for an edge case `ExactSizeIterator` callers shouldn't hit. In this case the initialized prefix is dropped via the RAII guard before the error returns; the buffer itself stays in `cx`. + /// + /// If `iter` yields **more** items than promised, the surplus items are dropped at the loop boundary and the function returns `Ok` with the first `len` elements. + pub fn from_iter_in(iter: I, cx: &MemCx<'mcx>) -> Result + where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, + { + let iter = iter.into_iter(); + let len = iter.len(); + let mut uninit = PBox::<[MaybeUninit]>::new_uninit_slice_in(len, cx)?; + + struct Guard { + base: *mut T, + initialized: usize, + } + impl Drop for Guard { + fn drop(&mut self) { + if self.initialized > 0 { + // SAFETY: the first `initialized` elements were written via ptr::write and are valid T. + unsafe { + let s = core::slice::from_raw_parts_mut(self.base, self.initialized); + core::ptr::drop_in_place(s); + } + } + } + } + + // Get a thin pointer to the first element of the uninit buffer. + let base: *mut T = uninit.as_mut_ptr().cast::(); + let mut guard = Guard { base, initialized: 0 }; + + let mut count = 0usize; + for item in iter { + if count >= len { + // Iterator over-yielded vs ExactSizeIterator promise. Stop here; + // we'll fail below since `count < len`. + break; + } + // SAFETY: `count < len`, so `base.add(count)` is in-bounds for the freshly allocated slice. The slot is uninit MaybeUninit; writing a T is sound. + unsafe { core::ptr::write(base.add(count), item) }; + count += 1; + guard.initialized = count; + } + + if count < len { + // Iterator under-yielded. Drop the initialized prefix via Guard, + // and return an error. The buffer itself stays in `cx`. + return Err(OutOfMemory::new()); + } + + // All elements initialized successfully — disarm the guard and assume_init. + core::mem::forget(guard); + // SAFETY: every slot in the slice was written above (count == len). + Ok(unsafe { uninit.assume_init() }) + } +} + unsafe impl<'mcx, T> BoxRet for PBox<'mcx, T> where T: ?Sized + BorrowDatum, @@ -57,8 +275,7 @@ where PassBy::Value => { // start with a zeroed Datum, just to minimize funny business let mut datum = pg_sys::Datum::null(); - // SAFETY: Due to BorrowDatum, this type has a definite size less than a Datum, - // and PBox must have an initialized pointee, so a copy is sound-by-construction. + // SAFETY: Due to BorrowDatum, this type has a definite size less than a Datum, and PBox must have an initialized pointee, so a copy is sound-by-construction. unsafe { let size = size_of_val(&*self.ptr.as_ptr()); debug_assert!(size <= size_of::()); @@ -75,8 +292,7 @@ where } } -/// SAFETY: SQL has no "pointers" so by-val and by-ref calling conventions are identical, -/// and all `PBox` truly does is enable pass-by-ref returns of unsized values. +/// SAFETY: SQL has no "pointers" so by-val and by-ref calling conventions are identical,and all `PBox` truly does is enable pass-by-ref returns of unsized values. unsafe impl<'mcx, T> SqlTranslatable for PBox<'mcx, T> where T: SqlTranslatable + ?Sized, @@ -108,3 +324,19 @@ where unsafe { self.ptr.as_mut() } } } + +impl<'mcx, T> Deref for PBox<'mcx, [T]> { + type Target = [T]; + + fn deref(&self) -> &[T] { + // SAFETY: `self.ptr` is a valid fat pointer to an initialized [T] slice (the Sized constructors only return PBox<[T]> after assume_init or from_slice_in/from_iter_in, all of which uphold initialization). + unsafe { self.ptr.as_ref() } + } +} + +impl<'mcx, T> DerefMut for PBox<'mcx, [T]> { + fn deref_mut(&mut self) -> &mut [T] { + // SAFETY: same as Deref, plus we hold `&mut self` so aliasing is exclusive. + unsafe { self.ptr.as_mut() } + } +} diff --git a/pgrx/src/varlena.rs b/pgrx/src/varlena.rs index 9cbe732da..c4c27cb80 100644 --- a/pgrx/src/varlena.rs +++ b/pgrx/src/varlena.rs @@ -14,8 +14,7 @@ use core::{ops::DerefMut, slice, str}; /// # Safety /// -/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer -/// that is aligned to 4 bytes, and that the `len` is a half of [`i32::MAX`] +/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer that is aligned to 4 bytes, and that the `len` is a half of [`i32::MAX`] #[inline(always)] pub unsafe fn set_varsize_4b(ptr: *mut pg_sys::varlena, len: i32) { // #ifdef WORDS_BIGENDIAN @@ -28,8 +27,7 @@ pub unsafe fn set_varsize_4b(ptr: *mut pg_sys::varlena, len: i32) { // SAFETY: A varlena can be safely cast to a varattrib_4b let header = &mut (*ptr.cast::()).va_4byte.deref_mut().va_header; - // Using core::ptr::write(), which never calls drop(), to prevent - // automatically dropping a field of a ManuallyDrop + // Using core::ptr::write(), which never calls drop(), to prevent automatically dropping a field of a ManuallyDrop core::ptr::write(header, encode_vlen_4b(len)) } @@ -51,8 +49,7 @@ pub(crate) fn encode_vlen_1b(len: i32) -> u8 { /// # Safety /// -/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer -/// that is aligned to 4 bytes. +/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer that is aligned to 4 bytes. #[inline(always)] #[deprecated(since = "0.12.0", note = "you probably meant set_varsize_4b")] pub unsafe fn set_varsize(ptr: *mut pg_sys::varlena, len: i32) { @@ -417,8 +414,7 @@ pub unsafe fn vardata_any(ptr: *const pg_sys::varlena) -> *const std::os::raw::c /// /// This function is unsafe because it blindly assumes the provided varlena pointer is non-null. /// -/// Note also that this function is zero-copy and the underlying Rust &str is backed by Postgres-allocated -/// memory. As such, the return value will become invalid the moment Postgres frees the varlena +/// Note also that this function is zero-copy and the underlying Rust &str is backed by Postgres-allocated memory. As such, the return value will become invalid the moment Postgres frees the varlena #[inline] pub unsafe fn text_to_rust_str<'a>( varlena: *const pg_sys::varlena, @@ -448,8 +444,7 @@ pub unsafe fn text_to_rust_str_unchecked<'a>(varlena: *const pg_sys::varlena) -> /// /// This function is unsafe because it blindly assumes the provided varlena pointer is non-null. /// -/// Note also that this function is zero-copy and the underlying Rust `&[u8]` slice is backed by Postgres-allocated -/// memory. As such, the return value will become invalid the moment Postgres frees the varlena +/// Note also that this function is zero-copy and the underlying Rust `&[u8]` slice is backed by Postgres-allocated memory. As such, the return value will become invalid the moment Postgres frees the varlena #[inline] pub unsafe fn varlena_to_byte_slice<'a>(varlena: *const pg_sys::varlena) -> &'a [u8] { let len = varsize_any_exhdr(varlena); From 74a60a6a1a69876bd4582017a1e3ff574f5a2401 Mon Sep 17 00:00:00 2001 From: danielshih Date: Tue, 16 Jun 2026 03:17:09 +0000 Subject: [PATCH 2/3] add example pbox_slice as a mini vector-embedding store Replace the synthetic xor_bytea / pack_i32_array / hash_chain demos with a focused mini vector-embedding extension that mirrors a real RAG / k-NN workload: - embedding_pack(real[]) -> bytea (build, write side) - embedding_dims(bytea) -> int (read side, layout check) - embedding_l2_squared(bytea, bytea) (read 2, no alloc) - embedding_dot(bytea, bytea) (read 2, no alloc) - embedding_cosine(bytea, bytea) (read 2, no alloc, NaN edges) - embedding_mean(bytea, bytea) -> bytea (read 2 + write 1) - embedding_normalize(bytea) -> bytea (read 1 + write 1) - build_bytea(int) -> bytea (kept as smoke-test fixture) The 4-byte-per-dim native-endian layout matches pgvector's binary convention, so the example reads as a stripped-down version of code extension authors actually write rather than a synthetic benchmark. README contrasts this against pgvector and points readers there for production workloads. * enhance varlena handling with improved safety checks and inline header support * add alignment assertions for slice allocation methods * add MAX_ALLOC_HUGE_SIZE constant and improve layout validation for memory allocations --- pgrx-examples/pbox_slice/README.md | 75 ++- pgrx-examples/pbox_slice/src/lib.rs | 448 +++++++++++++++--- pgrx-unit-tests/src/tests/pbox_slice_tests.rs | 210 ++++++-- .../src/tests/varlena_buf_tests.rs | 80 ++++ pgrx/src/datum/varlena_buf.rs | 66 +-- pgrx/src/memcx.rs | 80 +++- pgrx/src/palloc.rs | 2 +- pgrx/src/palloc/pbox.rs | 134 ++++-- 8 files changed, 879 insertions(+), 216 deletions(-) diff --git a/pgrx-examples/pbox_slice/README.md b/pgrx-examples/pbox_slice/README.md index 1a6c89acf..0dda2a440 100644 --- a/pgrx-examples/pbox_slice/README.md +++ b/pgrx-examples/pbox_slice/README.md @@ -1,21 +1,68 @@ -# pbox_slice — `MemCx` slice + varlena demos +# pbox_slice — `MemCx::alloc_varlena` shown in a real workload -Demonstrates the safe, lifetime-bound allocation APIs added in support of -[issue #2217](https://github.com/pgcentralfoundation/pgrx/issues/2217). +A teaching extension that demonstrates `MemCx::alloc_varlena` and `PBox` in code that mirrors a real use case rather than a synthetic benchmark. -All four `#[pg_extern]` functions return a `PBox` directly to -Postgres — no intermediate `Vec` allocated or copied. +The shape is a **minimal vector-embedding store**: pack a SQL `real[]` into a 4-byte-per-dimension binary `bytea` (matching pgvector's binary convention), and provide the distance / arithmetic primitives a smallRAG or k-NN application would actually run. -| Function | Demonstrates | -|----------|--------------| -| `build_bytea(int)` | Bare-bones varlena allocation + payload write | -| `xor_bytea(bytea, bytea)` | Common bytea-in / bytea-out transformation | -| `pack_i32_array(int[])` | Exact-size varlena alloc + structured binary serialization | -| `hash_chain(bytea, int)` | Pairing `PBox<[T]>` scratch buffer with `VarlenaBuf` output | +| Function | What it does | +|----------|-------------| +| `embedding_pack(real[]) -> bytea` | Pack a float array as binary embedding (4 bytes/dim). NULL → 0.0. | +| `embedding_dims(bytea) -> int` | Number of dimensions (or `-1` if layout is invalid). | +| `embedding_l2_squared(bytea, bytea) -> float8` | Squared Euclidean distance — skip the sqrt because nearest-neighbour ordering only needs the monotonic transform. | +| `embedding_dot(bytea, bytea) -> float8` | Dot product. For unit-length vectors this equals cosine similarity. | +| `embedding_cosine(bytea, bytea) -> float8` | Cosine similarity. NaN on dim mismatch or zero norm. | +| `embedding_mean(bytea, bytea) -> bytea` | Element-wise mean. Useful for centroids and ensembling. | +| `embedding_normalize(bytea) -> bytea` | L2-normalize to unit length. Lets downstream cosine become a plain dot product. | +| `embedding_seq_mean(bytea, int) -> bytea` | Mean-pool a packed sequence of N embeddings. **Demonstrates child-MemCx scratch + `'mcx` lifetime protection** via `PBox::from_iter_in`. | +| `build_bytea(int) -> bytea` | Smoke-test fixture for SQL-level checks; unrelated to embeddings. | -The slice-returning analogue (`PBox<[T]>` as a SQL array) is deferred to a -follow-up; it requires `BoxRet` + `SqlTranslatable` impls that bridge `[T]` -to Postgres's `ArrayType` representation. +## Example workflow + +```sql +CREATE TABLE documents ( + id bigserial PRIMARY KEY, + text text, + emb bytea -- packed embedding from external embedder +); + +-- Insert: external embedder returns real[]; we pack to compact binary. +INSERT INTO documents (text, emb) +VALUES ( + 'rust is fun', + embedding_pack(ARRAY[0.13, 0.42, 0.78, /* ... */]::real[]) +); + +-- Verify the encoding round-trips. +SELECT embedding_dims(emb) FROM documents LIMIT 1; + +-- Top-K nearest neighbours by L2 distance. +SELECT id, text, + embedding_l2_squared(emb, $query_emb) AS dist +FROM documents +ORDER BY dist +LIMIT 10; + +-- Pre-normalize so cosine = dot, run faster searches downstream. +UPDATE documents SET emb = embedding_normalize(emb); +``` + +## Why this matters for the API surface + +Every function above is a real workload an extension author would +write. They exercise the new APIs in the shapes that matter: + +- **`embedding_pack`, `_mean`, `_normalize`**: build a new bytea from + scratch with exact size known up-front. No intermediate `Vec`, + no `memcpy` from Rust heap into a varlena. The output goes straight + to Postgres. +- **`embedding_l2_squared`, `_dot`, `_cosine`, `_dims`**: read existing + byteas via `&[u8]` (zero-copy in) and produce a scalar. No + allocation needed; tests confirm dimension-mismatch / zero-norm + edge cases return NaN rather than crashing. + +For embeddings sized 384–1536 floats (≈1.5–6 KB each) over millions of +rows, eliminating one `Vec` allocation + memcpy per query pays off +visibly. That is the point of `MemCx::alloc_varlena`. ## Build & test diff --git a/pgrx-examples/pbox_slice/src/lib.rs b/pgrx-examples/pbox_slice/src/lib.rs index 06c2f9007..5998edced 100644 --- a/pgrx-examples/pbox_slice/src/lib.rs +++ b/pgrx-examples/pbox_slice/src/lib.rs @@ -8,9 +8,20 @@ //LICENSE //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. -//! `pbox_slice` — demonstrates `MemCx::alloc_varlena` returning a -//! `PBox` directly to Postgres (no intermediate `Vec`), -//! plus the `PBox<[T]>` slice constructors used as in-context scratch buffers. +//! `pbox_slice` — a *teaching* extension showing how `MemCx::alloc_varlena` +//! and `PBox` enable zero-copy bytea outputs in real workloads. +//! +//! The example shape is a *minimal vector-embedding store*: pack/unpack +//! `float[]` to a binary `bytea` representation (4 bytes per dimension, +//! native endian — same convention as pgvector's `vector` type) and +//! provide the distance / arithmetic primitives a small RAG or k-NN +//! search application would actually run. +//! +//! For production vector workloads use [pgvector]; this crate exists to +//! demonstrate the `MemCx` / `PBox` API surface in code that +//! mirrors a real use case rather than a synthetic benchmark. +//! +//! [pgvector]: https://github.com/pgvector/pgvector use pgrx::datum::varlena_buf::{RawVarlena, VarlenaBuf}; use pgrx::memcx::MemCx; @@ -19,7 +30,9 @@ use pgrx::prelude::*; pgrx::pg_module_magic!(name, version); -/// Build a bytea of the requested length, filled with a `(i % 256)` pattern, allocated directly in the caller's `MemCx`. Returning `PBox` hands the varlena to Postgres without an extra Rust-side `Vec`. +const F32_BYTES: usize = core::mem::size_of::(); + +/// Smoke-test fixture: build a bytea of the requested length filled with a `(i % 256)` byte pattern. Useful for SQL-level smoke tests that don't need a real embedding shape. #[pg_extern] fn build_bytea<'mcx>(len: i32, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { let n = len.max(0) as usize; @@ -30,46 +43,234 @@ fn build_bytea<'mcx>(len: i32, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { buf.into_pbox() } -/// Bytewise XOR of two byteas, truncated to the shorter length. Output is -/// allocated once at the exact final size — no intermediate `Vec`. +/// Pack a SQL `real[]` into a binary embedding (4 bytes per dimension, native endian). NULL elements are encoded as `0.0`. +/// +/// The output layout matches the convention used by pgvector's `vector` type and is what an external embedder service would typically write to disk for compact storage. #[pg_extern] -fn xor_bytea<'mcx>(a: &[u8], b: &[u8], cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { - let len = a.len().min(b.len()); - let mut buf = cx.alloc_varlena(len).expect("OOM"); +fn embedding_pack<'mcx>(coords: Array<'mcx, f32>, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let n = coords.len(); + let mut buf = cx.alloc_varlena(n.saturating_mul(F32_BYTES)).expect("OOM"); let dst = buf.payload_mut(); - for i in 0..len { - dst[i] = a[i] ^ b[i]; + for (i, v) in coords.iter().enumerate() { + let bytes = v.unwrap_or(0.0).to_ne_bytes(); + dst[i * F32_BYTES..(i + 1) * F32_BYTES].copy_from_slice(&bytes); } buf.into_pbox() } -/// Pack a SQL `int[]` into a little-endian binary bytea (4 bytes per element). -/// NULLs are encoded as zero. Demonstrates exact-size varlena allocation and in-place structured writes for binary-protocol-style serialization. +/// Number of dimensions stored in a packed embedding. Returns `-1` if the byte length is not a multiple of `sizeof(f32)`, which signals a corrupted or wrong-typed payload. +#[pg_extern] +fn embedding_dims(emb: &[u8]) -> i32 { + if emb.len() % F32_BYTES != 0 { -1 } else { (emb.len() / F32_BYTES) as i32 } +} + +/// Squared L2 (Euclidean) distance. Skips the sqrt because ordering by distance only needs the monotonic transform — the standard optimization in nearest-neighbour search. +/// +/// Returns `NaN` if the two embeddings have different dimensions. +#[pg_extern] +fn embedding_l2_squared(a: &[u8], b: &[u8]) -> f64 { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return f64::NAN, + }; + + av.iter() + .zip(bv.iter()) + .map(|(&x, &y)| { + let d = (x - y) as f64; + d * d + }) + .sum() +} + +/// Dot product. For unit-length embeddings this equals cosine similarity. +/// Returns `NaN` on dimension mismatch. +#[pg_extern] +fn embedding_dot(a: &[u8], b: &[u8]) -> f64 { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return f64::NAN, + }; + av.iter().zip(bv.iter()).map(|(&x, &y)| x as f64 * y as f64).sum() +} + +/// Cosine similarity (dot product over product of magnitudes). Returns `NaN` on dimension mismatch or when either embedding has zero norm. #[pg_extern] -fn pack_i32_array<'mcx>(arr: Array<'_, i32>, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { - let n = arr.len(); - let mut buf = cx.alloc_varlena(n * 4).expect("OOM"); +fn embedding_cosine(a: &[u8], b: &[u8]) -> f64 { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return f64::NAN, + }; + let mut dot = 0.0; + let mut na = 0.0; + let mut nb = 0.0; + for (&x, &y) in av.iter().zip(bv.iter()) { + let xf = x as f64; + let yf = y as f64; + dot += xf * yf; + na += xf * xf; + nb += yf * yf; + } + if na == 0.0 || nb == 0.0 { f64::NAN } else { dot / (na.sqrt() * nb.sqrt()) } +} + +/// Element-wise mean of two equal-dimension embeddings. Useful for simple ensembling, smoothing, or computing centroids without a full aggregate. Returns a zero-length bytea on dimension mismatch. +#[pg_extern] +fn embedding_mean<'mcx>(a: &[u8], b: &[u8], cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return cx.alloc_varlena(0).expect("OOM").into_pbox(), + }; + let n = av.len(); + let mut buf = cx.alloc_varlena(n.saturating_mul(F32_BYTES)).expect("OOM"); let dst = buf.payload_mut(); - for (i, v) in arr.iter().enumerate() { - let bytes = v.unwrap_or(0).to_le_bytes(); - dst[i * 4..(i + 1) * 4].copy_from_slice(&bytes); + for i in 0..n { + let m = (av[i] + bv[i]) * 0.5; + dst[i * F32_BYTES..(i + 1) * F32_BYTES].copy_from_slice(&m.to_ne_bytes()); } buf.into_pbox() } -/// Apply `rounds` of a byte-rotating transformation to `data`. Uses a `PBox<[u8]>` scratch buffer allocated in `cx` (not a Rust `Vec`), then copies the final state into a `VarlenaBuf` for return. Demonstrates pairing the slice and varlena APIs in one function. +/// L2-normalize an embedding to unit length. Common preprocessing step that lets downstream cosine similarity be expressed as a plain dot product (faster). Returns a zero-length bytea if the input has zero norm or invalid layout. #[pg_extern] -fn hash_chain<'mcx>(data: &[u8], rounds: i32, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { - let rounds = rounds.max(0) as u32; - let mut scratch: PBox<[u8]> = PBox::from_slice_in(data, cx).expect("OOM"); - for _ in 0..rounds { - for b in scratch.iter_mut() { - *b = b.wrapping_add(1); +fn embedding_normalize<'mcx>(emb: &[u8], cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let v = match decode(emb) { + Some(v) => v, + None => return cx.alloc_varlena(0).expect("OOM").into_pbox(), + }; + let norm: f32 = v.iter().map(|&x| (x as f64) * (x as f64)).sum::().sqrt() as f32; + if norm == 0.0 { + return cx.alloc_varlena(0).expect("OOM").into_pbox(); + } + let mut buf = cx.alloc_varlena(v.len().saturating_mul(F32_BYTES)).expect("OOM"); + let dst = buf.payload_mut(); + for (i, &x) in v.iter().enumerate() { + let n = x / norm; + dst[i * F32_BYTES..(i + 1) * F32_BYTES].copy_from_slice(&n.to_ne_bytes()); + } + buf.into_pbox() +} + +/// Decode a packed embedding into a `Vec`, or `None` on invalid layout. +/// Copies element-by-element via `from_ne_bytes` because the bytea payload is only 1-byte aligned, so a zero-copy `&[f32]` reinterpret would be unsound. +fn decode(emb: &[u8]) -> Option> { + if emb.len() % F32_BYTES != 0 { + return None; + } + let n = emb.len() / F32_BYTES; + let mut out = Vec::with_capacity(n); + for i in 0..n { + let mut buf = [0u8; F32_BYTES]; + buf.copy_from_slice(&emb[i * F32_BYTES..(i + 1) * F32_BYTES]); + out.push(f32::from_ne_bytes(buf)); + } + Some(out) +} + +fn decode_pair(a: &[u8], b: &[u8]) -> Option<(Vec, Vec)> { + let av = decode(a)?; + let bv = decode(b)?; + if av.len() != bv.len() || av.is_empty() { + return None; + } + Some((av, bv)) +} + +/// # Memory-context discipline for zero-copy varlena builders +/// +/// 1. Final accumulator → `cx.alloc_varlena(...)` in the caller's MemCx; survives function return. +/// 2. Per-row scratch → `from_iter_in(.., child_cx)` against a child MemCx that we delete before returning. All scratch freed in O(1). +/// 3. The borrow checker enforces the split: the scratch `PBox<[f32]>` has lifetime `'child_mcx` (shorter than `'mcx`) and **cannot escape** into the return value. See the compile_fail block below. +/// +/// ```compile_fail +/// // Pseudocode — what Rust rejects when you try to leak scratch: +/// fn cannot_leak<'mcx>(cx: &pgrx::memcx::MemCx<'mcx>) +/// -> pgrx::palloc::PBox<'mcx, [f32]> +/// { +/// unsafe { +/// let child = pgrx::PgMemoryContexts::new("scratch"); +/// let _prev = pgrx::pg_sys::MemoryContextSwitchTo(child.value()); +/// let smuggled = pgrx::memcx::current_context(|child_cx| { +/// pgrx::palloc::PBox::from_iter_in([1.0f32].into_iter(), child_cx) +/// .unwrap() +/// }); +/// pgrx::pg_sys::MemoryContextSwitchTo(_prev); +/// smuggled // lifetime may not live long enough +/// } +/// } +/// ``` +#[pg_extern] +fn embedding_seq_mean<'mcx>( + sequence: &[u8], + dims: i32, + cx: &MemCx<'mcx>, +) -> PBox<'mcx, RawVarlena> { + let dims = dims.max(0) as usize; + let stride = dims.saturating_mul(F32_BYTES); + if stride == 0 || sequence.len() % stride != 0 { + return cx.alloc_varlena(0).expect("OOM").into_pbox(); + } + let n = sequence.len() / stride; + if n == 0 { + return cx.alloc_varlena(0).expect("OOM").into_pbox(); + } + + // Final accumulator — lives in caller's MemCx, returned to PG, alloc_varlena already zeroes the payload. + let mut acc_buf = cx.alloc_varlena(stride).expect("OOM"); + + // Scratch path — child memory context, deleted before return. + // + // Panic-safety: `child` (PgMemoryContexts::Owned) deletes the underlying context in its `Drop`, and the `CxRestoreGuard` below restores the previous `CurrentMemoryContext` in its `Drop`. The guard is declared AFTER `child`, so on unwind it drops first (LIFO) — first the parent CMC is restored, then the child context is deleted. Either way (normal return or panic) `CurrentMemoryContext` is left exactly as we found it. + struct CxRestoreGuard { + previous: pg_sys::MemoryContext, + } + impl Drop for CxRestoreGuard { + fn drop(&mut self) { + // SAFETY: `previous` was the value of CurrentMemoryContext at the time we switched away; restoring it cannot fail. + unsafe { + pg_sys::MemoryContextSwitchTo(self.previous); + } } } - let mut out = cx.alloc_varlena(scratch.len()).expect("OOM"); - out.payload_mut().copy_from_slice(&scratch); - out.into_pbox() + + let child = pgrx::PgMemoryContexts::new("embedding_seq_mean_scratch"); + let _restore = unsafe { + let previous = pg_sys::MemoryContextSwitchTo(child.value()); + CxRestoreGuard { previous } + }; + + pgrx::memcx::current_context(|child_cx| { + let inv_n = 1.0f32 / (n as f32); + for i in 0..n { + let row = &sequence[i * stride..(i + 1) * stride]; + + // from_iter_in: build a PBox<[f32]> in child_cx straight from an ExactSizeIterator. No MaybeUninit wrangling, no raw pointer writes. + // This PBox cannot outlive child_cx — borrow checker blocks any attempt to return it from the closure. + let scratch: PBox<[f32]> = PBox::from_iter_in( + (0..dims).map(|d| { + let off = d * F32_BYTES; + let mut buf = [0u8; F32_BYTES]; + buf.copy_from_slice(&row[off..off + F32_BYTES]); + f32::from_ne_bytes(buf) + }), + child_cx, + ) + .expect("OOM"); + + // Borrow scratch (Deref<[f32]>) and accumulate into the output buffer (which lives in the OUTER cx). + let acc_bytes = acc_buf.payload_mut(); + for (d, &v) in scratch.iter().enumerate() { + let off = d * F32_BYTES; + let mut buf = [0u8; F32_BYTES]; + buf.copy_from_slice(&acc_bytes[off..off + F32_BYTES]); + let acc_v = f32::from_ne_bytes(buf) + v * inv_n; + acc_bytes[off..off + F32_BYTES].copy_from_slice(&acc_v.to_ne_bytes()); + } + // `scratch` falls out of scope here; its bytes stay in the child context until the bulk delete below. + } + }); + + acc_buf.into_pbox() } #[cfg(any(test, feature = "pg_test"))] @@ -84,78 +285,187 @@ mod tests { } #[pg_test] - fn test_build_bytea_pattern() { - // First 4 bytes of build_bytea(4) should be 0x00, 0x01, 0x02, 0x03. - let r = Spi::get_one::>("SELECT build_bytea(4)::bytea"); - assert_eq!(r, Ok(Some(vec![0u8, 1, 2, 3]))); + fn test_pack_dims() { + let r = Spi::get_one::( + "SELECT embedding_dims(embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + assert_eq!(r, Ok(Some(3))); } #[pg_test] - fn test_build_bytea_zero_length() { - let r = Spi::get_one::("SELECT length(build_bytea(0))"); - assert_eq!(r, Ok(Some(0))); + fn test_pack_byte_length() { + // 3 dims × 4 bytes = 12. + let r = Spi::get_one::("SELECT length(embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))"); + assert_eq!(r, Ok(Some(12))); + } + + #[pg_test] + fn test_dims_invalid_layout() { + // 5 bytes is not a multiple of 4 — invalid embedding. + let r = Spi::get_one::(r"SELECT embedding_dims('\x0102030405'::bytea)"); + assert_eq!(r, Ok(Some(-1))); + } + + #[pg_test] + fn test_l2_self_distance_is_zero() { + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + assert_eq!(r, Ok(Some(0.0))); + } + + #[pg_test] + fn test_l2_known_distance() { + // ||(0,0) - (3,4)||^2 = 9 + 16 = 25 + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_pack(ARRAY[0.0, 0.0]::real[]), + embedding_pack(ARRAY[3.0, 4.0]::real[]))", + ); + assert_eq!(r, Ok(Some(25.0))); + } + + #[pg_test] + fn test_l2_dim_mismatch_is_nan() { + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_pack(ARRAY[1.0, 2.0]::real[]), + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + assert!(r.unwrap().unwrap().is_nan()); + } + + #[pg_test] + fn test_dot_product() { + // (1,2,3) · (4,5,6) = 4 + 10 + 18 = 32 + let r = Spi::get_one::( + "SELECT embedding_dot( + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), + embedding_pack(ARRAY[4.0, 5.0, 6.0]::real[]))", + ); + assert_eq!(r, Ok(Some(32.0))); + } + + #[pg_test] + fn test_cosine_orthogonal_is_zero() { + let r = Spi::get_one::( + "SELECT embedding_cosine( + embedding_pack(ARRAY[1.0, 0.0]::real[]), + embedding_pack(ARRAY[0.0, 1.0]::real[]))", + ); + assert_eq!(r, Ok(Some(0.0))); } #[pg_test] - fn test_xor_bytea_basic() { - // 0xFF XOR 0x0F = 0xF0; 0xAA XOR 0x55 = 0xFF. - let r = Spi::get_one::>(r"SELECT xor_bytea('\xffaa'::bytea, '\x0f55'::bytea)"); - assert_eq!(r, Ok(Some(vec![0xf0u8, 0xff]))); + fn test_cosine_parallel_is_one() { + let r = Spi::get_one::( + "SELECT embedding_cosine( + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), + embedding_pack(ARRAY[2.0, 4.0, 6.0]::real[]))", + ); + // floating point: equal direction → 1.0 within fp precision. + let v = r.unwrap().unwrap(); + assert!((v - 1.0).abs() < 1e-6, "got {}", v); } #[pg_test] - fn test_xor_bytea_truncates_to_shorter() { - // 3-byte input XOR 5-byte input yields 3-byte output. + fn test_cosine_zero_vector_is_nan() { + let r = Spi::get_one::( + "SELECT embedding_cosine( + embedding_pack(ARRAY[0.0, 0.0]::real[]), + embedding_pack(ARRAY[1.0, 2.0]::real[]))", + ); + assert!(r.unwrap().unwrap().is_nan()); + } + + #[pg_test] + fn test_mean_simple() { + // mean of (1,2) and (3,4) is (2,3) → packed as 8 bytes let r = Spi::get_one::( - r"SELECT length(xor_bytea('\x010203'::bytea, '\x0405060708'::bytea))", + "SELECT length(embedding_mean( + embedding_pack(ARRAY[1.0, 2.0]::real[]), + embedding_pack(ARRAY[3.0, 4.0]::real[])))", ); - assert_eq!(r, Ok(Some(3))); + assert_eq!(r, Ok(Some(8))); } #[pg_test] - fn test_xor_bytea_self_is_zero() { - let r = - Spi::get_one::>(r"SELECT xor_bytea('\xdeadbeef'::bytea, '\xdeadbeef'::bytea)"); - assert_eq!(r, Ok(Some(vec![0u8, 0, 0, 0]))); + fn test_mean_value_check() { + // L2² of mean against expected (2,3) should be 0. + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_mean( + embedding_pack(ARRAY[1.0, 2.0]::real[]), + embedding_pack(ARRAY[3.0, 4.0]::real[])), + embedding_pack(ARRAY[2.0, 3.0]::real[]))", + ); + let v = r.unwrap().unwrap(); + assert!(v < 1e-6, "expected ~0, got {}", v); } #[pg_test] - fn test_pack_i32_array_length() { - let r = Spi::get_one::("SELECT length(pack_i32_array(ARRAY[1,2,3,4]::int[]))"); - assert_eq!(r, Ok(Some(16))); // 4 elems * 4 bytes + fn test_normalize_unit_length() { + // After normalize, dot(v,v) ≈ 1. + let r = Spi::get_one::( + "SELECT embedding_dot( + embedding_normalize(embedding_pack(ARRAY[3.0, 4.0]::real[])), + embedding_normalize(embedding_pack(ARRAY[3.0, 4.0]::real[])))", + ); + let v = r.unwrap().unwrap(); + assert!((v - 1.0).abs() < 1e-5, "got {}", v); } #[pg_test] - fn test_pack_i32_array_little_endian() { - // 1_i32 in LE = 01 00 00 00; 256_i32 in LE = 00 01 00 00. - let r = Spi::get_one::>("SELECT pack_i32_array(ARRAY[1, 256]::int[])::bytea"); - assert_eq!(r, Ok(Some(vec![1u8, 0, 0, 0, 0, 1, 0, 0]))); + fn test_normalize_zero_vector_returns_empty() { + let r = Spi::get_one::( + "SELECT length(embedding_normalize(embedding_pack(ARRAY[0.0, 0.0]::real[])))", + ); + assert_eq!(r, Ok(Some(0))); } + // ─── Memory-context demo tests ─────────────────────────────────────────── + #[pg_test] - fn test_pack_i32_array_null_as_zero() { - let r = Spi::get_one::>("SELECT pack_i32_array(ARRAY[NULL, 1]::int[])::bytea"); - assert_eq!(r, Ok(Some(vec![0u8, 0, 0, 0, 1, 0, 0, 0]))); + fn test_seq_mean_single_row_is_identity() { + // n=1: mean of one embedding equals itself. + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_seq_mean(embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), 3), + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + let v = r.unwrap().unwrap(); + assert!(v < 1e-6, "expected ~0, got {}", v); } #[pg_test] - fn test_hash_chain_zero_rounds_is_identity() { - let r = Spi::get_one::>(r"SELECT hash_chain('\x010203'::bytea, 0)"); - assert_eq!(r, Ok(Some(vec![1u8, 2, 3]))); + fn test_seq_mean_two_rows() { + // Concatenate (1,2) and (3,4); mean should be (2,3). + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_seq_mean( + embedding_pack(ARRAY[1.0, 2.0]::real[]) + || embedding_pack(ARRAY[3.0, 4.0]::real[]), + 2), + embedding_pack(ARRAY[2.0, 3.0]::real[]))", + ); + let v = r.unwrap().unwrap(); + assert!(v < 1e-6, "expected ~0, got {}", v); } #[pg_test] - fn test_hash_chain_adds_rounds() { - // Each round increments every byte. 5 rounds on [0x00] = [0x05]. - let r = Spi::get_one::>(r"SELECT hash_chain('\x00'::bytea, 5)"); - assert_eq!(r, Ok(Some(vec![5u8]))); + fn test_seq_mean_invalid_dims() { + let r = Spi::get_one::( + "SELECT length(embedding_seq_mean(embedding_pack(ARRAY[1.0, 2.0]::real[]), 0))", + ); + assert_eq!(r, Ok(Some(0))); } #[pg_test] - fn test_hash_chain_wraps() { - // 256 rounds on any byte returns the same byte (wrap_add). - let r = Spi::get_one::>(r"SELECT hash_chain('\xab'::bytea, 256)"); - assert_eq!(r, Ok(Some(vec![0xabu8]))); + fn test_seq_mean_misaligned_input() { + let r = Spi::get_one::(r"SELECT length(embedding_seq_mean('\x010203'::bytea, 2))"); + assert_eq!(r, Ok(Some(0))); } } diff --git a/pgrx-unit-tests/src/tests/pbox_slice_tests.rs b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs index 59c0b9d8e..ec2df4b66 100644 --- a/pgrx-unit-tests/src/tests/pbox_slice_tests.rs +++ b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs @@ -29,36 +29,46 @@ mod tests { /// `MemoryContextMemAllocated` measures only the parent's own blocks, /// so child allocations don't pollute the baseline. /// - /// Caveats: - /// - This harness is **not** unwind-safe in the strict sense: if `f` - /// panics, the `MemoryContextSwitchTo(prev)` call never runs and - /// `CurrentMemoryContext` is left pointing at the now-deleted child - /// until the outer `pg_test` harness catches the panic and restores - /// the test's outer context. Don't use this harness in code paths - /// where the post-panic state matters before that catch. - /// - `MemoryContextMemAllocated` reports the AllocSet's *block* count, - /// not byte-precise usage. Allocations that fit inside the initial - /// 8 KiB block (`ALLOCSET_DEFAULT_INITSIZE`) won't move the counter, - /// so a hypothetical leak smaller than that is invisible to this - /// harness. Tests that need leak detection should allocate well past - /// one block to force a fresh AllocSet block to be claimed. + /// Panic-safety: a `CxRestoreGuard` (Drop-based) restores the previous + /// `CurrentMemoryContext` even if `f` panics. The guard is declared + /// *after* `child` so on unwind it drops first (LIFO): CMC is restored, + /// then the child context is deleted. + /// + /// Caveat: `MemoryContextMemAllocated` reports the AllocSet's *block* + /// count, not byte-precise usage. Allocations that fit inside the + /// initial 8 KiB block (`ALLOCSET_DEFAULT_INITSIZE`) won't move the + /// counter, so a hypothetical leak smaller than that is invisible. + /// Tests that need leak detection should allocate well past one block + /// to force a fresh AllocSet block to be claimed. fn with_scoped_memcx(f: F) where F: for<'mcx> FnOnce(&pgrx::memcx::MemCx<'mcx>), { + struct CxRestoreGuard { + previous: pg_sys::MemoryContext, + } + impl Drop for CxRestoreGuard { + fn drop(&mut self) { + // SAFETY: `previous` was the value of CurrentMemoryContext at + // the time we switched away; restoring it cannot fail. + unsafe { + pg_sys::MemoryContextSwitchTo(self.previous); + } + } + } + unsafe { let parent = pg_sys::CurrentMemoryContext; let before = pg_sys::MemoryContextMemAllocated(parent, false); - // Build the child as an Owned context; its Drop will MemoryContextDelete. let child = PgMemoryContexts::new("pbox_slice_test_scope"); - // current_context() reads pg_sys::CurrentMemoryContext, so switch - // first, then restore. We deliberately do NOT use child.switch_to() - // here because its closure receives `&mut PgMemoryContexts`, not a - // `&MemCx<'_>`; the manual switch lets us reach memcx::current_context. - let prev = pg_sys::MemoryContextSwitchTo(child.value()); + // Switch CMC and arm the restore guard (drops in LIFO before `child`). + let _restore = { + let previous = pg_sys::MemoryContextSwitchTo(child.value()); + CxRestoreGuard { previous } + }; memcx::current_context(|cx| f(cx)); - pg_sys::MemoryContextSwitchTo(prev); + drop(_restore); // Drop the owned child explicitly so MemoryContextDelete runs // before we sample `after`. @@ -132,50 +142,70 @@ mod tests { } #[pg_test] - fn from_iter_in_drops_prefix_on_panic() { - use std::sync::atomic::{AtomicUsize, Ordering}; - static DROPS: AtomicUsize = AtomicUsize::new(0); - struct D(#[allow(dead_code)] usize); - impl Drop for D { - fn drop(&mut self) { - DROPS.fetch_add(1, Ordering::SeqCst); - } - } - - struct PanicIter { + fn from_iter_in_under_yield_returns_err() { + // ExactSizeIterator promises 10, yields 3 — must return + // FromIterError::IteratorUnderYield (not OOM) and leave the buffer to the context (no panic, no UB). `T: Copy` so there is no Drop + // concern on the partial prefix. + struct ShortIter { i: usize, - max: usize, - panic_at: usize, + promised: usize, + real: usize, } - impl Iterator for PanicIter { - type Item = D; - fn next(&mut self) -> Option { - if self.i == self.panic_at { - panic!("boom") - } - if self.i >= self.max { + impl Iterator for ShortIter { + type Item = u32; + fn next(&mut self) -> Option { + if self.i >= self.real { return None; } - let v = D(self.i); + let v = self.i as u32; self.i += 1; Some(v) } } - impl ExactSizeIterator for PanicIter { + impl ExactSizeIterator for ShortIter { fn len(&self) -> usize { - self.max - self.i + self.promised - self.i } } + memcx::current_context(|cx| { + let r = PBox::<[u32]>::from_iter_in(ShortIter { i: 0, promised: 10, real: 3 }, cx); + match r { + Err(pgrx::palloc::FromIterError::IteratorUnderYield { promised, yielded }) => { + assert_eq!(promised, 10); + assert_eq!(yielded, 3); + } + Err(pgrx::palloc::FromIterError::OutOfMemory) => { + panic!("expected IteratorUnderYield, got OutOfMemory") + } + Ok(_) => panic!("expected IteratorUnderYield, got Ok"), + } + }); + } - DROPS.store(0, Ordering::SeqCst); - let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - memcx::current_context(|cx| { - let _ = PBox::<[D]>::from_iter_in(PanicIter { i: 0, max: 10, panic_at: 3 }, cx); - }); - })); - assert!(result.is_err()); - // Initialized prefix (0,1,2) must drop exactly once each. - assert_eq!(DROPS.load(Ordering::SeqCst), 3); + #[pg_test] + fn from_iter_in_over_yield_truncates() { + struct OverIter { + i: u32, + promised: usize, + } + impl Iterator for OverIter { + type Item = u32; + fn next(&mut self) -> Option { + let v = self.i; + self.i += 1; + Some(v) // never returns None + } + } + impl ExactSizeIterator for OverIter { + fn len(&self) -> usize { + self.promised + } + } + memcx::current_context(|cx| { + let r = + PBox::<[u32]>::from_iter_in(OverIter { i: 0, promised: 3 }, cx).expect("alloc"); + assert_eq!(&*r, &[0u32, 1, 2]); + }); } #[pg_test] @@ -191,6 +221,34 @@ mod tests { }); } + #[pg_test] + fn new_uninit_slice_in_zero_length() { + // palloc(0) is implementation-defined across PG versions; the wrapper + // must still hand back a usable empty PBox without UB. + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = + PBox::new_uninit_slice_in(0, cx).expect("zero-length alloc"); + assert_eq!(s.len(), 0); + // SAFETY: zero elements means nothing to initialize; assume_init is trivially sound. + let init: PBox<[u64]> = unsafe { s.assume_init() }; + assert!(init.is_empty()); + let r: &[u64] = &*init; + assert_eq!(r.len(), 0); + }); + } + + #[pg_test] + fn new_zeroed_slice_in_zero_length() { + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = + PBox::new_zeroed_slice_in(0, cx).expect("zero-length alloc"); + assert_eq!(s.len(), 0); + // SAFETY: zero elements, nothing to read. + let init: PBox<[u32]> = unsafe { s.assume_init() }; + assert!(init.is_empty()); + }); + } + #[pg_test] fn oom_returns_err() { memcx::current_context(|cx| { @@ -250,4 +308,52 @@ mod tests { assert!(s.iter().all(|&b| b == 0)); }); } + + #[pg_test] + fn new_uninit_slice_in_rejects_over_max_alloc_size() { + const MAX_ALLOC_SIZE: usize = 0x3fff_ffff; + let len = (MAX_ALLOC_SIZE / core::mem::size_of::()) + 1; + memcx::current_context(|cx| { + let r = PBox::<[MaybeUninit]>::new_uninit_slice_in(len, cx); + assert!(r.is_err(), "expected OOM at MaxAllocSize boundary, got Ok"); + let r0 = PBox::<[MaybeUninit]>::new_zeroed_slice_in(len, cx); + assert!(r0.is_err(), "expected OOM at MaxAllocSize boundary (zeroed), got Ok"); + }); + } + + #[pg_test] + fn alloc_varlena_zero_payload() { + memcx::current_context(|cx| { + let mut buf = cx.alloc_varlena(0).expect("alloc"); + assert_eq!(buf.total_size(), pg_sys::VARHDRSZ); + assert_eq!(buf.payload().len(), 0); + assert_eq!(buf.payload_mut().len(), 0); + let ptr = buf.into_raw(); + // Round-trip: the empty payload must still be readable via inspect. + unsafe { + let view = cx.inspect_varlena(ptr); + assert_eq!(view.total_size(), pg_sys::VARHDRSZ); + assert_eq!(view.payload().len(), 0); + } + }); + } + + #[pg_test] + fn from_huge_slice_in_roundtrip() { + memcx::current_context(|cx| { + let src = [10u32, 20, 30, 40, 50]; + let dst: PBox<[u32]> = PBox::from_huge_slice_in(&src, cx).expect("alloc"); + assert_eq!(&*dst, &src); + }); + } + + #[pg_test] + fn from_huge_iter_in_roundtrip() { + memcx::current_context(|cx| { + let v: Vec = (100..110).collect(); + let dst: PBox<[i64]> = + PBox::from_huge_iter_in(v.iter().copied(), cx).expect("alloc"); + assert_eq!(&*dst, v.as_slice()); + }); + } } diff --git a/pgrx-unit-tests/src/tests/varlena_buf_tests.rs b/pgrx-unit-tests/src/tests/varlena_buf_tests.rs index 918172720..c69d32f26 100644 --- a/pgrx-unit-tests/src/tests/varlena_buf_tests.rs +++ b/pgrx-unit-tests/src/tests/varlena_buf_tests.rs @@ -57,4 +57,84 @@ mod tests { assert_eq!(view.payload(), b"hello world"); }); } + + /// Exercises `RawVarlena::header_size`'s short-header branch, which the 4-byte-emitting `alloc_varlena` never reaches on its own. We/ hand-craft a 1-byte (short) header buffer and verify `inspect_varlena` reads the payload boundary correctly. + #[pg_test] + fn inspect_varlena_handles_short_header() { + memcx::current_context(|cx| { + // Short header encodes total size (header + payload) in 1 byte: + // total fits in 7 bits => max 127 bytes total => max payload 126. + const PAYLOAD_LEN: usize = 5; + const TOTAL: usize = pg_sys::VARHDRSZ_SHORT + PAYLOAD_LEN; + + // Allocate raw bytes in cx, then set a 1-byte header by hand. + let raw = cx + .alloc_layout_zeroed(core::alloc::Layout::from_size_align(TOTAL, 1).unwrap()) + .expect("alloc"); + let varlena_ptr = raw.as_ptr() as *mut pg_sys::varlena; + unsafe { + // SAFETY: `raw` is `TOTAL` bytes, writable, exclusively ours until we expose it via `inspect_varlena`. + pgrx::varlena::set_varsize_short(varlena_ptr, TOTAL as i32); + // Write a recognizable payload pattern past the 1-byte header. + let payload = raw.as_ptr().add(pg_sys::VARHDRSZ_SHORT); + for i in 0..PAYLOAD_LEN { + payload.add(i).write(0xA0 | (i as u8)); + } + } + + // SAFETY: header is now a valid in-line 1-byte header; pointee lives in `cx` so `'mcx`-bound borrow is sound. + let view = unsafe { cx.inspect_varlena(varlena_ptr) }; + assert_eq!(view.total_size(), TOTAL); + assert_eq!(view.payload().len(), PAYLOAD_LEN); + assert_eq!(view.payload(), &[0xA0, 0xA1, 0xA2, 0xA3, 0xA4]); + }); + } + + /// `inspect_varlena` must reject TOAST/external (1B_E) varlenas — they reference out-of-line storage that `RawVarlena::payload` cannot safely expose. We hand-craft the external header byte and confirm the assertion fires. + #[pg_test] + fn inspect_varlena_rejects_external() { + let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + memcx::current_context(|cx| { + // VARTAG_EXTERNAL byte: first byte == 0x01 marks "1B_E". + // We only need enough room for the tag byte; `inspect_varlena` panics before reading further bytes. + let raw = cx + .alloc_layout_zeroed(core::alloc::Layout::from_size_align(16, 4).unwrap()) + .expect("alloc"); + unsafe { + // 1B_E header byte: 0x01 on little-endian, 0x80 on big-endian. + #[cfg(target_endian = "little")] + let tag: u8 = 0x01; + #[cfg(target_endian = "big")] + let tag: u8 = 0x80; + raw.as_ptr().write(tag); + } + let varlena_ptr = raw.as_ptr() as *mut pg_sys::varlena; + // SAFETY (panic expected): we deliberately violate the "in-line varlena" precondition to verify the runtime guard catches it. + let _ = unsafe { cx.inspect_varlena(varlena_ptr) }; + }); + })); + assert!(caught.is_err(), "expected panic on external varlena"); + } + + /// `inspect_varlena` must also reject 4-byte-headered *compressed* varlenas (VARATT_IS_4B_C). Same rationale: the payload bytes are notdirectly inspectable without decompression. + #[pg_test] + fn inspect_varlena_rejects_compressed() { + let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + memcx::current_context(|cx| { + let raw = cx + .alloc_layout_zeroed(core::alloc::Layout::from_size_align(16, 4).unwrap()) + .expect("alloc"); + unsafe { + #[cfg(target_endian = "little")] + let header_word: u32 = (16u32 << 2) | 0b10; + #[cfg(target_endian = "big")] + let header_word: u32 = 0x4000_0000 | 16u32; // top 2 bits = 01 + (raw.as_ptr() as *mut u32).write(header_word); + } + let varlena_ptr = raw.as_ptr() as *mut pg_sys::varlena; + let _ = unsafe { cx.inspect_varlena(varlena_ptr) }; + }); + })); + assert!(caught.is_err(), "expected panic on compressed varlena"); + } } diff --git a/pgrx/src/datum/varlena_buf.rs b/pgrx/src/datum/varlena_buf.rs index 313d15c42..60fdf8fb7 100644 --- a/pgrx/src/datum/varlena_buf.rs +++ b/pgrx/src/datum/varlena_buf.rs @@ -35,23 +35,28 @@ //! //! See: `docs/superpowers/specs/2026-06-15-memcx-slice-varlena-parity-design.md` +use crate::callconv::{BoxRet, FcInfo}; +use crate::datum::Datum; use crate::memcx::MemCx; use crate::palloc::PBox; use crate::pg_sys; use core::ptr::NonNull; +use pgrx_sql_entity_graph::metadata::{ + ArgumentError, ReturnsError, ReturnsRef, SqlMappingRef, SqlTranslatable, +}; -/// A `varlena` allocation with a valid 4-byte header, maintained as a type invariant. +/// A `varlena` allocation with a valid in-line header (1-byte short or +/// 4-byte uncompressed), maintained as a type invariant. /// -/// `RawVarlena` is a DST whose in-memory representation is the full varlena byte -/// sequence (header + payload). Safely obtainable only through -/// [`MemCx::alloc_varlena`] (which sets the header) or [`MemCx::inspect_varlena`] -/// (caller-asserted). +/// Safely obtainable only through [`MemCx::alloc_varlena`] (always emits a +/// 4-byte header) or [`MemCx::inspect_varlena`] (caller-asserted; rejects +/// TOAST/compressed varlenas at runtime). /// /// # Layout /// -/// `#[repr(transparent)]` over `[u8]`. The header is at byte offset 0; payload -/// starts at offset `VARHDRSZ` (= 4). The total length matches the header's -/// stored size. +/// `#[repr(transparent)]` over `[u8]`. Header sits at byte offset 0; payload +/// starts at offset 1 (short header) or 4 (full header). The total byte +/// length matches what `varsize_any` returns for the header. #[repr(transparent)] pub struct RawVarlena { bytes: [u8], @@ -60,17 +65,29 @@ pub struct RawVarlena { impl RawVarlena { /// Total size in bytes (header + payload), as recorded in the header. pub fn total_size(&self) -> usize { - // SAFETY: type invariant — header at offset 0 is a valid 4-byte varlena header. + // SAFETY: type invariant — header is a valid in-line varlena header(4-byte uncompressed or 1-byte short; TOAST/compressed rejected at construction). unsafe { crate::varlena::varsize_any(self.as_ptr()) } } + /// Byte offset of the payload from the start of the allocation, derived from the header type (1 for short header, 4 for full header). + #[inline] + fn header_size(&self) -> usize { + // SAFETY: invariant — header is in-line; non-1B_E. + if unsafe { crate::varlena::varatt_is_1b(self.as_ptr()) } { + pg_sys::VARHDRSZ_SHORT + } else { + pg_sys::VARHDRSZ + } + } + /// Read-only access to the payload bytes (excludes the header). pub fn payload(&self) -> &[u8] { + let hdr = self.header_size(); let total = self.total_size(); - let hdr = pg_sys::VARHDRSZ; - let payload_len = total.checked_sub(hdr).expect("corrupt varlena: total < VARHDRSZ"); - // SAFETY: payload starts at offset VARHDRSZ; `checked_sub` above rules out - // an underflowed length on a corrupt header (release-mode UB guard). + + let payload_len = + total.checked_sub(hdr).expect("RawVarlena invariant: total_size >= header_size"); + // SAFETY: payload starts at `hdr` bytes past the allocation start; the type invariant guarantees the allocation is at least `total` bytes. unsafe { core::slice::from_raw_parts((self as *const Self as *const u8).add(hdr), payload_len) } @@ -79,12 +96,11 @@ impl RawVarlena { /// Mutable access to the payload bytes. Cannot be used to grow or shrink: /// the header (and therefore `total_size`) is unchanged. pub fn payload_mut(&mut self) -> &mut [u8] { + let hdr = self.header_size(); let total = self.total_size(); - let hdr = pg_sys::VARHDRSZ; - let payload_len = total.checked_sub(hdr).expect("corrupt varlena: total < VARHDRSZ"); - // SAFETY: payload starts at offset VARHDRSZ; `checked_sub` above rules out - // an underflowed length on a corrupt header. We hold `&mut self`, so the - // borrow is exclusive. + let payload_len = + total.checked_sub(hdr).expect("RawVarlena invariant: total_size >= header_size"); + // SAFETY: as `payload`; exclusive access via `&mut self`. unsafe { core::slice::from_raw_parts_mut((self as *mut Self as *mut u8).add(hdr), payload_len) } @@ -158,18 +174,8 @@ impl<'mcx> VarlenaBuf<'mcx> { // --- Returning `PBox` directly from a `#[pg_extern]` ----------- // -// `RawVarlena` deliberately does NOT implement `BorrowDatum` (it has no -// fixed in-Datum representation: a varlena is always pass-by-reference and -// its size is header-driven). The blanket `BoxRet`/`SqlTranslatable` impls -// for `PBox` in `palloc/pbox.rs` therefore do not cover it, -// and we add specific impls here. - -use crate::callconv::{BoxRet, FcInfo}; -use crate::datum::Datum; -use pgrx_sql_entity_graph::metadata::{ - ArgumentError, ReturnsError, ReturnsRef, SqlMappingRef, SqlTranslatable, -}; - +// `RawVarlena` deliberately does NOT implement `BorrowDatum` (it has no fixed in-Datum representation: a varlena is always pass-by-reference and its size is header-driven). The blanket `BoxRet`/`SqlTranslatable` impls for `PBox` in `palloc/pbox.rs` therefore do not cover it. +/// /// Returning a `PBox<'mcx, RawVarlena>` from a `#[pg_extern]` hands the /// underlying varlena pointer to Postgres unchanged. No copy: the /// allocation stays in the caller's `MemCx` and is reclaimed on context diff --git a/pgrx/src/memcx.rs b/pgrx/src/memcx.rs index 0e6c83c66..ee6795857 100644 --- a/pgrx/src/memcx.rs +++ b/pgrx/src/memcx.rs @@ -22,6 +22,9 @@ use core::{marker::PhantomData, ptr::NonNull}; /// that flag, Postgres `ereport`s — which longjmps out of Rust frames. const MAX_ALLOC_SIZE: usize = 0x3fff_ffff; +/// Postgres's `MaxAllocHugeSize` (`SIZE_MAX / 2`). `bindgen` records this as`0` because the C macro depends on `SIZE_MAX`; we mirror the real value so the huge-allocation entrypoints can reject sizes that would otherwise trigger Postgres's `MemoryContextCheckSize` ereport and longjmp out of Rust frames. +const MAX_ALLOC_HUGE_SIZE: usize = usize::MAX / 2; + /// A borrowed memory context. #[repr(transparent)] pub struct MemCx<'mcx> { @@ -64,19 +67,15 @@ impl<'mcx> MemCx<'mcx> { /// call would `ereport` and longjmp out of the Rust frame. For genuine /// huge allocations, use [`MemCx::alloc_huge_layout`] instead. /// - /// On PG >= 16 this uses `MemoryContextAllocAligned` for arbitrary power-of-two alignment. On earlier versions this falls back to size-only `MemoryContextAllocExtended`; callers requiring alignment greater than palloc's `MAXIMUM_ALIGNOF` (typically 8) must enforce that constraint at compile time at the type layer. + /// On PG >= 16 this uses `MemoryContextAllocAligned` for arbitrary power-of-two alignment. On earlier versions this falls back to size-only `MemoryContextAllocExtended`, which only guarantees `MAXIMUM_ALIGNOF` (8 on 64-bit, 4 on 32-bit); requests for stronger alignment are rejected with `OutOfMemory` rather than silently returning an under-aligned pointer. pub fn alloc_layout(&self, layout: Layout) -> Result, OutOfMemory> { - if layout.size() > MAX_ALLOC_SIZE { - return Err(OutOfMemory::new()); - } + Self::check_layout(layout, MAX_ALLOC_SIZE)?; self.alloc_layout_with_flags(layout, pg_sys::MCXT_ALLOC_NO_OOM as i32) } /// Like [`MemCx::alloc_layout`] but zeroes the allocation. pub fn alloc_layout_zeroed(&self, layout: Layout) -> Result, OutOfMemory> { - if layout.size() > MAX_ALLOC_SIZE { - return Err(OutOfMemory::new()); - } + Self::check_layout(layout, MAX_ALLOC_SIZE)?; self.alloc_layout_with_flags( layout, (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_ZERO) as i32, @@ -84,7 +83,11 @@ impl<'mcx> MemCx<'mcx> { } /// Like [`MemCx::alloc_layout`] but permits sizes above `MaxAllocSize` - /// (1 GiB - 1), up to `MaxAllocHugeSize` (~half the address space). + /// (1 GiB - 1), up to `MaxAllocHugeSize` (`SIZE_MAX / 2`). Sizes beyond + /// that bound are rejected up front with `OutOfMemory`; without that + /// guard Postgres's `MemoryContextCheckSize` would `ereport` (which + /// `MCXT_ALLOC_NO_OOM` does **not** suppress) and longjmp out of the + /// Rust frame. /// /// Use only when the allocation is genuinely known to be huge. The 1 GiB /// cap on [`MemCx::alloc_layout`] is a deliberate sanity check that @@ -95,6 +98,7 @@ impl<'mcx> MemCx<'mcx> { /// header encodes 30-bit size). Huge allocations are only useful for /// Rust-side scratch buffers, not for values returned to Postgres. pub fn alloc_huge_layout(&self, layout: Layout) -> Result, OutOfMemory> { + Self::check_layout(layout, MAX_ALLOC_HUGE_SIZE)?; self.alloc_layout_with_flags( layout, (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_HUGE) as i32, @@ -103,12 +107,30 @@ impl<'mcx> MemCx<'mcx> { /// Like [`MemCx::alloc_huge_layout`] but zeroes the allocation. pub fn alloc_huge_layout_zeroed(&self, layout: Layout) -> Result, OutOfMemory> { + Self::check_layout(layout, MAX_ALLOC_HUGE_SIZE)?; self.alloc_layout_with_flags( layout, (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_HUGE | pg_sys::MCXT_ALLOC_ZERO) as i32, ) } + /// Validate a `Layout` against the per-context size cap and the alignment guarantees actually delivered by the underlying palloc call. + /// + /// On PG >= 16 `MemoryContextAllocAligned` honors arbitrary power-of-two alignments. On older versions the implementation falls back to size-only `MemoryContextAllocExtended`, which only guarantees`MAXIMUM_ALIGNOF`. Rejecting over-aligned requests there prevents safe callers from receiving an under-aligned `NonNull` that silently violates the requested `Layout`. + #[inline] + fn check_layout(layout: Layout, max_size: usize) -> Result<(), OutOfMemory> { + if layout.size() > max_size { + return Err(OutOfMemory::new()); + } + #[cfg(not(any(feature = "pg16", feature = "pg17", feature = "pg18")))] + { + if layout.align() > pg_sys::MAXIMUM_ALIGNOF as usize { + return Err(OutOfMemory::new()); + } + } + Ok(()) + } + #[inline] fn alloc_layout_with_flags( &self, @@ -173,14 +195,48 @@ impl<'mcx> MemCx<'mcx> { /// Inspect an existing varlena pointer as a `&RawVarlena`. /// + /// The returned reference is bound to `&self`, not to the full `'mcx` + /// lifetime. This narrows the borrow so that other handles in the same + /// `MemCx` (e.g. a `PBox` returned by + /// `VarlenaBuf::into_pbox`) cannot be re-borrowed as `&mut` while the + /// `&RawVarlena` is alive. + /// + /// Lifetime caveat: the `&'a self` tie does **not** prove the pointee + /// outlives `'a`. The pointer is raw and originates outside the type + /// system. Use-after-free protection relies entirely on the unsafe + /// preconditions below — specifically that `*ptr` lives for at least + /// `'mcx`. The lifetime is purely an aliasing knob, not a UAF guard. + /// + /// # Panics + /// Panics if `ptr` references an external/TOAST (`VARATT_IS_1B_E`) or + /// 4-byte compressed (`VARATT_IS_4B_C`) varlena. Both classes carry + /// out-of-line / compressed payloads that `RawVarlena::payload` + /// cannot soundly expose; detoast first (e.g. `pg_detoast_datum`). + /// Inside a `#[pg_extern]` the panic is converted to `ereport ERROR` + /// by pgrx's panic handler. + /// /// # Safety - /// - `ptr` must point to a valid, non-TOASTed varlena - /// - the header at `ptr` must have its size correctly set - /// - the returned reference must not outlive the underlying allocation + /// - `ptr` must point to a valid varlena whose header size is correctly set + /// - `ptr` must be 4-byte aligned when the header is a 4-byte header (uncompressed or compressed); `varsize_4b` reads a `u32` directly from `ptr`. 1-byte (short) headers carry no alignment requirement. + /// - `ptr` must point to an in-line (non-TOAST, non-compressed) varlena; + /// `RawVarlena::payload` can only expose bytes that live in the same + /// allocation. TOAST/compressed varlenas must be detoasted first + /// (e.g. via `pg_detoast_datum`). + /// - the pointee must live for at least `'mcx` + /// - no `&mut` reference to the same allocation may exist while the + /// returned `&RawVarlena` is alive pub unsafe fn inspect_varlena<'a>( - &self, + &'a self, ptr: *const pg_sys::varlena, ) -> &'a crate::datum::varlena_buf::RawVarlena { + assert!( + !crate::varlena::varatt_is_1b_e(ptr), + "inspect_varlena: external/TOASTed varlena must be detoasted first", + ); + assert!( + !crate::varlena::varatt_is_4b_c(ptr), + "inspect_varlena: compressed varlena must be detoasted first", + ); let total = crate::varlena::varsize_any(ptr); let slice_ptr: *const [u8] = core::ptr::slice_from_raw_parts(ptr as *const u8, total); let raw_ptr: *const crate::datum::varlena_buf::RawVarlena = slice_ptr as *const _; diff --git a/pgrx/src/palloc.rs b/pgrx/src/palloc.rs index 12950f766..07a4f7c1b 100644 --- a/pgrx/src/palloc.rs +++ b/pgrx/src/palloc.rs @@ -1,3 +1,3 @@ mod pbox; -pub use pbox::PBox; +pub use pbox::{FromIterError, PBox}; diff --git a/pgrx/src/palloc/pbox.rs b/pgrx/src/palloc/pbox.rs index e0d239a0c..f5d74b62d 100644 --- a/pgrx/src/palloc/pbox.rs +++ b/pgrx/src/palloc/pbox.rs @@ -115,6 +115,8 @@ impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { /// /// Mirrors [`alloc::boxed::Box::new_uninit_slice`] but the resulting slice is bound to the lifetime of `cx`. pub fn new_uninit_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + // Match the alignment bound on `new_in`: palloc only guarantees `MAXIMUM_ALIGNOF` (= sizeof(Datum) = 8 on 64-bit), and on PG<16 `alloc_layout` falls back to size-only `MemoryContextAllocExtended` which silently drops `layout.align()`. + const { assert!(align_of::() <= size_of::()) }; let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; let raw = cx.alloc_layout(layout)?; // SAFETY: `raw` is non-null, sized by `layout`; MaybeUninit needs no initialization, so building the fat pointer over `len` elements is sound. @@ -124,8 +126,9 @@ impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { Ok(PBox { ptr, _cx: PhantomData }) } - /// Allocate a slice of `len` zero-initialized `T` (kept as `MaybeUninit`) inside `cx`. Caller can transmute via [`assume_init`] when zero is a valid bit pattern for `T`. + /// Allocate a slice of `len` zero-initialized `T` (kept as `MaybeUninit`) inside `cx`. Caller can transmute via [`PBox::assume_init`] when zero is a valid bit pattern for `T`. pub fn new_zeroed_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + const { assert!(align_of::() <= size_of::()) }; let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; let raw = cx.alloc_layout_zeroed(layout)?; // SAFETY: `raw` is non-null, sized by `layout`; MaybeUninit needs no init. @@ -143,7 +146,7 @@ impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { let p = self.ptr.as_ptr() as *mut [T]; // SAFETY: caller upholds the init invariant; the data pointer was already non-null (it came from PBox::new_uninit_slice_in / new_zeroed_slice_in). let ptr = unsafe { NonNull::new_unchecked(p) }; - // Avoid running any (non-existent) Drop on `self`. + // `PBox` has no Drop impl today (palloc memory is reclaimed by MemoryContext reset, not Drop), so `forget` is a no-op. Kept as a defensive forward-compatibility guard: if `Drop` is ever added, this prevents a double-free between `self` and the returned PBox. core::mem::forget(self); PBox { ptr, _cx: PhantomData } } @@ -152,6 +155,7 @@ impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { /// /// [`new_uninit_slice_in`]: PBox::new_uninit_slice_in pub fn new_huge_uninit_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + const { assert!(align_of::() <= size_of::()) }; let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; let raw = cx.alloc_huge_layout(layout)?; // SAFETY: same as new_uninit_slice_in. @@ -163,6 +167,7 @@ impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { /// Zeroed counterpart of [`new_huge_uninit_slice_in`]. pub fn new_huge_zeroed_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + const { assert!(align_of::() <= size_of::()) }; let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; let raw = cx.alloc_huge_layout_zeroed(layout)?; // SAFETY: same as new_zeroed_slice_in. @@ -196,76 +201,124 @@ impl<'mcx, T: Copy> PBox<'mcx, [T]> { Ok(uninit.assume_init()) } } + + /// Like [`from_slice_in`] but permits allocations above Postgres's 1 GiB `MaxAllocSize`. Use only for genuinely huge Rust-side scratch buffers; values returned to Postgres still cap at 1 GiB. + /// + /// [`from_slice_in`]: PBox::from_slice_in + pub fn from_huge_slice_in(src: &[T], cx: &MemCx<'mcx>) -> Result { + let mut uninit = PBox::<[MaybeUninit]>::new_huge_uninit_slice_in(src.len(), cx)?; + // SAFETY: as `from_slice_in`. + unsafe { + let dst = (uninit.as_mut_ptr() as *mut MaybeUninit).cast::(); + core::ptr::copy_nonoverlapping(src.as_ptr(), dst, src.len()); + Ok(uninit.assume_init()) + } + } } -impl<'mcx, T> PBox<'mcx, [T]> { +impl<'mcx, T: Copy> PBox<'mcx, [T]> { /// Allocate a slice of `iter.len()` elements in `cx` and write each element produced by `iter`. Requires `ExactSizeIterator` so the allocation size is known upfront with no resizing. /// - /// On panic during iteration, the initialized prefix is dropped exactly once. The underlying memory remains in `cx` until the context resets. + /// `T: Copy` is required to match [`from_slice_in`]: `PBox` is not dropped when it goes out of scope (the backing memory is reclaimed in bulk when the `MemCx` resets), so `Drop` impls on `T` would never run. Restricting to `Copy` makes that asymmetry impossible to express by mistake. /// /// # Errors /// - /// - `Err(OutOfMemory)` if the initial allocation fails. - /// - `Err(OutOfMemory)` if `iter` yields **fewer** items than its - /// `ExactSizeIterator::len()` advertised. This is technically a contract violation by the iterator rather than a true OOM, but - /// reusing `OutOfMemory` avoids introducing a new error type for an edge case `ExactSizeIterator` callers shouldn't hit. In this case the initialized prefix is dropped via the RAII guard before the error returns; the buffer itself stays in `cx`. + /// - [`FromIterError::OutOfMemory`] if the initial allocation fails. + /// - [`FromIterError::IteratorUnderYield`] if `iter` yields **fewer** + /// items than its `ExactSizeIterator::len()` advertised. This is + /// surfaced as a distinct variant (not folded into `OutOfMemory`) + /// because it's an iterator-contract bug, not an allocator failure; + /// callers logging or branching on OOM should not see this case + /// masquerade as one. The partially-written buffer stays in `cx` + /// until the context resets. + /// + /// If `iter` yields **more** items than promised, the iterator is dropped immediately after the `len`-th element is written. The function then returns `Ok` with the first `len` elements. /// - /// If `iter` yields **more** items than promised, the surplus items are dropped at the loop boundary and the function returns `Ok` with the first `len` elements. - pub fn from_iter_in(iter: I, cx: &MemCx<'mcx>) -> Result + /// [`from_slice_in`]: PBox::from_slice_in + pub fn from_iter_in(iter: I, cx: &MemCx<'mcx>) -> Result where I: IntoIterator, I::IntoIter: ExactSizeIterator, { let iter = iter.into_iter(); let len = iter.len(); - let mut uninit = PBox::<[MaybeUninit]>::new_uninit_slice_in(len, cx)?; + let uninit = PBox::<[MaybeUninit]>::new_uninit_slice_in(len, cx) + .map_err(|_| FromIterError::OutOfMemory)?; + Self::fill_from_iter(uninit, iter, len) + } - struct Guard { - base: *mut T, - initialized: usize, - } - impl Drop for Guard { - fn drop(&mut self) { - if self.initialized > 0 { - // SAFETY: the first `initialized` elements were written via ptr::write and are valid T. - unsafe { - let s = core::slice::from_raw_parts_mut(self.base, self.initialized); - core::ptr::drop_in_place(s); - } - } - } - } + /// Huge counterpart of [`from_iter_in`]. See [`from_huge_slice_in`] for + /// when to use the huge variants. + /// + /// [`from_iter_in`]: PBox::from_iter_in + /// [`from_huge_slice_in`]: PBox::from_huge_slice_in + pub fn from_huge_iter_in(iter: I, cx: &MemCx<'mcx>) -> Result + where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, + { + let iter = iter.into_iter(); + let len = iter.len(); + let uninit = PBox::<[MaybeUninit]>::new_huge_uninit_slice_in(len, cx) + .map_err(|_| FromIterError::OutOfMemory)?; + Self::fill_from_iter(uninit, iter, len) + } - // Get a thin pointer to the first element of the uninit buffer. + fn fill_from_iter( + mut uninit: PBox<'mcx, [MaybeUninit]>, + iter: I, + len: usize, + ) -> Result + where + I: Iterator, + { + // Thin pointer to the first element of the uninit buffer. let base: *mut T = uninit.as_mut_ptr().cast::(); - let mut guard = Guard { base, initialized: 0 }; let mut count = 0usize; for item in iter { if count >= len { - // Iterator over-yielded vs ExactSizeIterator promise. Stop here; - // we'll fail below since `count < len`. + // Iterator over-yielded vs ExactSizeIterator promise; stop here. break; } - // SAFETY: `count < len`, so `base.add(count)` is in-bounds for the freshly allocated slice. The slot is uninit MaybeUninit; writing a T is sound. + // SAFETY: `count < len`, so `base.add(count)` is in-bounds for the freshly allocated slice. The slot is uninit MaybeUninit; writing a T is sound. On panic the buffer stays uninitialized — safe because `T: Copy` has no Drop and `PBox<[MaybeUninit]>` has no Drop either; the memory is reclaimed when `cx` resets. unsafe { core::ptr::write(base.add(count), item) }; count += 1; - guard.initialized = count; } if count < len { - // Iterator under-yielded. Drop the initialized prefix via Guard, - // and return an error. The buffer itself stays in `cx`. - return Err(OutOfMemory::new()); + // Iterator under-yielded; the buffer stays in `cx` until reset. + return Err(FromIterError::IteratorUnderYield { promised: len, yielded: count }); } - // All elements initialized successfully — disarm the guard and assume_init. - core::mem::forget(guard); // SAFETY: every slot in the slice was written above (count == len). Ok(unsafe { uninit.assume_init() }) } } +/// Error returned by [`PBox::from_iter_in`]. +#[derive(Debug)] +pub enum FromIterError { + /// The initial slice allocation in the [`MemCx`] failed. + OutOfMemory, + /// The iterator yielded fewer items than its [`ExactSizeIterator::len`] advertised. This indicates a buggy iterator implementation; the buffer was allocated but only partially initialized. + IteratorUnderYield { promised: usize, yielded: usize }, +} + +impl core::fmt::Display for FromIterError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + FromIterError::OutOfMemory => f.write_str("out of memory"), + FromIterError::IteratorUnderYield { promised, yielded } => write!( + f, + "ExactSizeIterator under-yielded: promised {promised} items, got {yielded}", + ), + } + } +} + +impl std::error::Error for FromIterError {} + unsafe impl<'mcx, T> BoxRet for PBox<'mcx, T> where T: ?Sized + BorrowDatum, @@ -278,7 +331,12 @@ where // SAFETY: Due to BorrowDatum, this type has a definite size less than a Datum, and PBox must have an initialized pointee, so a copy is sound-by-construction. unsafe { let size = size_of_val(&*self.ptr.as_ptr()); - debug_assert!(size <= size_of::()); + // Hard assert (not debug_assert): a `?Sized` `BorrowDatum` whose dynamic size exceeds `Datum` would otherwise overrun the on-stack `Datum` slot via `copy_from_nonoverlapping`. Real PassBy::Value types are all `Sized` and ≤ 8 bytes today, but a future fat-pointee implementing `BorrowDatum` with `PASS = PassBy::Value` must crash, not corrupt the stack. + assert!( + size <= size_of::(), + "PBox::box_into: PassBy::Value pointee size {size} exceeds Datum size {}", + size_of::() + ); // using `BorrowDatum::point_from` handles endianness let datum_ptr = T::point_from(NonNull::from_mut(&mut datum).cast::()); datum_ptr.cast::().copy_from_nonoverlapping(self.ptr.cast(), size); From 34ac70d423baea57ac9432b0990cf7b630978da4 Mon Sep 17 00:00:00 2001 From: danielshih Date: Tue, 16 Jun 2026 12:29:14 +0000 Subject: [PATCH 3/3] cargo fmt --- pgrx-unit-tests/src/tests/pbox_slice_tests.rs | 21 +++++++++---------- pgrx/src/palloc/pbox.rs | 2 ++ 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/pgrx-unit-tests/src/tests/pbox_slice_tests.rs b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs index ec2df4b66..4af6ee896 100644 --- a/pgrx-unit-tests/src/tests/pbox_slice_tests.rs +++ b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs @@ -62,17 +62,18 @@ mod tests { let before = pg_sys::MemoryContextMemAllocated(parent, false); let child = PgMemoryContexts::new("pbox_slice_test_scope"); - // Switch CMC and arm the restore guard (drops in LIFO before `child`). - let _restore = { + // Switch CMC and arm the restore guard. The guard is held in a + // named binding (NOT `let _ = ...`, which would drop immediately) + // so it lives until the end of the unsafe block, then drops in + // LIFO order before `child` — restoring CMC before the child + // context is deleted. `#[allow(unused_variables)]` silences the + // dead-binding lint since the guard's effect is in its Drop impl. + #[allow(unused_variables)] + let restore_guard = { let previous = pg_sys::MemoryContextSwitchTo(child.value()); CxRestoreGuard { previous } }; memcx::current_context(|cx| f(cx)); - drop(_restore); - - // Drop the owned child explicitly so MemoryContextDelete runs - // before we sample `after`. - drop(child); let after = pg_sys::MemoryContextMemAllocated(parent, false); assert_eq!( @@ -202,8 +203,7 @@ mod tests { } } memcx::current_context(|cx| { - let r = - PBox::<[u32]>::from_iter_in(OverIter { i: 0, promised: 3 }, cx).expect("alloc"); + let r = PBox::<[u32]>::from_iter_in(OverIter { i: 0, promised: 3 }, cx).expect("alloc"); assert_eq!(&*r, &[0u32, 1, 2]); }); } @@ -351,8 +351,7 @@ mod tests { fn from_huge_iter_in_roundtrip() { memcx::current_context(|cx| { let v: Vec = (100..110).collect(); - let dst: PBox<[i64]> = - PBox::from_huge_iter_in(v.iter().copied(), cx).expect("alloc"); + let dst: PBox<[i64]> = PBox::from_huge_iter_in(v.iter().copied(), cx).expect("alloc"); assert_eq!(&*dst, v.as_slice()); }); } diff --git a/pgrx/src/palloc/pbox.rs b/pgrx/src/palloc/pbox.rs index f5d74b62d..37649cd88 100644 --- a/pgrx/src/palloc/pbox.rs +++ b/pgrx/src/palloc/pbox.rs @@ -166,6 +166,8 @@ impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { } /// Zeroed counterpart of [`new_huge_uninit_slice_in`]. + /// + /// [`new_huge_uninit_slice_in`]: PBox::new_huge_uninit_slice_in pub fn new_huge_zeroed_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { const { assert!(align_of::() <= size_of::()) }; let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?;