From fb9a72c4dd1d98ddbd03cfded511e9d749439838 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 02:49:06 -0700 Subject: [PATCH] Fix plugin-state leaks, callback-table lifetime, and unload gating in nxmem MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Revision of the Phase 6 plugin memory ABI after review rejection. The architecture and the ABI design are unchanged; this closes the specific safety defects and the test holes that let four mutations pass unnoticed. The two lifetime defects were the same mistake at two scales. `open_allocator` took on a debt when the plugin returned `Ok` and then dropped it on six rejection paths, and `AllocatorCore` pinned the plugin's *code* while freeing the callback context the plugin still held a pointer into. Both are now carried by the type system rather than by remembering: an `AbandonOnDrop` guard covers rejection paths that do not exist yet, and the bridge and callback table live in one boxed unit whose teardown is gated on the outstanding-release count. The abandon path could not have worked as written. It re-read the vtable through `read_prefix`, which validates, and the only way to reach it was a `read_prefix` failure — so the re-read failed identically and returned before calling `release`. Splitting out `read_prefix_unvalidated` is what makes the release reachable at all; the guard alone would have fixed nothing. `MemoryPlugin` now has a `Drop`. Without one, every early return and unwind skipped the unload gate silently and left `dlclose` free to unmap a module with live objects in it. Whether that actually unmaps is a property of the loader, not of this code, so the drop keeps the module mapped and counts it. The test holes shared a root cause: a fixture that lies about its size cannot exhibit an out-of-bounds read, because the bytes past the declaration are still valid. The short-struct and poisoned-tail fixtures are now backed by allocations that really end where they say they do, which is what lets Miri see the over-read. Assertions are on status codes and sizes rather than on substrings of human-readable messages. `live_views` is wired rather than removed: a shared prefix committed into a live allocation is a genuine plugin-side object, and views retire when the block is removed rather than when its bytes are freed. The header now pins 49 field offsets as well as 23 sizes, and the layout tests no longer skip on Windows — MSVC is the toolchain most likely to disagree about `#[repr(C)]` packing, which is the whole point of the test. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../include/nxmem_memory_abi.h | 79 ++++ crates/onnx-runtime-memory-abi/src/vtable.rs | 372 +++++++++++++++++- .../tests/header_contract.rs | 266 ++++++++++--- .../onnx-runtime-memory-host/src/allocator.rs | 352 ++++++++++++++--- crates/onnx-runtime-memory-host/src/error.rs | 14 + crates/onnx-runtime-memory-host/src/loader.rs | 82 +++- .../tests/nxmem_abi.rs | 233 ++++++++++- .../tests/nxmem_abi_callback_pinning.rs | 52 +++ .../tests/nxmem_abi_plugin_gate.rs | 96 +++++ .../tests/nxmem_abi_unload_gate.rs | 37 +- .../onnx-runtime-memory-testplugin/src/lib.rs | 279 ++++++++++++- docs/memory/MEMORY_PLUGIN_ABI.md | 46 ++- 12 files changed, 1765 insertions(+), 143 deletions(-) create mode 100644 crates/onnx-runtime-memory-host/tests/nxmem_abi_callback_pinning.rs create mode 100644 crates/onnx-runtime-memory-host/tests/nxmem_abi_plugin_gate.rs diff --git a/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h b/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h index 3e36c0ba4c..ec4e318631 100644 --- a/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h +++ b/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h @@ -412,6 +412,17 @@ typedef struct NxmemSharedPrefixCommitInfo { * A plugin must not hold any of its own locks across either callback. * * NXMEM_LAYOUT: NxmemHostCallbacks size=32 + * + * Every field offset below is part of the contract too, not just the total + * size. A prefix is only meaningful if the fields inside it stay where they + * are: MIN_STRUCT_SIZE_MINOR_0 is derived from a field offset, so inserting a + * field mid-struct would silently move it and every older peer would be + * misread. Pinning the offsets makes that a build failure. + * NXMEM_LAYOUT_FIELD: NxmemHostCallbacks.struct_size offset=0 + * NXMEM_LAYOUT_FIELD: NxmemHostCallbacks.abi_minor offset=4 + * NXMEM_LAYOUT_FIELD: NxmemHostCallbacks.host_ctx offset=8 + * NXMEM_LAYOUT_FIELD: NxmemHostCallbacks.request_reclaim offset=16 + * NXMEM_LAYOUT_FIELD: NxmemHostCallbacks.release_completed offset=24 */ typedef struct NxmemHostCallbacks { uint32_t struct_size; @@ -452,6 +463,22 @@ typedef struct NxmemOpenRequest { * allocator that published it and released with it — never separately. * * NXMEM_LAYOUT: NxmemVirtualBackingVtable size=72 + * + * Every field offset below is part of the contract too, not just the total + * size. A prefix is only meaningful if the fields inside it stay where they + * are: MIN_STRUCT_SIZE_MINOR_0 is derived from a field offset, so inserting a + * field mid-struct would silently move it and every older peer would be + * misread. Pinning the offsets makes that a build failure. + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.struct_size offset=0 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.abi_minor offset=4 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.mechanism_id offset=8 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.ctx offset=16 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.allocate_committed offset=24 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.commit_range offset=32 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.decommit_range offset=40 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.mapped_bytes_for_ranges offset=48 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.mapped_bytes_for_allocation offset=56 + * NXMEM_LAYOUT_FIELD: NxmemVirtualBackingVtable.committed_bytes offset=64 */ typedef struct NxmemVirtualBackingVtable { uint32_t struct_size; @@ -479,6 +506,21 @@ typedef struct NxmemVirtualBackingVtable { * Owned by the allocator that published it, exactly as virtual backing is. * * NXMEM_LAYOUT: NxmemSharedMappingVtable size=64 + * + * Every field offset below is part of the contract too, not just the total + * size. A prefix is only meaningful if the fields inside it stay where they + * are: MIN_STRUCT_SIZE_MINOR_0 is derived from a field offset, so inserting a + * field mid-struct would silently move it and every older peer would be + * misread. Pinning the offsets makes that a build failure. + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.struct_size offset=0 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.abi_minor offset=4 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.mechanism_id offset=8 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.ctx offset=16 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.create_shared_prefix offset=24 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.retain_shared_prefix offset=32 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.release_shared_prefix offset=40 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.incremental_owned_bytes offset=48 + * NXMEM_LAYOUT_FIELD: NxmemSharedMappingVtable.commit_shared_prefix offset=56 */ typedef struct NxmemSharedMappingVtable { uint32_t struct_size; @@ -512,6 +554,29 @@ typedef struct NxmemSharedMappingVtable { * back unchanged to every slot. * * NXMEM_LAYOUT: NxmemAllocatorVtable size=128 + * + * Every field offset below is part of the contract too, not just the total + * size. A prefix is only meaningful if the fields inside it stay where they + * are: MIN_STRUCT_SIZE_MINOR_0 is derived from a field offset, so inserting a + * field mid-struct would silently move it and every older peer would be + * misread. Pinning the offsets makes that a build failure. + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.struct_size offset=0 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.abi_minor offset=4 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.mechanism_id offset=8 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.device offset=16 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.capability_flags offset=24 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.name offset=32 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.ctx offset=40 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.allocate offset=48 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.deallocate offset=56 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.retain offset=64 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.release offset=72 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.virtual_backing offset=80 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.shared_mapping offset=88 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.enqueue_release offset=96 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.drain_releases offset=104 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.pending_release_count offset=112 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorVtable.release_allocation offset=120 */ typedef struct NxmemAllocatorVtable { uint32_t struct_size; @@ -552,6 +617,20 @@ typedef struct NxmemAllocatorVtable { * been released. `name` must outlive the factory. * * NXMEM_LAYOUT: NxmemAllocatorFactoryVtable size=56 + * + * Every field offset below is part of the contract too, not just the total + * size. A prefix is only meaningful if the fields inside it stay where they + * are: MIN_STRUCT_SIZE_MINOR_0 is derived from a field offset, so inserting a + * field mid-struct would silently move it and every older peer would be + * misread. Pinning the offsets makes that a build failure. + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.struct_size offset=0 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.abi_minor offset=4 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.name offset=8 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.device offset=16 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.capability_flags offset=24 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.ctx offset=32 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.open_allocator offset=40 + * NXMEM_LAYOUT_FIELD: NxmemAllocatorFactoryVtable.release offset=48 */ typedef struct NxmemAllocatorFactoryVtable { uint32_t struct_size; diff --git a/crates/onnx-runtime-memory-abi/src/vtable.rs b/crates/onnx-runtime-memory-abi/src/vtable.rs index 9ac1ecf118..77fc57c614 100644 --- a/crates/onnx-runtime-memory-abi/src/vtable.rs +++ b/crates/onnx-runtime-memory-abi/src/vtable.rs @@ -291,6 +291,36 @@ impl NxmemAllocatorVtable { pub unsafe fn read_prefix( ptr: *const Self, negotiated_minor: u32, + ) -> Result { + // SAFETY: delegated verbatim to this function's contract. + let vtable = unsafe { Self::read_prefix_unvalidated(ptr, negotiated_minor) }?; + vtable.validate_required()?; + Ok(vtable) + } + + /// Read the prefix while skipping [`Self::validate_required`]. + /// + /// This performs **every memory-safety check** [`Self::read_prefix`] does — + /// null, alignment, self-consistent `struct_size`, bounded copy, and the + /// clamp — and differs only in that it does not insist the required slots + /// are populated or that `mechanism_id` is non-zero. + /// + /// It exists for exactly one caller: the host's abandon path, which must + /// still find `release` in a vtable it has *already decided to refuse*. + /// Refusing a vtable does not cancel the plugin's `Ok` from + /// `open_allocator`, so the host still owes it a `release`; re-reading + /// through the validating entry point would fail for the very reason the + /// vtable was refused and the plugin's state would be stranded. + /// + /// Do not use this to decide whether a vtable may be *called*. The only + /// slot it is safe to invoke on the result is `release`. + /// + /// # Safety + /// + /// Same contract as [`Self::read_prefix`]. + pub unsafe fn read_prefix_unvalidated( + ptr: *const Self, + negotiated_minor: u32, ) -> Result { check_vtable_ptr(ptr, "allocator")?; // SAFETY: `ptr` was proved non-null and aligned; `struct_size` is the @@ -333,7 +363,6 @@ impl NxmemAllocatorVtable { if effective_minor < 1 { vtable.release_allocation = None; } - vtable.validate_required()?; Ok(vtable) } @@ -970,4 +999,345 @@ mod tests { assert_eq!(read.mechanism_id, 42); assert!(read.release_allocation.is_some()); } + + /// A `T`-sized, `T`-aligned buffer whose every byte the test controls. + /// + /// Used to build senders that a plain `#[repr(C)]` literal cannot express: + /// a struct whose bytes *beyond the size it declares* hold a known poison + /// pattern rather than plausible values. + fn poisoned_buffer(prefix: &T, declared_size: usize, poison: u8) -> Vec { + let size = core::mem::size_of::(); + assert!(declared_size <= size); + assert!(core::mem::align_of::() <= core::mem::align_of::()); + // A `Vec` gives 8-byte alignment, which is at least `align_of::()` + // for every nxmem vtable, without depending on the global allocator + // over-aligning a byte buffer. + let mut storage = vec![0u64; size.div_ceil(core::mem::size_of::())]; + let base = storage.as_mut_ptr().cast::(); + // SAFETY: `storage` owns at least `size` bytes at `base`, and `prefix` + // is a live `T`, so both regions are valid for `size` bytes. They + // cannot overlap: `storage` was allocated in this function. + unsafe { + core::ptr::write_bytes(base, poison, size); + core::ptr::copy_nonoverlapping((prefix as *const T).cast::(), base, declared_size); + } + storage + } + + /// Every byte of `value`, including padding, which is always initialised + /// here because the value came from a whole-struct byte copy. + fn as_bytes(value: &T) -> &[u8] { + // SAFETY: `T` is `#[repr(C)]` plain data and every byte of `value` was + // initialised by `copy_prefix`, which starts from a zeroed buffer. + unsafe { + core::slice::from_raw_parts((value as *const T).cast::(), core::mem::size_of::()) + } + } + + /// **The bounded read itself**, tested without going through any of the + /// checks layered on top of it. + /// + /// `copy_prefix` is the memory-safety core of prefix negotiation: it is + /// the single line that stops the host reading past the end of a sender's + /// allocation. Asserting on it directly is the only way to pin the bound + /// independently of the *nulling* that happens afterwards — a test that + /// only checks `release_allocation.is_none()` stays green when the bound + /// is destroyed, because the clamp nulls that slot either way. + #[test] + fn copy_prefix_reads_exactly_the_declared_prefix_and_zeroes_the_rest() { + const POISON: u8 = 0xAB; + let declared = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0; + let source = conforming_allocator(1); + let storage = poisoned_buffer(&source, declared, POISON); + let ptr = storage.as_ptr().cast::(); + + // SAFETY: `storage` holds `size_of::()` readable, + // correctly aligned, fully initialised bytes, which is more than + // `declared`. + let read = unsafe { copy_prefix(ptr, declared) }; + let read_bytes = as_bytes(&read); + let source_bytes = as_bytes(&source); + + assert_eq!( + &read_bytes[..declared], + &source_bytes[..declared], + "the declared prefix must be copied verbatim" + ); + assert!( + read_bytes[declared..].iter().all(|byte| *byte == 0), + "every byte past the declared size must read back as zero, but the copy brought back \ + {:x?}; the bound was not honoured", + &read_bytes[declared..] + ); + } + + /// The destination is bounded too: a sender from a *newer* build declares + /// more bytes than this build defines, and the copy must stop at this + /// build's struct rather than overrun the local. + #[test] + fn copy_prefix_never_writes_more_than_this_builds_struct() { + let source = conforming_allocator(1); + let size = core::mem::size_of::(); + let storage = poisoned_buffer(&source, size, 0); + let ptr = storage.as_ptr().cast::(); + + // A newer peer declares a struct twice this size. Copying that many + // bytes into a local `NxmemAllocatorVtable` would smash the stack. + // SAFETY: `copy_prefix` clamps to `size_of::()`, so it reads only + // the `size` initialised bytes `storage` really has. + let read = unsafe { copy_prefix(ptr, size * 2) }; + assert_eq!(read.mechanism_id, source.mechanism_id); + assert_eq!(as_bytes(&read), as_bytes(&source)); + } + + /// The same property, observed through the public entry point. + /// + /// A minor-0 sender is followed in memory by poison. Every byte the host + /// hands back past the declared prefix must be zero — that is what makes + /// "a slot the sender does not define reads back as `None`" true rather + /// than merely usually true. + #[test] + fn read_prefix_zeroes_every_byte_past_the_senders_declared_size() { + const POISON: u8 = 0xAB; + let declared = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0; + let source = NxmemAllocatorVtable { + struct_size: declared as u32, + abi_minor: 0, + ..conforming_allocator(0) + }; + let storage = poisoned_buffer(&source, declared, POISON); + let ptr = storage.as_ptr().cast::(); + + // SAFETY: `storage` holds a fully initialised, correctly aligned + // `NxmemAllocatorVtable`-sized region whose first `declared` bytes are + // a well-formed minor-0 vtable. + let read = unsafe { NxmemAllocatorVtable::read_prefix(ptr, 1) } + .expect("a well-formed minor-0 sender is acceptable to a minor-1 host"); + + assert_eq!(read.abi_minor, 0); + assert!( + read.release_allocation.is_none(), + "a slot the sender never defined must not be callable" + ); + assert!( + as_bytes(&read)[declared..].iter().all(|byte| *byte == 0), + "the host read {:x?} from past the sender's struct; the poison beyond {declared} \ + bytes must never reach the host's copy", + &as_bytes(&read)[declared..] + ); + } + + /// A vtable that is exactly as big as it declares, and no bigger. + /// + /// The point is not the return value — it is that reading even one byte + /// past `declared` is a genuine out-of-bounds access of a heap allocation, + /// so a sanitiser or Miri can see it. A fixture backed by a full-size + /// struct that merely *lies* about its size cannot exhibit that. + #[test] + fn a_vtable_sized_exactly_to_its_declaration_is_read_within_bounds() { + let declared = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0; + let source = NxmemAllocatorVtable { + struct_size: declared as u32, + abi_minor: 0, + ..conforming_allocator(0) + }; + // Exactly `declared` bytes, not a byte more. Backed by `u64` rather + // than `u8` so the allocation is aligned for the vtable by + // construction: a `Vec` is only byte-aligned, and while the system + // allocator happens to hand back an 8-aligned block, Miri does not, so + // a byte-backed fixture fails under Miri for the wrong reason. + assert!( + declared.is_multiple_of(core::mem::size_of::()), + "the minor-0 prefix must be a whole number of words for an exact fixture" + ); + assert!(core::mem::align_of::() <= core::mem::align_of::()); + let mut exact = vec![0u64; declared / core::mem::size_of::()]; + // SAFETY: both regions are valid for `declared` bytes and do not + // overlap; `exact` was allocated in this function and holds exactly + // `declared` bytes. + unsafe { + core::ptr::copy_nonoverlapping( + (&source as *const NxmemAllocatorVtable).cast::(), + exact.as_mut_ptr().cast::(), + declared, + ); + } + + // SAFETY: the buffer holds exactly `declared` readable bytes and + // declares exactly that size, so a conforming reader stays in bounds. + let read = + unsafe { NxmemAllocatorVtable::read_prefix(exact.as_ptr().cast(), 1) }.expect("valid"); + assert_eq!(read.abi_minor, 0); + assert!(read.release_allocation.is_none()); + } + + /// A refused vtable can still be handed back. + /// + /// `read_prefix` rejects a vtable whose required slots are null, so the + /// host cannot use it to find `release` on the abandon path — the read + /// would fail for exactly the reason the vtable was refused, and the + /// plugin's state would be stranded. + #[test] + fn read_prefix_unvalidated_still_finds_release_in_a_refused_vtable() { + let refused = NxmemAllocatorVtable { + // A plugin that forgot `allocate`: refused, but it has already + // created state and `release` is right there. + allocate: None, + ..conforming_allocator(1) + }; + // SAFETY: `refused` is a live, aligned local of exactly this type. + let status = + unsafe { NxmemAllocatorVtable::read_prefix(&refused, 1) }.expect_err("missing slot"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::ShortStruct)); + + // SAFETY: as above. + let unvalidated = unsafe { NxmemAllocatorVtable::read_prefix_unvalidated(&refused, 1) } + .expect("the abandon path must still be able to reach `release`"); + assert!( + unvalidated.release.is_some(), + "a refusal must not strand the state the plugin already created" + ); + + // The same is true of a zero `mechanism_id`, the other validation. + let zero_id = NxmemAllocatorVtable { + mechanism_id: 0, + ..conforming_allocator(1) + }; + // SAFETY: `zero_id` is a live, aligned local. + assert_eq!( + unsafe { NxmemAllocatorVtable::read_prefix(&zero_id, 1) } + .expect_err("zero mechanism id") + .status_code(), + Some(NxmemStatusCode::InvalidArgument) + ); + // SAFETY: as above. + assert!( + unsafe { NxmemAllocatorVtable::read_prefix_unvalidated(&zero_id, 1) } + .expect("the memory-safety checks still pass") + .release + .is_some() + ); + } + + /// The unvalidated read is only a *validation* shortcut: every + /// memory-safety check still applies. + #[test] + fn read_prefix_unvalidated_still_refuses_unsafe_pointers_and_short_structs() { + // SAFETY: a null pointer is exactly what is under test. + assert_eq!( + unsafe { NxmemAllocatorVtable::read_prefix_unvalidated(core::ptr::null(), 1) } + .expect_err("null") + .status_code(), + Some(NxmemStatusCode::InvalidArgument) + ); + + let vtable = conforming_allocator(1); + let misaligned = + (&vtable as *const NxmemAllocatorVtable as usize + 1) as *const NxmemAllocatorVtable; + // SAFETY: rejected on alignment before any dereference. + assert_eq!( + unsafe { NxmemAllocatorVtable::read_prefix_unvalidated(misaligned, 1) } + .expect_err("misaligned") + .status_code(), + Some(NxmemStatusCode::InvalidArgument) + ); + + let short = NxmemAllocatorVtable { + struct_size: (NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 - 1) as u32, + ..conforming_allocator(0) + }; + // SAFETY: `short` is a live, aligned local; only `struct_size` and + // `abi_minor` are read before the refusal. + assert_eq!( + unsafe { NxmemAllocatorVtable::read_prefix_unvalidated(&short, 1) } + .expect_err("short struct") + .status_code(), + Some(NxmemStatusCode::ShortStruct) + ); + } + + /// Field **offsets**, not just sizes. + /// + /// `MIN_STRUCT_SIZE_MINOR_0` is derived from `offset_of!`, so inserting a + /// field in the middle of a vtable silently moves the constant and every + /// size-based assertion keeps passing while the wire format has changed + /// underneath every already-built plugin. Pinning the offsets makes that + /// edit fail here, where it is cheap, instead of in a third party's + /// process. The same numbers are pinned against the C header, and against + /// a real C compiler, in `tests/header_contract.rs`. + #[test] + fn the_allocator_vtable_field_offsets_are_pinned() { + use core::mem::offset_of; + + for (field, actual, expected) in [ + ( + "struct_size", + offset_of!(NxmemAllocatorVtable, struct_size), + 0, + ), + ("abi_minor", offset_of!(NxmemAllocatorVtable, abi_minor), 4), + ( + "mechanism_id", + offset_of!(NxmemAllocatorVtable, mechanism_id), + 8, + ), + ("device", offset_of!(NxmemAllocatorVtable, device), 16), + ( + "capability_flags", + offset_of!(NxmemAllocatorVtable, capability_flags), + 24, + ), + ("name", offset_of!(NxmemAllocatorVtable, name), 32), + ("ctx", offset_of!(NxmemAllocatorVtable, ctx), 40), + ("allocate", offset_of!(NxmemAllocatorVtable, allocate), 48), + ( + "deallocate", + offset_of!(NxmemAllocatorVtable, deallocate), + 56, + ), + ("retain", offset_of!(NxmemAllocatorVtable, retain), 64), + ("release", offset_of!(NxmemAllocatorVtable, release), 72), + ( + "virtual_backing", + offset_of!(NxmemAllocatorVtable, virtual_backing), + 80, + ), + ( + "shared_mapping", + offset_of!(NxmemAllocatorVtable, shared_mapping), + 88, + ), + ( + "enqueue_release", + offset_of!(NxmemAllocatorVtable, enqueue_release), + 96, + ), + ( + "drain_releases", + offset_of!(NxmemAllocatorVtable, drain_releases), + 104, + ), + ( + "pending_release_count", + offset_of!(NxmemAllocatorVtable, pending_release_count), + 112, + ), + ( + "release_allocation", + offset_of!(NxmemAllocatorVtable, release_allocation), + 120, + ), + ] { + assert_eq!( + actual, expected, + "NxmemAllocatorVtable.{field} moved from {expected} to {actual}; every plugin \ + already built against this ABI reads the old offset, so this is a breaking \ + change and needs a major bump, not an edit here" + ); + } + + // And the derived prefix constants follow from those offsets. + assert_eq!(NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0, 120); + assert_eq!(NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_1, 128); + } } diff --git a/crates/onnx-runtime-memory-abi/tests/header_contract.rs b/crates/onnx-runtime-memory-abi/tests/header_contract.rs index 12f6b9991f..6b2bb9a0e4 100644 --- a/crates/onnx-runtime-memory-abi/tests/header_contract.rs +++ b/crates/onnx-runtime-memory-abi/tests/header_contract.rs @@ -39,6 +39,72 @@ fn annotated_sizes() -> Vec<(String, usize)> { .collect() } +/// Every `NXMEM_LAYOUT_FIELD: . offset=` annotation. +/// +/// Offsets are annotated only for the structs the ABI reads by *prefix*, where +/// a field's position is load-bearing rather than incidental. +fn annotated_offsets() -> Vec<(String, String, usize)> { + let header = std::fs::read_to_string(header_path()).expect("the public header is readable"); + header + .lines() + .filter_map(|line| { + let rest = line.split_once("NXMEM_LAYOUT_FIELD:")?.1.trim(); + let (path, offset) = rest.split_once("offset=")?; + let (type_name, field) = path.trim().split_once('.').unwrap_or_else(|| { + panic!("malformed field annotation, expected `Type.field`: {line}") + }); + let offset = offset + .trim() + .parse::() + .unwrap_or_else(|_| panic!("malformed field annotation: {line}")); + Some(( + type_name.trim().to_string(), + field.trim().to_string(), + offset, + )) + }) + .collect() +} + +/// The Rust offset of one annotated field, or `None` if this test does not +/// know the field. +/// +/// Spelled out rather than derived, because the point is to notice when the +/// Rust definition moves: a table generated from the definition would agree +/// with it by construction and prove nothing. +fn rust_offset(type_name: &str, field: &str) -> Option { + use std::mem::offset_of; + macro_rules! table { + ($($ty:ident { $($field:ident),+ $(,)? }),+ $(,)?) => { + match (type_name, field) { + $($((stringify!($ty), stringify!($field)) => Some(offset_of!($ty, $field)),)+)+ + _ => None, + } + }; + } + table! { + NxmemHostCallbacks { + struct_size, abi_minor, host_ctx, request_reclaim, release_completed, + }, + NxmemVirtualBackingVtable { + struct_size, abi_minor, mechanism_id, ctx, allocate_committed, commit_range, + decommit_range, mapped_bytes_for_ranges, mapped_bytes_for_allocation, committed_bytes, + }, + NxmemSharedMappingVtable { + struct_size, abi_minor, mechanism_id, ctx, create_shared_prefix, retain_shared_prefix, + release_shared_prefix, incremental_owned_bytes, commit_shared_prefix, + }, + NxmemAllocatorVtable { + struct_size, abi_minor, mechanism_id, device, capability_flags, name, ctx, allocate, + deallocate, retain, release, virtual_backing, shared_mapping, enqueue_release, + drain_releases, pending_release_count, release_allocation, + }, + NxmemAllocatorFactoryVtable { + struct_size, abi_minor, name, device, capability_flags, ctx, open_allocator, release, + }, + } +} + /// The Rust size of every type the header annotates. /// /// Adding a `#[repr(C)]` type to the header without adding it here fails the @@ -131,78 +197,164 @@ fn header_layout_matches_the_rust_definitions() { } } +/// A C compiler and the flag dialect it speaks. +struct CCompiler { + path: PathBuf, + /// MSVC takes `/c`, `/I`, `/Fo` where everything else takes `-c`, `-I`, + /// `-o`. Nothing else about these invocations is platform-specific. + msvc: bool, +} + +impl CCompiler { + /// Compile `source` to an object file, returning the compiler's stderr on + /// failure. + fn compile(&self, source: &Path, include_dir: &Path, out: &Path) -> Result<(), String> { + let mut command = Command::new(&self.path); + if self.msvc { + command + .arg("/c") + .arg("/std:c11") + .arg("/W4") + .arg("/WX") + .arg("/nologo") + .arg(format!("/I{}", include_dir.display())) + .arg(source) + .arg(format!("/Fo{}", out.display())); + } else { + command + .arg("-c") + .arg("-std=c11") + .arg("-Wall") + .arg("-Wextra") + .arg("-Werror") + .arg("-fPIC") + .arg("-I") + .arg(include_dir) + .arg(source) + .arg("-o") + .arg(out); + } + let output = command.output().expect("the C compiler runs"); + if output.status.success() { + return Ok(()); + } + // MSVC writes diagnostics to stdout, everyone else to stderr. + Err(format!( + "{}{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + )) + } +} + /// The first C compiler on this machine, if any. -fn find_cc() -> Option { +/// +/// `cl` is tried too, and deliberately so: MSVC is the toolchain most likely +/// to disagree with Rust about `#[repr(C)]` packing, and this crate's entire +/// deliverable is a binary contract that has to hold across toolchains. A +/// layout test that never runs on the compiler most likely to break it is not +/// a layout test. +fn find_cc() -> Option { if let Ok(cc) = std::env::var("CC") && !cc.is_empty() { - return Some(PathBuf::from(cc)); + let msvc = Path::new(&cc) + .file_stem() + .and_then(|stem| stem.to_str()) + .is_some_and(|stem| stem.eq_ignore_ascii_case("cl")); + return Some(CCompiler { + path: PathBuf::from(cc), + msvc, + }); } for candidate in ["cc", "clang", "gcc"] { if let Ok(output) = Command::new(candidate).arg("--version").output() && output.status.success() { - return Some(PathBuf::from(candidate)); + return Some(CCompiler { + path: PathBuf::from(candidate), + msvc: false, + }); } } + // `cl` has no `--version`; it prints a banner and its usage for `/?`. + if let Ok(output) = Command::new("cl").arg("/?").output() + && output.status.success() + { + return Some(CCompiler { + path: PathBuf::from("cl"), + msvc: true, + }); + } None } +/// The message every compiler-backed test fails with when there is none. +/// +/// Skipping instead would let the binary contract go unchecked on exactly the +/// platform where it is most likely to be wrong, while the suite still +/// reported success. +const NO_CC: &str = "no C compiler found (tried $CC, cc, clang, gcc, cl); the public C header \ + cannot be checked without one"; + /// A real C compiler accepts the header and the example plugin. /// /// This is what backs the "one minimal plugin example that documents the /// contract without requiring workspace linking" claim: the example is /// compiled with nothing on the include path but the header itself. #[test] -#[cfg(unix)] fn nxmem_c_example_compiles() { - let cc = find_cc().expect( - "no C compiler found (tried $CC, cc, clang, gcc); \ - the public C header cannot be checked without one", - ); + let cc = find_cc().expect(NO_CC); let out = std::env::temp_dir().join(format!("nxmem_minimal_plugin_{}.o", std::process::id())); - let output = Command::new(&cc) - .arg("-c") - .arg("-std=c11") - .arg("-Wall") - .arg("-Wextra") - .arg("-Werror") - .arg("-fPIC") - .arg("-I") - .arg(header_path().parent().expect("include directory")) - .arg(example_path()) - .arg("-o") - .arg(&out) - .output() - .expect("the C compiler runs"); - - let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); - let _ = std::fs::remove_file(&out); - assert!( - output.status.success(), - "the minimal C plugin must compile against the public header alone:\n{stderr}" + let result = cc.compile( + &example_path(), + header_path().parent().expect("include directory"), + &out, ); + let _ = std::fs::remove_file(&out); + if let Err(diagnostics) = result { + panic!("the minimal C plugin must compile against the public header alone:\n{diagnostics}"); + } } -/// A C compiler agrees with Rust about every annotated struct size. +/// A C compiler agrees with Rust about every annotated struct size **and every +/// annotated field offset**. /// /// The layout test above compares Rust against a *comment*. This one compares /// Rust against the compiler that will actually build a plugin, which is the /// opinion that decides whether a real plugin works. +/// +/// Offsets matter as much as sizes here. Prefix negotiation is built on the +/// claim that a newer struct is a byte-for-byte superset of an older one, and +/// the minimum prefix sizes are *derived* from field offsets — so a field +/// inserted in the middle keeps the total size plausible while silently moving +/// the boundary every older peer reads up to. +/// +/// Not gated on Unix. These are plain `-c`/`/c` invocations, and MSVC is +/// precisely the compiler most likely to disagree. #[test] -#[cfg(all(unix, target_pointer_width = "64"))] +#[cfg(target_pointer_width = "64")] fn a_c_compiler_agrees_with_the_rust_layouts() { - let cc = find_cc().expect( - "no C compiler found (tried $CC, cc, clang, gcc); \ - the public C header cannot be checked without one", - ); + let cc = find_cc().expect(NO_CC); - let mut source = String::from("#include \"nxmem_memory_abi.h\"\n"); + let mut source = String::from("#include \n#include \"nxmem_memory_abi.h\"\n"); for (name, size) in rust_sizes() { source.push_str(&format!( "_Static_assert(sizeof({name}) == {size}, \"{name} disagrees with Rust\");\n" )); } + let offsets = annotated_offsets(); + assert!( + !offsets.is_empty(), + "the header must annotate field offsets; finding none means the annotation format \ + changed and this test stopped checking anything" + ); + for (type_name, field, offset) in &offsets { + source.push_str(&format!( + "_Static_assert(offsetof({type_name}, {field}) == {offset}, \ + \"{type_name}.{field} disagrees with the header\");\n" + )); + } source.push_str("int nxmem_layout_probe(void) { return 0; }\n"); let dir = std::env::temp_dir().join(format!("nxmem_layout_{}", std::process::id())); @@ -210,21 +362,37 @@ fn a_c_compiler_agrees_with_the_rust_layouts() { let probe = dir.join("probe.c"); std::fs::write(&probe, source).expect("probe source is writable"); - let output = Command::new(&cc) - .arg("-c") - .arg("-std=c11") - .arg("-I") - .arg(header_path().parent().expect("include directory")) - .arg(&probe) - .arg("-o") - .arg(dir.join("probe.o")) - .output() - .expect("the C compiler runs"); - - let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let result = cc.compile( + &probe, + header_path().parent().expect("include directory"), + &dir.join("probe.o"), + ); let _ = std::fs::remove_dir_all(&dir); + if let Err(diagnostics) = result { + panic!("a C compiler disagrees with the Rust struct layouts:\n{diagnostics}"); + } +} + +/// Rust agrees with every field offset the header annotates. +/// +/// The C probe above proves the *header* and a C compiler agree. This proves +/// the header and *Rust* agree, without needing a C compiler at all — so the +/// offsets stay pinned even on a machine that cannot run the probe. +#[test] +fn the_header_field_offsets_match_rust() { + let annotated = annotated_offsets(); assert!( - output.status.success(), - "a C compiler disagrees with the Rust struct layouts:\n{stderr}" + !annotated.is_empty(), + "the header must annotate field offsets; finding none means the annotation format \ + changed and this test stopped checking anything" ); + for (type_name, field, offset) in annotated { + let actual = rust_offset(&type_name, &field).unwrap_or_else(|| { + panic!("the header annotates {type_name}.{field}, which this test does not know") + }); + assert_eq!( + actual, offset, + "{type_name}.{field}: the header says offset {offset}, Rust says {actual}" + ); + } } diff --git a/crates/onnx-runtime-memory-host/src/allocator.rs b/crates/onnx-runtime-memory-host/src/allocator.rs index ca3474b4a3..edf37350b6 100644 --- a/crates/onnx-runtime-memory-host/src/allocator.rs +++ b/crates/onnx-runtime-memory-host/src/allocator.rs @@ -88,6 +88,55 @@ struct HostBridge { reclaim_failures: AtomicU64, retired: Mutex>, module: Arc, + /// Deferred releases queued against **this allocator** that have not yet + /// reported completion. + /// + /// The module-wide tally in [`PluginModule`] answers "may the module + /// unload"; this one answers the narrower question "may this callback + /// table be freed", which is what the plugin's stored + /// `*const NxmemHostCallbacks` depends on. They are deliberately separate: + /// an allocator can be dropped long before the module is, and it is the + /// allocator's drop that would otherwise free the table. + outstanding_releases: AtomicU64, +} + +/// The callback table and the state behind it, as one heap-resident unit. +/// +/// The plugin captures `&callbacks` at open and dereferences it later, from +/// its own threads. `callbacks.host_ctx` points at `bridge`, in the same +/// allocation. The ABI contract ([`onnx_runtime_memory_abi`], `NxmemHostCallbacks`) +/// says the host keeps this alive past the allocator's final `release` **and +/// past every queued release**, so its lifetime cannot simply be the +/// allocator's: a queued release outlives the allocator by construction. +/// +/// Boxing the whole thing gives both members one stable address for the life +/// of the context, and lets [`AllocatorCore::drop`] decide its fate rather +/// than having the compiler free it unconditionally. +#[derive(Debug)] +struct HostCallbackContext { + bridge: HostBridge, + callbacks: NxmemHostCallbacks, +} + +impl HostCallbackContext { + /// Build a context whose callback table already points at its own bridge. + fn new(minor: u32, bridge: HostBridge) -> Box { + let mut context = Box::new(Self { + bridge, + callbacks: NxmemHostCallbacks { + struct_size: core::mem::size_of::() as u32, + abi_minor: minor, + // Patched below, once the box has its final address. + host_ctx: core::ptr::null_mut(), + request_reclaim: Some(host_request_reclaim), + release_completed: Some(host_release_completed), + }, + }); + // Boxing first and patching second is what makes the pointer stable: + // moving the `Box` moves the pointer, never the allocation it names. + context.callbacks.host_ctx = (&raw const context.bridge) as *mut core::ffi::c_void; + context + } } /// `request_reclaim` trampoline. @@ -195,6 +244,10 @@ unsafe extern "C" fn host_release_completed( } bridge.module.release_retired(); bridge.module.allocation_closed(); + // Last, because it is what permits the table this callback is running + // against to be freed. Decrementing it earlier would open a window in + // which `AllocatorCore::drop` could free the bridge under our feet. + bridge.outstanding_releases.fetch_sub(1, Ordering::AcqRel); NxmemStatus::ok() }) } @@ -207,6 +260,17 @@ struct LiveAllocation { align: usize, } +/// How many times [`AllocatorCore::drop`] asks a plugin to drain before giving +/// up and leaking the callback table instead. +/// +/// Draining is a courtesy, not an obligation: the plugin owns the queue and is +/// entitled to hold it. The bound exists so a plugin that reports progress it +/// is not making cannot hang a `drop`. +const MAX_DRAIN_PASSES_ON_DROP: usize = 16; + +/// How many host callback tables this process has deliberately leaked. +static LEAKED_CALLBACK_TABLES: AtomicU64 = AtomicU64::new(0); + /// The shared core of a plugin-backed mechanism. /// /// Every capability view holds an `Arc` of this, so the plugin allocator is @@ -225,12 +289,14 @@ pub struct AllocatorCore { capability_flags: u64, next_allocation_id: AtomicU64, live: Mutex>, - /// Boxed so the plugin's stored pointer stays valid across moves of the - /// core. Declared before `module` and dropped after `Drop::drop` has - /// already made the final `release` call, so the plugin can still touch - /// the table from inside that call. - bridge: Box, - callbacks: Box, + /// The callback table the plugin holds a raw pointer to. + /// + /// `Option` so [`Drop`] can *take* it: the table must survive not only + /// this allocator's final `release` but every queued release that can + /// still name it, and a queued release can outlive the allocator. Taking + /// it lets `drop` decide between freeing and deliberately leaking, which a + /// plain field cannot express. + context: Option>, module: Arc, } @@ -262,14 +328,39 @@ impl AllocatorCore { self.vtable.abi_minor.min(self.negotiated_minor) } + /// The callback context, which exists for the whole life of the core. + /// + /// It is only ever `None` inside [`Drop::drop`], after the context has + /// been taken to decide its fate. + fn context(&self) -> &HostCallbackContext { + self.context + .as_ref() + .expect("the callback context is taken only while the allocator core is being dropped") + } + + fn bridge(&self) -> &HostBridge { + &self.context().bridge + } + /// How many times the plugin has called back into the host for reclaim. pub fn reclaim_calls(&self) -> u64 { - self.bridge.reclaim_calls.load(Ordering::Acquire) + self.bridge().reclaim_calls.load(Ordering::Acquire) } /// How many of those reclaim calls the host refused. pub fn reclaim_failures(&self) -> u64 { - self.bridge.reclaim_failures.load(Ordering::Acquire) + self.bridge().reclaim_failures.load(Ordering::Acquire) + } + + /// Deferred releases queued against this allocator that have not yet + /// reported completion. + /// + /// While this is non-zero the plugin may still dereference the host + /// callback table from one of its own threads, so the table must not be + /// freed. It is exposed because "may this be torn down" is a question a + /// caller is entitled to ask before dropping the allocator. + pub fn outstanding_releases(&self) -> u64 { + self.bridge().outstanding_releases.load(Ordering::Acquire) } /// The callback table handed to the plugin. @@ -277,12 +368,35 @@ impl AllocatorCore { /// Exposed so a caller can confirm the table the plugin holds is the one /// this host installed, and that its `abi_minor` matches negotiation. pub fn host_callbacks(&self) -> &NxmemHostCallbacks { - &self.callbacks + &self.context().callbacks + } + + /// Whether the plugin published a structured-release slot that survived + /// prefix negotiation. + /// + /// This reports the *slot*, not the capability: a mechanism that declares + /// minor 0 must read back with the slot nulled even if the bytes at that + /// offset in the sender's memory were something else entirely. + pub fn publishes_structured_release_slot(&self) -> bool { + self.vtable.release_allocation.is_some() + } + + /// How many host callback tables this process has deliberately leaked. + /// + /// One per allocator that was dropped while a queued release could still + /// name its callback table. Process-wide, because the tables are freed — + /// or not — independently of any one allocator. + /// + /// A non-zero value is not corruption; it is the cost of refusing to hand + /// a plugin thread a dangling pointer. Treat it as a leak to be fixed by + /// draining releases before dropping the allocator. + pub fn leaked_callback_tables() -> u64 { + LEAKED_CALLBACK_TABLES.load(Ordering::Acquire) } /// Deferred releases the plugin has reported as retired, in report order. pub fn retired_releases(&self) -> Vec { - self.bridge + self.bridge() .retired .lock() .map(|retired| retired.clone()) @@ -294,6 +408,37 @@ impl AllocatorCore { self.live.lock().map(|live| live.len()).unwrap_or(0) } + /// Ask the plugin to retire everything it has queued, before teardown. + /// + /// Bounded rather than looping to completion: a mechanism is entitled to + /// refuse to drain (the `sticky` mechanism in the test plugin does exactly + /// that, and a real stream-ordered mechanism may simply not be ready), and + /// spinning on it inside `drop` would hang the caller. When the plugin + /// declines, the outstanding count stays non-zero and the caller's + /// callback table is leaked rather than freed, which is the honest + /// outcome. + fn retire_queued_releases_on_drop(&self) { + let Some(drain) = self.vtable.drain_releases else { + return; + }; + // Each pass must make progress or the loop stops. The bound is a + // belt-and-braces guard against a plugin that reports a non-zero + // retirement without ever emptying its queue. + for _ in 0..MAX_DRAIN_PASSES_ON_DROP { + if self.bridge().outstanding_releases.load(Ordering::Acquire) == 0 { + return; + } + let mut retired = 0u64; + // SAFETY: `ctx` came from this vtable and `retired` is a valid + // local. No host lock is held, which is required because the + // plugin calls back into the host from inside this call. + let status = unsafe { drain(self.vtable.ctx, u64::MAX, &raw mut retired) }; + if !status.is_ok() || retired == 0 { + return; + } + } + } + fn next_allocation_id(&self) -> u64 { // Monotonic and explicitly not pointer-derived, so a reused address // can never alias a previous allocation's identity. @@ -574,6 +719,12 @@ impl AllocatorCore { impl Drop for AllocatorCore { fn drop(&mut self) { + // Give the plugin the chance to retire what it has queued. Without + // this, an ordinary drop leaves the module's queued tally non-zero + // forever and the unload gate can never reopen — a deferral that + // never resolves is indistinguishable from a leak. + self.retire_queued_releases_on_drop(); + // Any allocation still in the map is a host-side leak, not something to // silently free: the plugin may already have been told about it. Report // through the plugin's terminal slot so the bytes are not stranded. @@ -593,11 +744,32 @@ impl Drop for AllocatorCore { if let Some(release) = self.vtable.release { // SAFETY: `ctx` came from this vtable and `release` is called // exactly once, here. The callback table and bridge are still - // alive: they are fields of `self`, which is only destructured - // after this `drop` body returns. + // alive: the context is taken below, after this call returns, so + // the plugin may touch the table from inside `release` itself. unsafe { release(self.vtable.ctx) }; } self.module.allocator_closed(); + + // Now decide the callback table's fate. The synchronous `release` + // above is *not* the last thing that can reach it: the ABI explicitly + // permits a plugin to report a queued release from its own worker + // thread, which is the whole point of deferred release. So the table + // may only be freed once no queued release can still name it. + let Some(context) = self.context.take() else { + return; + }; + if context.bridge.outstanding_releases.load(Ordering::Acquire) == 0 { + drop(context); + return; + } + // A release the plugin never retired still holds a raw pointer into + // this table. Freeing it would be a use-after-free the moment the + // plugin reports; leaking it is bounded by exactly the condition that + // already keeps the module mapped forever, and a leak is the safe + // direction — the same reasoning the plugin side applies to a block it + // cannot account for. + LEAKED_CALLBACK_TABLES.fetch_add(1, Ordering::AcqRel); + core::mem::forget(context); } } @@ -685,11 +857,17 @@ impl PluginAllocator { // arrives synchronously from inside `enqueue_release` cannot decrement // a counter that was never incremented. core.module.release_queued(); + core.bridge() + .outstanding_releases + .fetch_add(1, Ordering::AcqRel); // SAFETY: `ctx` came from this vtable; both pointers address valid // locals. No host lock is held. let status = unsafe { enqueue(core.vtable.ctx, &raw const allocation, &raw mut ticket) }; if !status.is_ok() { core.module.release_retired(); + core.bridge() + .outstanding_releases + .fetch_sub(1, Ordering::AcqRel); core.insert_allocation(ptr, record); return Err(status_to_memory_error( "enqueue_release", @@ -1314,7 +1492,6 @@ impl Drop for PluginSharedPrefix { } } -/// Open an allocator from a factory, wiring up the host callback table. /// Give back an allocator the host has decided it cannot use. /// /// `open_allocator` returning `Ok` means the plugin created state, so the host @@ -1325,13 +1502,22 @@ impl Drop for PluginSharedPrefix { /// publishes a vtable too malformed to locate `release` cannot be given /// anything back — such a plugin must not allocate state before returning. /// +/// The re-read deliberately goes through +/// [`NxmemAllocatorVtable::read_prefix_unvalidated`] rather than +/// `read_prefix`. Every memory-safety check still runs; what is skipped is +/// `validate_required`, which is the very check most refusals failed. Reading +/// through the validating entry point here would fail a second time, for the +/// same reason, and the plugin's state would be stranded — a null `allocate` +/// slot or a zero `mechanism_id` would leak rather than be handed back. +/// /// # Safety /// /// `ptr` must be null or the pointer the plugin just returned. unsafe fn abandon_allocator(ptr: *const NxmemAllocatorVtable, minor: u32) { - // SAFETY: delegated to this function's contract; `read_prefix` performs - // every null, alignment, and size check before reading a field. - let Ok(vtable) = (unsafe { NxmemAllocatorVtable::read_prefix(ptr, minor) }) else { + // SAFETY: delegated to this function's contract; the unvalidated read + // still performs every null, alignment, and size check before reading a + // field, and copies rather than borrowing. + let Ok(vtable) = (unsafe { NxmemAllocatorVtable::read_prefix_unvalidated(ptr, minor) }) else { return; }; let Some(release) = vtable.release else { @@ -1341,6 +1527,74 @@ unsafe fn abandon_allocator(ptr: *const NxmemAllocatorVtable, minor: u32) { unsafe { release(vtable.ctx) }; } +/// Holds the host's debt to the plugin until the allocator is accepted. +/// +/// From the moment `open_allocator` returns `Ok`, the plugin has created state +/// and the host owes it exactly one `release`. Between that point and the +/// construction of the [`AllocatorCore`] the host applies a series of +/// independent checks, any of which may refuse the vtable — and every one of +/// those refusals still owes the release. +/// +/// Discharging that debt with an explicit call on each early return is what +/// the previous shape did, and it is why five of the six refusal paths leaked: +/// each new check is a new opportunity to forget. A guard inverts the default, +/// so a path that forgets to think about the debt pays it anyway, and a +/// seventh check added later is covered before it is written. +struct AbandonOnDrop { + vtable: *const NxmemAllocatorVtable, + minor: u32, + armed: bool, +} + +impl AbandonOnDrop { + /// Take on the debt for a vtable the plugin has just handed over. + /// + /// # Safety + /// + /// `vtable` must be null or the pointer the plugin returned from a + /// successful `open_allocator`, and must stay valid until the guard is + /// dropped or disarmed. + unsafe fn new(vtable: *const NxmemAllocatorVtable, minor: u32) -> Self { + Self { + vtable, + minor, + armed: true, + } + } + + /// The host accepted the allocator: the debt now belongs to + /// [`AllocatorCore::drop`], which makes the matching `release` call. + fn accepted(mut self) { + self.armed = false; + } +} + +impl Drop for AbandonOnDrop { + fn drop(&mut self) { + if !self.armed { + return; + } + // SAFETY: the pointer is the one the plugin returned, as promised by + // `AbandonOnDrop::new`; `abandon_allocator` re-reads it defensively + // and calls nothing but `release`. No host lock is held: this runs on + // an early return from `open_allocator`, which takes none. + unsafe { abandon_allocator(self.vtable, self.minor) }; + } +} + +/// Open an allocator from a factory, wiring up the host callback table. +/// +/// # The post-`Ok` obligation +/// +/// Once the plugin returns `Ok` it has created state, and the host owes it a +/// `release` — even if the host then refuses the vtable on its own terms. The +/// `Ok` is what creates the debt; the host's later opinion of the vtable does +/// not cancel it. A refusal that just returns `Err` strands plugin state and, +/// because the module's live-allocator tally gates unload, permanently +/// disables the gate for the life of the process. +/// +/// That obligation is discharged by [`AbandonOnDrop`] rather than by a call on +/// each rejection path, so a path added later is covered before it is written. pub(crate) fn open_allocator( factory: &PluginFactory, required_capability_flags: u64, @@ -1361,29 +1615,26 @@ pub(crate) fn open_allocator( // The bridge and the callback table must live at stable heap addresses // *before* the plugin sees them, and must outlive the allocator's final - // release. Boxing them here and moving the boxes into the core satisfies - // both: moving a `Box` never moves the heap allocation. - let bridge = Box::new(HostBridge { - device: factory.device(), - reclaim, - reclaim_calls: AtomicU64::new(0), - reclaim_failures: AtomicU64::new(0), - retired: Mutex::new(Vec::new()), - module: Arc::clone(&module), - }); - let callbacks = Box::new(NxmemHostCallbacks { - struct_size: core::mem::size_of::() as u32, - abi_minor: minor, - host_ctx: (&raw const *bridge) as *mut core::ffi::c_void, - request_reclaim: Some(host_request_reclaim), - release_completed: Some(host_release_completed), - }); + // release **and every queued release**. `HostCallbackContext` owns both in + // one boxed unit whose teardown `AllocatorCore::drop` controls. + let context = HostCallbackContext::new( + minor, + HostBridge { + device: factory.device(), + reclaim, + reclaim_calls: AtomicU64::new(0), + reclaim_failures: AtomicU64::new(0), + retired: Mutex::new(Vec::new()), + module: Arc::clone(&module), + outstanding_releases: AtomicU64::new(0), + }, + ); let request = NxmemOpenRequest::new( minor, factory_vtable.device, required_capability_flags, - &raw const *callbacks, + &raw const context.callbacks, ); let mut raw_allocator = core::ptr::null::(); // SAFETY: `ctx` came from the factory vtable; `request` and @@ -1400,24 +1651,19 @@ pub(crate) fn open_allocator( return Err(PluginError::call("open_allocator", status)); } + // From here on the plugin has created state and the host owes it exactly + // one `release`, whatever it decides about the vtable. The guard carries + // that debt so no rejection path below has to remember it. + // SAFETY: `raw_allocator` is whatever the plugin wrote in response to the + // successful call above, which is precisely what the guard promises to + // hand back. + let abandon = unsafe { AbandonOnDrop::new(raw_allocator, minor) }; + // SAFETY: the plugin wrote this pointer in response to the call above. // `read_prefix` null-, alignment- and size-checks it before reading any // field, and copies rather than borrowing. - let vtable = - unsafe { NxmemAllocatorVtable::read_prefix(raw_allocator, minor) }.map_err(|status| { - // SAFETY: the vtable was refused, so nothing may be called through - // it; `abandon_allocator` re-reads the prefix defensively and only - // calls `release` when the struct is well-formed enough to locate - // that slot. - unsafe { abandon_allocator(raw_allocator, minor) }; - PluginError::call("allocator vtable", status) - })?; - if let Err(status) = vtable.validate_required() { - // SAFETY: as above. Here the prefix did read, so `release` is known. - unsafe { abandon_allocator(raw_allocator, minor) }; - return Err(PluginError::call("allocator vtable", status)); - } - + let vtable = unsafe { NxmemAllocatorVtable::read_prefix(raw_allocator, minor) } + .map_err(|status| PluginError::call("allocator vtable", status))?; let device = device_key(vtable.device).ok_or_else(|| PluginError::Contract { path: module.path().display().to_string(), reason: format!( @@ -1501,6 +1747,11 @@ pub(crate) fn open_allocator( let name = unsafe { read_optional_c_string(vtable.name) } .unwrap_or_else(|| factory.name().to_string()); + // Every check has passed, so the host is taking the allocator. The debt + // for the plugin's `release` transfers from the guard to `AllocatorCore`, + // whose `Drop` makes the matching call. + abandon.accepted(); + module.allocator_opened(); let core = Arc::new(AllocatorCore { vtable, @@ -1513,8 +1764,7 @@ pub(crate) fn open_allocator( capability_flags: vtable.capability_flags, next_allocation_id: AtomicU64::new(0), live: Mutex::new(HashMap::new()), - bridge, - callbacks, + context: Some(context), module, }); diff --git a/crates/onnx-runtime-memory-host/src/error.rs b/crates/onnx-runtime-memory-host/src/error.rs index 2ee686174f..d6046b5c38 100644 --- a/crates/onnx-runtime-memory-host/src/error.rs +++ b/crates/onnx-runtime-memory-host/src/error.rs @@ -65,6 +65,20 @@ impl PluginError { status: Box::new(status), } } + + /// The concrete status code the plugin reported, when there was one. + /// + /// Exposed so a caller — a test above all — can assert on the code rather + /// than on the wording of a human-readable message. A message is written + /// for a person reading a log; matching on its text turns a rewording into + /// a false pass or a false failure, and makes a test that never checked + /// the real condition look like it did. + pub fn status_code(&self) -> Option { + match self { + Self::Call { status, .. } => status.status_code(), + _ => None, + } + } } /// Turn a failing plugin status into a [`MemoryError`]. diff --git a/crates/onnx-runtime-memory-host/src/loader.rs b/crates/onnx-runtime-memory-host/src/loader.rs index 4c59a7d022..f7b4891cea 100644 --- a/crates/onnx-runtime-memory-host/src/loader.rs +++ b/crates/onnx-runtime-memory-host/src/loader.rs @@ -222,6 +222,20 @@ impl Drop for PluginFactory { } /// A loaded plugin and the mechanisms it publishes. +/// +/// # Shutting one down +/// +/// [`MemoryPlugin::try_unload`] is the only clean shutdown. **Dropping a +/// `MemoryPlugin` is not**: a drop cannot report a refusal, so when live +/// objects remain it has no honest option but to keep the module mapped +/// forever. That is the safe direction — unmapping a module whose code a +/// queued release is about to run is a use-after-free that no amount of +/// host-side care can recover from — but it is a leak, and it is silent apart +/// from the [`MemoryPlugin::forced_module_leaks`] counter. +/// +/// So: call `try_unload`, and retire what it names. Reach a drop only on paths +/// where the gate is already clear, or where leaking the module is genuinely +/// the intended outcome (process teardown, for instance). #[derive(Debug)] pub struct MemoryPlugin { /// Declared before `module` so factories release before the module (and @@ -230,6 +244,9 @@ pub struct MemoryPlugin { module: Arc, } +/// How many times a `MemoryPlugin` was dropped while its unload gate was shut. +static FORCED_MODULE_LEAKS: AtomicU64 = AtomicU64::new(0); + impl MemoryPlugin { /// Load a plugin using this host's current supported version range. pub fn load(path: impl AsRef) -> Result { @@ -353,7 +370,12 @@ impl MemoryPlugin { /// Unload is refused while **either** side still reports live objects. On /// refusal the plugin is handed back, so a caller can retire the /// outstanding work and try again — this is a deferral, not a leak. - pub fn try_unload(self) -> Result<(), UnloadRejection> { + /// + /// This is the only clean shutdown. Dropping the plugin instead skips + /// nothing — the same gate is evaluated in [`Drop`] — but a drop has no + /// way to report a refusal, so it can only keep the module mapped. See + /// the type-level documentation. + pub fn try_unload(mut self) -> Result<(), UnloadRejection> { let host = self.module.host_live_counts(); // Ask the plugin first, unconditionally, so every rejection carries // both tallies. A refusal that reports only one side leaves the caller @@ -405,10 +427,14 @@ impl MemoryPlugin { // Both sides agree nothing is live. Release the factories, then check // that nothing else still pins the module before letting it unmap. - let Self { - factories, module, .. - } = self; - drop(factories); + self.factories.clear(); + + // Hold the module across `self`'s own drop. `Drop` re-checks the gate + // — it passes, since nothing is live — and releases this handle's + // reference; the clone below is what keeps the library mapped until + // the deliberate unmap attempt. + let module = Arc::clone(&self.module); + drop(self); match Arc::try_unwrap(module) { Ok(module) => { @@ -432,6 +458,52 @@ impl MemoryPlugin { } } } + + /// How many times a plugin has been dropped while its unload gate was shut. + /// + /// Each one is a module this process deliberately kept mapped rather than + /// risk unmapping code a queued release still needs. Process-wide, because + /// unloading a module is a process-wide act. + /// + /// A non-zero value is not a crash and not corruption; it means some + /// caller took the drop path instead of [`Self::try_unload`] and the + /// module can never be reclaimed. Treat it as a leak to be fixed at the + /// call site. + pub fn forced_module_leaks() -> u64 { + FORCED_MODULE_LEAKS.load(Ordering::Acquire) + } +} + +impl Drop for MemoryPlugin { + fn drop(&mut self) { + let host = self.module.host_live_counts(); + // A module that cannot answer is treated exactly as `try_unload` + // treats it: unmapping it would be a guess, so it is not unmapped. + let plugin_total = match self.module.unload_report() { + Ok(report) => report.total(), + Err(_) => u64::MAX, + }; + if host.total() == 0 && plugin_total == 0 { + // The gate is open. Ordinary field drops run: the factories + // release, this handle's `Arc` goes, and if it was the last one + // the library unmaps. That is the intended outcome. + return; + } + + // The gate is shut, and a drop has no channel to refuse through. The + // only two options are to unmap a module whose code a live object may + // still enter, or to keep it mapped forever; nothing here can retire + // the outstanding work on the caller's behalf. + // + // Whether `dlclose` *actually* unmaps is a platform accident — glibc + // unmaps a refcount-zero DSO without `DF_1_NODELETE`, macOS commonly + // declines to — so relying on it not happening is not a safety + // argument. Leaking one strong reference makes the outcome the same + // everywhere: `library` is never dropped, so `dlclose` is never + // called. + FORCED_MODULE_LEAKS.fetch_add(1, Ordering::AcqRel); + core::mem::forget(Arc::clone(&self.module)); + } } /// Why an unload was refused or deferred. diff --git a/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs b/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs index 1d1d71c8d7..ef60c182d8 100644 --- a/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs +++ b/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs @@ -34,12 +34,14 @@ use std::sync::atomic::{AtomicU64, Ordering}; use onnx_runtime_memory_abi::{ NXMEM_ABI_VERSION_MAJOR, NXMEM_ABI_VERSION_MINOR, NXMEM_CAP_ALLOCATOR, NXMEM_CAP_DEFERRED_RELEASE, NXMEM_CAP_SHARED_MAPPING, NXMEM_CAP_STRUCTURED_RELEASE, - NXMEM_CAP_VIRTUAL_BACKING, NxmemVersionRange, + NXMEM_CAP_VIRTUAL_BACKING, NxmemAllocatorVtable, NxmemStatusCode, NxmemVersionRange, }; use onnx_runtime_memory_api::{ AllocationCommitRange, AllocationReleaseOutcome, DeviceAllocator, DeviceKey, MemoryError, }; -use onnx_runtime_memory_host::{HostReclaim, MemoryPlugin, PluginAllocator, PluginError}; +use onnx_runtime_memory_host::{ + AllocatorCore as PluginAllocatorCore, HostReclaim, MemoryPlugin, PluginAllocator, PluginError, +}; // ─── locating the plugin ──────────────────────────────────────────────────── @@ -249,8 +251,16 @@ fn a_minor_0_mechanism_works_under_a_minor_1_host() { /// **Short struct.** /// /// The `short-struct` mechanism publishes an allocator vtable whose -/// `struct_size` is smaller than the baseline prefix. The host must refuse it -/// *before* reading any function pointer out of it. +/// `struct_size` is smaller than the baseline prefix, in an allocation that is +/// genuinely that short. The host must refuse it on the size check, *before* +/// reading any function pointer out of it. +/// +/// The assertion is on the status code and on the two sizes, not on the +/// wording. Matching the word "short" in the message cannot distinguish this +/// from a refusal that happened for a completely different reason: the status +/// code's own name contains it, so a host that skipped the size check entirely +/// and tripped over a null `allocate` slot would produce a message containing +/// "short" too, and the test would stay green while testing nothing. #[test] fn a_short_allocator_vtable_is_refused_before_any_slot_is_read() { let _serial = serial(); @@ -261,12 +271,217 @@ fn a_short_allocator_vtable_is_refused_before_any_slot_is_read() { .open(NXMEM_CAP_ALLOCATOR, None) .expect_err("an undersized vtable must be refused"); + assert_eq!( + error.status_code(), + Some(NxmemStatusCode::ShortStruct), + "the refusal must be the size check, not something downstream of it: {error}" + ); + + // The message still has to name the real numbers, because a human reading + // a log needs to know *which* struct disagreed and by how much. The + // declared size is the fixture's own declaration; the required size is the + // baseline prefix this host was built against. + let declared = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 / 2; + let required = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0; let text = error.to_string(); assert!( - text.contains("struct_size") - || text.to_lowercase().contains("short") - || text.contains("small"), - "the refusal must explain that the struct is too short, got: {text}" + text.contains(&declared.to_string()) && text.contains(&required.to_string()), + "the refusal must report the declared size {declared} and the required size {required}, \ + got: {text}" + ); +} + +/// **The declared prefix bounds the read.** +/// +/// `poisoned-tail` publishes a minor-0 allocator vtable out of a full-size +/// heap allocation whose bytes past the declaration are `0xAB` — including a +/// populated `release_allocation` slot the declaration excludes. +/// +/// A host that honours `struct_size` reads the baseline prefix and zeroes the +/// rest, so the structured-release slot comes back absent. A host that reads +/// the whole struct *and* skips the level clamp comes back holding +/// `0xABAB_ABAB_ABAB_ABAB` in a function-pointer slot — and would call it. +/// +/// The assertion is on the raw post-negotiation slot rather than on release +/// behaviour, because the structured-release path short-circuits on +/// `abi_minor < 1` before it ever looks at the slot: asserting through +/// behaviour would pass no matter what the slot contained. +#[test] +fn a_poisoned_tail_never_reaches_the_hosts_view_of_the_vtable() { + let _serial = serial(); + let plugin = load(); + let allocator = plugin + .factory("poisoned-tail") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, None) + .expect("a minor-0 vtable that declares its own size honestly is acceptable"); + + assert!( + !allocator.core().publishes_structured_release_slot(), + "a slot past the sender's declared size must never survive into the host's copy" + ); + + // And the allocator is otherwise entirely ordinary, so the refusal above + // is about the tail and nothing else. + let ptr = allocator + .allocate(4096, 256) + .expect("the mechanism allocates"); + // SAFETY: `ptr` came from this allocator with exactly these parameters. + unsafe { allocator.deallocate(ptr, 4096, 256) }; +} + +/// **A refused vtable still owes the plugin a release.** +/// +/// `bad-tier` creates real allocator state, returns `Ok`, and publishes a +/// device tier code no host knows. The plugin has therefore done everything +/// right up to the point where the host applies its *own* policy and refuses. +/// +/// The `Ok` is what creates the debt: from the plugin's side an allocator +/// exists and will exist until somebody calls `release`. A host that just +/// returns `Err` strands it — and because the module's live-allocator tally is +/// what gates unload, a single refusal like this permanently disables the gate +/// for the life of the process. +/// +/// Asserting through `try_unload` rather than through a counter is deliberate: +/// the gate is the thing that actually breaks, so that is what the test +/// exercises. +#[test] +fn a_vtable_the_host_refuses_after_ok_is_still_released() { + let _serial = serial(); + let plugin = load(); + + let error = plugin + .factory("bad-tier") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, None) + .expect_err("a tier code this host does not know must be refused"); + assert!( + matches!(error, PluginError::Contract { .. }), + "the refusal is the host's own policy, not a plugin failure: {error}" + ); + + // The refusal must have left the module exactly as it found it. If the + // plugin's state were stranded the report would still show one live + // allocator and this would come back as a rejection. + plugin + .try_unload() + .expect("a refused open must leave no plugin state behind"); +} + +/// **A vtable the host cannot even parse is still released.** +/// +/// The sibling of the test above, and the harder half. There the host refused +/// a vtable it had read successfully, so it could find `release` the ordinary +/// way. Here `read_prefix` itself refuses — a required slot is null — so the +/// host has no validated view to look in, and the only way to honour the +/// post-`Ok` debt is to read the vtable *again* without validating it and call +/// `release` out of that. +/// +/// This is the path that decides whether the abandon route works at all. A +/// re-read that validates is deterministic over the same bytes: it fails +/// identically, returns before reaching `release`, and the plugin's state is +/// stranded for the life of the process with no error anywhere. +#[test] +fn a_vtable_the_host_cannot_parse_is_still_released() { + let _serial = serial(); + let plugin = load(); + + let error = plugin + .factory("missing-slot") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, None) + .expect_err("a vtable missing a required slot must be refused"); + assert_eq!( + error.status_code(), + Some(NxmemStatusCode::ShortStruct), + "a null required slot is reported as a short struct, the same code the header \ + contract uses for a vtable that cannot supply what the level requires: {error}" + ); + + // The plugin created real state and took its own reference before + // returning `Ok`. If the host did not reach `release`, that reference is + // still held and this comes back as a rejection naming a live allocator. + plugin + .try_unload() + .expect("a vtable the host could not parse must still have been released"); +} + +/// **A retirable queued release does not leak anything.** +/// +/// The counterpart to the test above: leaking is the fallback, not the policy. +/// `lazy` retires what it is asked to, so dropping its allocator drains the +/// queue first and the table is freed normally — the module's queued tally +/// comes back to zero and the gate reopens. +#[test] +fn dropping_an_allocator_drains_what_the_plugin_will_retire() { + let _serial = serial(); + let plugin = load(); + let allocator = plugin + .factory("lazy") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR | NXMEM_CAP_DEFERRED_RELEASE, None) + .expect("the mechanism opens"); + + let ptr = allocator.allocate(4096, 64).expect("allocation"); + // SAFETY: a live allocation from this allocator with matching parameters. + let _ticket = unsafe { allocator.enqueue_release(ptr, 4096, 64) }.expect("queued release"); + assert_eq!(allocator.core().outstanding_releases(), 1); + + let before = PluginAllocatorCore::leaked_callback_tables(); + drop(allocator); + assert_eq!( + PluginAllocatorCore::leaked_callback_tables(), + before, + "a release the plugin is willing to retire must not cost a leaked table" + ); + + // And the gate really did reopen, which is the point of draining. + plugin + .try_unload() + .expect("a drained allocator leaves nothing outstanding"); +} + +/// **Dropping a plugin with live objects must not unmap it.** +/// +/// `try_unload` consumes the plugin, so every early return, `?` and unwind +/// reaches `drop` instead — and `drop` has no channel to refuse through. Its +/// only safe option is to keep the module mapped. +/// +/// Whether `dlclose` would *actually* unmap is a platform accident: glibc +/// unmaps a refcount-zero DSO without `DF_1_NODELETE`, macOS commonly +/// declines. Relying on the platform declining is not a safety property, so +/// the assertion is on the host's own decision to leak rather than on whether +/// the mapping survived. +#[test] +fn dropping_a_plugin_with_live_objects_keeps_the_module_mapped() { + let _serial = serial(); + let before = MemoryPlugin::forced_module_leaks(); + + { + let plugin = load(); + let allocator = plugin + .factory("eager") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, None) + .expect("the mechanism opens"); + // The allocator is still live when the plugin goes out of scope, and + // nothing here calls `try_unload`. + drop(plugin); + assert_eq!( + MemoryPlugin::forced_module_leaks() - before, + 1, + "dropping a plugin whose allocator is still live must keep the module mapped" + ); + drop(allocator); + } + + // An idle plugin, by contrast, unloads cleanly and costs nothing. + let after_live_drop = MemoryPlugin::forced_module_leaks(); + load().try_unload().expect("an idle plugin unloads"); + assert_eq!( + MemoryPlugin::forced_module_leaks(), + after_live_drop, + "a clean unload must not be counted as a forced leak" ); } @@ -985,7 +1200,7 @@ fn every_factory_is_released_exactly_once() { let before = releases(&observer); let plugin = load(); let count = plugin.factories().len() as u64; - assert_eq!(count, 8, "the test plugin publishes eight mechanisms"); + assert_eq!(count, 12, "the test plugin publishes twelve mechanisms"); plugin.try_unload().expect("an idle plugin unloads"); assert_eq!( diff --git a/crates/onnx-runtime-memory-host/tests/nxmem_abi_callback_pinning.rs b/crates/onnx-runtime-memory-host/tests/nxmem_abi_callback_pinning.rs new file mode 100644 index 0000000000..ade01f684e --- /dev/null +++ b/crates/onnx-runtime-memory-host/tests/nxmem_abi_callback_pinning.rs @@ -0,0 +1,52 @@ +//! The callback table's pinning, in a process of its own. +//! +//! Like the unload gate's terminal case, this uses the `sticky` mechanism, +//! which deliberately never retires a queued release. That permanently poisons +//! the module's process-wide counters, so any two such tests sharing a process +//! would read each other's residue and the exact assertions below would have to +//! be softened into inequalities — which is how a test stops testing anything. +//! Hence a binary of its own. + +use onnx_runtime_memory_abi::{NXMEM_CAP_ALLOCATOR, NXMEM_CAP_DEFERRED_RELEASE}; +use onnx_runtime_memory_api::DeviceAllocator; +use onnx_runtime_memory_host::{AllocatorCore as PluginAllocatorCore, MemoryPlugin}; + +#[path = "support/testplugin.rs"] +mod testplugin; + +/// **A queued release keeps the host's callback table alive.** +/// +/// The plugin holds a raw pointer to the callback table, captured at open, and +/// the contract explicitly permits it to report a completion from one of its +/// own threads — that is the entire reason deferred release exists. So the +/// table may only be freed once no queued release can still name it. +/// +/// `sticky` never retires anything, so dropping its allocator is precisely the +/// case where freeing would be wrong. The host must leak the table instead, +/// and say so. +#[test] +fn dropping_an_allocator_with_a_queued_release_leaks_its_callback_table() { + let plugin = MemoryPlugin::load(testplugin::testplugin_path()).expect("the test plugin loads"); + let allocator = plugin + .factory("sticky") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR | NXMEM_CAP_DEFERRED_RELEASE, None) + .expect("the mechanism opens"); + + let ptr = allocator.allocate(4096, 64).expect("allocation"); + // SAFETY: a live allocation from this allocator with matching parameters. + let _ticket = unsafe { allocator.enqueue_release(ptr, 4096, 64) }.expect("queued release"); + assert_eq!( + allocator.core().outstanding_releases(), + 1, + "the queued release is what pins the table" + ); + + let before = PluginAllocatorCore::leaked_callback_tables(); + drop(allocator); + assert_eq!( + PluginAllocatorCore::leaked_callback_tables() - before, + 1, + "a table a queued release can still name must be leaked, not freed" + ); +} diff --git a/crates/onnx-runtime-memory-host/tests/nxmem_abi_plugin_gate.rs b/crates/onnx-runtime-memory-host/tests/nxmem_abi_plugin_gate.rs new file mode 100644 index 0000000000..b5ddbc9a56 --- /dev/null +++ b/crates/onnx-runtime-memory-host/tests/nxmem_abi_plugin_gate.rs @@ -0,0 +1,96 @@ +//! The *plugin's* half of the unload gate, in a process of its own. +//! +//! The gate has two independent halves: what the host knows it is holding, and +//! what the plugin says it still owns. The host's half is easy to exercise +//! because the host counts it. The plugin's half is the one that matters most +//! and is the hardest to reach, because in every ordinary scenario the host is +//! still holding something too — so the host's check fires first and the +//! plugin's check is never the reason for the refusal. +//! +//! `self-retaining` closes that gap. It takes a reference to its own allocator +//! state at open and never drops it, modelling a plugin that keeps an +//! allocator alive for its own worker thread after the host has let go. Once +//! the host drops its handle the host's tallies are all zero, so the *only* +//! thing that can refuse the unload is the plugin's report. +//! +//! That state is never released, so the module's process-wide counters stay +//! poisoned for the rest of the process. Hence a binary of its own. + +use onnx_runtime_memory_abi::NXMEM_CAP_ALLOCATOR; +use onnx_runtime_memory_host::MemoryPlugin; + +#[path = "support/testplugin.rs"] +mod testplugin; + +/// **Unload is refused on the plugin's word alone.** +/// +/// With the host's own tallies at zero, a host that trusts only itself would +/// happily `dlclose` a module that still owns an allocator — and the plugin's +/// worker thread would then run unmapped code. +#[test] +fn unload_is_refused_when_only_the_plugin_still_owns_something() { + let plugin = MemoryPlugin::load(testplugin::testplugin_path()).expect("the test plugin loads"); + let allocator = plugin + .factory("self-retaining") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, None) + .expect("the mechanism opens"); + + // Let go of everything the host tracks. The plugin's extra reference means + // its state survives, and `release` therefore does not tear it down. + drop(allocator); + + let rejection = plugin + .try_unload() + .expect_err("a plugin that still owns an allocator must block unload"); + + // This is the whole point: the host has nothing left to complain about, so + // if the refusal did not come from the plugin's report it did not happen. + assert_eq!( + rejection.host.allocators, 0, + "the host must have let go of its allocator: {:?}", + rejection.host + ); + assert_eq!( + rejection.host.allocations, 0, + "and of its allocations: {:?}", + rejection.host + ); + assert_eq!( + rejection.host.queued_releases, 0, + "and of its queued releases: {:?}", + rejection.host + ); + assert_eq!( + rejection.host.total(), + 0, + "so the host's half of the gate is open: {:?}", + rejection.host + ); + + assert_eq!( + rejection.report.live_allocators, 1, + "the plugin still owns the allocator, and that must be what refuses: {:?}", + rejection.report + ); + assert!( + rejection.report.total() >= 1, + "the plugin's report is the only thing keeping the module mapped: {:?}", + rejection.report + ); + + // And the refusal is honest about being unrecoverable from here: the + // plugin comes back, but nothing the host can do will retire the + // plugin-side reference, so the only exit is a drop that keeps the module + // mapped. + let leaks_before = MemoryPlugin::forced_module_leaks(); + let plugin = rejection + .into_plugin() + .expect("the refusal hands the plugin back"); + drop(plugin); + assert_eq!( + MemoryPlugin::forced_module_leaks() - leaks_before, + 1, + "a module the plugin still owns objects in must stay mapped" + ); +} diff --git a/crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs b/crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs index 50005e22b6..0aac9b9459 100644 --- a/crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs +++ b/crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs @@ -45,10 +45,14 @@ fn unload_is_refused_while_a_queued_release_has_not_retired() { let rejection = plugin .try_unload() .expect_err("a queued release must block unload"); - assert!( - rejection.host.queued_releases >= 1 || rejection.report.queued_releases >= 1, - "either side may notice first, but somebody must: host={:?} plugin={:?}", - rejection.host, + assert_eq!( + rejection.host.queued_releases, 1, + "the host's own tally must show the queued release: {:?}", + rejection.host + ); + assert_eq!( + rejection.report.queued_releases, 1, + "and so must the plugin's: {:?}", rejection.report ); let plugin = rejection @@ -57,10 +61,27 @@ fn unload_is_refused_while_a_queued_release_has_not_retired() { // Dropping the allocator lets the plugin reclaim its own storage, but the // host's queued tally is what keeps the gate shut, and it never clears. + let leaks_before = MemoryPlugin::forced_module_leaks(); drop(allocator); - let outcome = plugin.try_unload(); - assert!( - outcome.is_err(), - "the host must not unmap a module that still owes a free" + let rejection = plugin + .try_unload() + .expect_err("the host must not unmap a module that still owes a free"); + assert_eq!( + rejection.report.queued_releases, 1, + "the plugin is what still owes the free: {:?}", + rejection.report + ); + + // The refusal hands the plugin back, and nothing here can retire the + // release, so the only remaining exit is a drop — which must keep the + // module mapped rather than unmap code the free will run. + let plugin = rejection + .into_plugin() + .expect("the refusal returns the plugin"); + drop(plugin); + assert_eq!( + MemoryPlugin::forced_module_leaks() - leaks_before, + 1, + "dropping a plugin that still owes a free must keep its module mapped" ); } diff --git a/crates/onnx-runtime-memory-testplugin/src/lib.rs b/crates/onnx-runtime-memory-testplugin/src/lib.rs index 54368672b0..ffcd58a3c7 100644 --- a/crates/onnx-runtime-memory-testplugin/src/lib.rs +++ b/crates/onnx-runtime-memory-testplugin/src/lib.rs @@ -56,6 +56,12 @@ use onnx_runtime_memory_abi::{ struct ModuleCounters { live_allocators: AtomicU64, live_allocations: AtomicU64, + /// Shared prefixes committed into a live allocation. + /// + /// This is the module's referent for `NxmemUnloadReport::live_views`: a + /// window onto shared physical bytes, owned by the plugin, whose life is + /// bounded by the allocation it looks into. Incremented by + /// `commit_shared_prefix`, retired by [`retire_block_views`]. live_views: AtomicU64, live_capabilities: AtomicU64, queued_releases: AtomicU64, @@ -92,6 +98,27 @@ struct Behaviour { abi_minor: u32, /// Publish an allocator vtable that is deliberately too short. short_allocator_struct: bool, + /// Publish a minor-0 allocator vtable whose declared prefix is followed by + /// poisoned bytes inside the same allocation. + poisoned_tail_struct: bool, + /// Publish a vtable with a required slot missing, *after* creating real + /// state. + /// + /// This is the one post-`Ok` rejection the host cannot diagnose by reading + /// the vtable normally: `read_prefix` refuses it, so the host has no + /// validated view to find `release` in. It must still find `release`, or + /// the state is stranded for the life of the process. + missing_required_slot: bool, + /// Publish a device tier code no host knows, *after* creating real state. + /// + /// This is the shape of bug that tests the host's post-`Ok` obligation: the + /// plugin succeeded, so it created state and expects a `release`, and the + /// host then refuses the vtable on its own terms. + unknown_device_tier: bool, + /// Take an extra reference to the allocator state at open and never drop + /// it, modelling a plugin that keeps its allocator alive for its own + /// worker thread after the host has let go. + self_retaining: bool, /// Call the host's `request_reclaim` on every allocation, and fail the /// allocation when the host refuses. probe_host_callback: bool, @@ -111,6 +138,10 @@ impl Behaviour { capability_flags: NXMEM_CAP_ALLOCATOR | NXMEM_CAP_STRUCTURED_RELEASE, abi_minor: 1, short_allocator_struct: false, + poisoned_tail_struct: false, + missing_required_slot: false, + unknown_device_tier: false, + self_retaining: false, probe_host_callback: false, never_drain: false, quarantine_on_release: false, @@ -157,6 +188,49 @@ const MECHANISMS: &[Mechanism] = &[ ..Behaviour::base() }, }, + Mechanism { + name: "poisoned-tail", + c_name: c"poisoned-tail", + behaviour: Behaviour { + // Declares the baseline prefix while its allocation holds a + // populated, wrong-for-this-level `release_allocation` past that + // declaration. A host that reads past the declaration and skips + // the level clamp ends up holding it. + capability_flags: NXMEM_CAP_ALLOCATOR, + abi_minor: 0, + poisoned_tail_struct: true, + ..Behaviour::base() + }, + }, + Mechanism { + name: "missing-slot", + c_name: c"missing-slot", + behaviour: Behaviour { + // Real state, then a vtable the host cannot validate. The host has + // to reach `release` through an unvalidated read or lose the state. + capability_flags: NXMEM_CAP_ALLOCATOR, + missing_required_slot: true, + ..Behaviour::base() + }, + }, + Mechanism { + name: "bad-tier", + c_name: c"bad-tier", + behaviour: Behaviour { + capability_flags: NXMEM_CAP_ALLOCATOR, + unknown_device_tier: true, + ..Behaviour::base() + }, + }, + Mechanism { + name: "self-retaining", + c_name: c"self-retaining", + behaviour: Behaviour { + capability_flags: NXMEM_CAP_ALLOCATOR, + self_retaining: true, + ..Behaviour::base() + }, + }, Mechanism { name: "callback-probe", c_name: c"callback-probe", @@ -218,6 +292,10 @@ pub const MECHANISM_NAMES: &[&str] = &[ "eager", "lazy", "short-struct", + "poisoned-tail", + "missing-slot", + "bad-tier", + "self-retaining", "callback-probe", "legacy-1-0", "quarantining", @@ -233,6 +311,12 @@ struct Block { bytes: usize, align: usize, committed: usize, + /// Shared prefixes currently committed into this block. + /// + /// Each one is a plugin-side *view*: a window onto shared physical bytes + /// that lives inside this allocation and dies with it. This is what + /// `NxmemUnloadReport::live_views` counts. + views: u64, } #[derive(Debug, Clone, Copy)] @@ -366,6 +450,22 @@ fn free_block(block: Block) { unsafe { dealloc(block.address as *mut u8, layout) }; } +/// Retire the views a block carried, at the moment it stops being live. +/// +/// A view's life is bounded by the allocation it looks into: once the host has +/// handed the allocation back, the window into it is gone whether or not the +/// physical bytes have been freed yet. Retiring here rather than in +/// [`free_block`] keeps the count honest for a *deferred* release, where the +/// allocation leaves the live set long before the bytes do. +fn retire_block_views(block: &Block) { + if block.views == 0 { + return; + } + counters() + .live_views + .fetch_sub(block.views, Ordering::AcqRel); +} + /// # Safety /// /// `ctx` must be a live `AllocatorState` pointer produced by `open_allocator`. @@ -492,6 +592,7 @@ unsafe fn allocate_inner( bytes, align, committed: committed.min(bytes), + views: 0, }; match state.inner.lock() { Ok(mut inner) => { @@ -552,6 +653,7 @@ unsafe extern "C" fn plugin_deallocate( "testplugin: that allocation id is not live on this mechanism", ); }; + retire_block_views(&block); let unmapped = block.committed as u64; free_block(block); counters().live_allocations.fetch_sub(1, Ordering::AcqRel); @@ -603,6 +705,7 @@ unsafe extern "C" fn plugin_release_allocation( return NxmemStatus::ok(); }; + retire_block_views(&block); let unmapped = block.committed as u64; let bytes = block.bytes as u64; @@ -674,6 +777,7 @@ unsafe extern "C" fn plugin_enqueue_release( "testplugin: that allocation id is not live on this mechanism", ); }; + retire_block_views(&block); inner.queue.push(QueuedRelease { ticket, allocation_id: allocation.allocation_id, @@ -742,6 +846,7 @@ unsafe extern "C" fn plugin_drain_releases( bytes: queued.bytes, align: queued.align, committed: queued.bytes, + views: 0, }); counters().live_allocations.fetch_sub(1, Ordering::AcqRel); counters().queued_releases.fetch_sub(1, Ordering::AcqRel); @@ -851,6 +956,7 @@ fn release_state(state: &AllocatorState) { // Reclaim anything still outstanding rather than leaking it. if let Ok(mut inner) = state.inner.lock() { for (_, block) in inner.blocks.drain() { + retire_block_views(&block); free_block(block); counters().live_allocations.fetch_sub(1, Ordering::AcqRel); } @@ -861,6 +967,7 @@ fn release_state(state: &AllocatorState) { bytes: entry.bytes, align: entry.align, committed: entry.bytes, + views: 0, }); counters().live_allocations.fetch_sub(1, Ordering::AcqRel); counters().queued_releases.fetch_sub(1, Ordering::AcqRel); @@ -879,6 +986,7 @@ fn release_state(state: &AllocatorState) { bytes: prefix.bytes, align: PREFIX_ALIGN, committed: prefix.bytes, + views: 0, }); } } @@ -1256,6 +1364,7 @@ unsafe extern "C" fn plugin_release_shared_prefix( bytes: prefix.bytes, align: PREFIX_ALIGN, committed: prefix.bytes, + views: 0, }); } NxmemStatus::ok() @@ -1361,6 +1470,11 @@ unsafe extern "C" fn plugin_commit_shared_prefix( ); } block.committed = block.committed.max(end); + // A committed prefix is a live plugin-side view into this + // allocation. It is exactly the object `live_views` names, and + // it retires when the allocation does. + block.views += 1; + counters().live_views.fetch_add(1, Ordering::AcqRel); prefix.bytes as u64 } Err(_) => { @@ -1513,36 +1627,171 @@ unsafe extern "C" fn plugin_open_allocator( .as_ref() .map_or(core::ptr::null(), |shared| &raw const **shared); + if behaviour.missing_required_slot { + // The state above is real, the refcount is held, and the plugin is + // about to return `Ok`. `release` stays populated precisely + // because it is the one slot the host is still entitled to call on + // a vtable it has refused. + state.allocator_vtable.allocate = None; + } + if behaviour.unknown_device_tier { + // The state above is real and the plugin is about to return `Ok`, + // so the host owes it a `release` — even though the host is about + // to refuse this vtable on a tier code it has never heard of. + state.allocator_vtable.device.tier = UNKNOWN_DEVICE_TIER; + } + if behaviour.self_retaining { + // A reference the host never learns about and can never drop. + state.refcount.fetch_add(1, Ordering::AcqRel); + } + counters().live_allocators.fetch_add(1, Ordering::AcqRel); + let published: *const NxmemAllocatorVtable = if behaviour.poisoned_tail_struct { + // Publish out of a separate allocation whose tail is poisoned. + // The state itself is entirely ordinary, so the allocator behaves + // normally once the host has read the prefix correctly — the only + // thing under test is what the host does with the bytes past the + // declaration. + poisoned_tail_allocator_vtable(mechanism_id, request.device, allocator_ctx) + .cast::() + } else { + &raw const *state.allocator_vtable + }; // SAFETY: checked non-null above. - unsafe { *allocator_out = &raw const *state.allocator_vtable }; + unsafe { *allocator_out = published }; NxmemStatus::ok() }) } -/// A deliberately undersized allocator vtable. +/// A vtable published in an allocation sized to its own declaration. +/// +/// The point is that the memory really *ends* where `struct_size` says it +/// does. A full-size static that merely lies about its size cannot exhibit an +/// out-of-bounds read at all: the bytes past the declaration are valid, so a +/// host that ignores the bound reads garbage rather than tripping a +/// sanitiser, and a test named for the over-read proves nothing. /// -/// `struct_size` claims fewer bytes than even the baseline prefix needs, so -/// the host must refuse it before reading a single function pointer. It is -/// stateless and `'static`, so refusing it leaks nothing. -struct ShortVtable(NxmemAllocatorVtable); +/// Leaked deliberately. The host is expected to refuse or clamp these, and a +/// refusal path must not depend on the fixture staying alive; one leaked +/// allocation per process is a fair price for a fixture that can actually +/// fail. +struct TruncatedVtable { + ptr: *const NxmemAllocatorVtable, +} -// SAFETY: immutable after construction and holding only null pointers. -unsafe impl Send for ShortVtable {} +// SAFETY: the pointed-to bytes are written once, before the pointer escapes, +// and never mutated afterwards. +unsafe impl Send for TruncatedVtable {} // SAFETY: as above. -unsafe impl Sync for ShortVtable {} +unsafe impl Sync for TruncatedVtable {} + +/// Leak `bytes` of memory aligned for an allocator vtable, filled with `fill`. +/// +/// `bytes` is the *exact* allocation size, so reading past it is a genuine +/// heap over-read that ASan, Miri or `MallocScribble` can see. +fn leak_vtable_bytes(bytes: usize, fill: u8) -> *mut u8 { + let layout = Layout::from_size_align(bytes, align_of::()) + .expect("the fixture layout is a compile-time constant and is valid"); + // SAFETY: `bytes` is non-zero for every caller below. + let raw = unsafe { alloc(layout) }; + assert!( + !raw.is_null(), + "testplugin: could not allocate the {bytes}-byte vtable fixture" + ); + // SAFETY: `raw` owns exactly `bytes` writable bytes. + unsafe { core::ptr::write_bytes(raw, fill, bytes) }; + raw +} +/// A vtable whose allocation is genuinely too short for even the baseline +/// prefix. +/// +/// `struct_size` declares half the baseline, and the allocation is exactly +/// that size. The host must refuse it on the size check, before it reads a +/// single function pointer — and if it ever stops doing so, the read runs off +/// the end of a real heap allocation. fn short_allocator_vtable() -> *const NxmemAllocatorVtable { - static VTABLE: OnceLock = OnceLock::new(); - &VTABLE + static VTABLE: OnceLock = OnceLock::new(); + VTABLE .get_or_init(|| { - let mut vtable = NxmemAllocatorVtable::zeroed(); - vtable.struct_size = (NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 as u32) / 2; - ShortVtable(vtable) + let declared = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 / 2; + let raw = leak_vtable_bytes(declared, 0); + // Only `struct_size` and `abi_minor` fit; that is the point. + // SAFETY: `raw` owns at least eight writable bytes and both fields + // are `u32` at offsets 0 and 4 of every version of this struct. + unsafe { + core::ptr::write_unaligned(raw.cast::(), declared as u32); + core::ptr::write_unaligned(raw.cast::().add(1), 0); + } + TruncatedVtable { + ptr: raw.cast::(), + } }) - .0 + .ptr } +/// A minor-0 vtable followed by poisoned bytes *inside the same allocation*. +/// +/// The allocation is a full `size_of::()`, but only the +/// baseline prefix is written; everything past it is `0xAB`. A host that +/// honours the declared size reads the prefix and zeroes the rest, so +/// `release_allocation` comes back null. A host that reads the whole struct +/// *and* skips the level clamp comes back holding `0xABAB_ABAB_ABAB_ABAB` in +/// a function-pointer slot — deterministically, on every platform, with no +/// dependence on what happened to be adjacent in the heap. +/// +/// This is the value-observable half of the bound. The pure-bound half is not +/// observable through any value here: `release_allocation` is the only field +/// past the baseline prefix, and the clamp nulls it independently, so +/// destroying only the bound changes no value the host can see. It is a +/// silent over-read, which is why the direct `copy_prefix` unit test in the +/// ABI crate exists alongside this. +fn poisoned_tail_allocator_vtable( + mechanism_id: u64, + device: NxmemDeviceId, + ctx: *mut c_void, +) -> *mut u8 { + let full = size_of::(); + let declared = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0; + debug_assert!( + declared < full, + "the baseline prefix must be a strict prefix" + ); + let raw = leak_vtable_bytes(full, 0xAB); + + let mut prefix = NxmemAllocatorVtable::zeroed(); + prefix.struct_size = declared as u32; + prefix.abi_minor = 0; + prefix.mechanism_id = mechanism_id; + prefix.device = device; + prefix.capability_flags = NXMEM_CAP_ALLOCATOR; + prefix.ctx = ctx; + prefix.name = POISONED_TAIL_NAME.as_ptr().cast(); + prefix.allocate = Some(plugin_allocate); + prefix.deallocate = Some(plugin_deallocate); + prefix.retain = Some(plugin_retain); + prefix.release = Some(plugin_release); + // Deliberately populated even though the declaration excludes it: the + // bytes the host must not read have to be *wrong*, not merely absent. + prefix.release_allocation = Some(plugin_release_allocation); + + // SAFETY: `raw` owns `full` writable bytes and `declared <= full`, so + // copying the first `declared` bytes of a `full`-sized value is in bounds + // at both ends. The tail keeps its poison. + unsafe { + core::ptr::copy_nonoverlapping((&raw const prefix).cast::(), raw, declared); + } + raw +} + +const POISONED_TAIL_NAME: &std::ffi::CStr = c"poisoned-tail"; + +/// A device tier code no released host knows. +/// +/// Chosen far above anything the tier enum will plausibly reach, so a host +/// that grows new tiers still rejects it. +const UNKNOWN_DEVICE_TIER: u32 = 4242; + unsafe extern "C" fn plugin_factory_release(_ctx: *mut c_void) { // Factories are `&'static Mechanism`, so there is nothing to free. The // slot still exists because the contract requires an explicit release for diff --git a/docs/memory/MEMORY_PLUGIN_ABI.md b/docs/memory/MEMORY_PLUGIN_ABI.md index 62f6c0ae73..23d67b8fdc 100644 --- a/docs/memory/MEMORY_PLUGIN_ABI.md +++ b/docs/memory/MEMORY_PLUGIN_ABI.md @@ -142,8 +142,38 @@ ticket. Until that ticket retires: - the host's callback table stays alive, because `release_completed` will still be called through it. -This is what "deferred release keeps the plugin module and provider/context -pinned until callbacks and frees retire" means concretely. +The pinning is at two different granularities and both matter. An +`Arc` keeps the plugin's *code* mapped; it says nothing about the +host's *callback context*, which is a separate heap object the plugin captured +a raw pointer to at open. A queued release holds both, so an allocator dropped +with releases still outstanding drains what the plugin is willing to retire +and, for anything left, deliberately leaks its callback table rather than free +memory a plugin thread may still write to. `AllocatorCore::leaked_callback_tables` +counts those; a non-zero value is a leak to be fixed at the call site, not a +crash. + +**Not yet implemented:** the *provider/context* half. No execution provider is +wired to this ABI in this phase, so "deferred release keeps the provider and +its context pinned" has nothing to pin. Only the intra-boundary half — module +and callback table — is implemented and tested here. + +## Unloading a plugin + +`MemoryPlugin::try_unload` is the only clean shutdown. It consumes the plugin, +asks both sides what is still live, and hands the plugin back on refusal so the +caller can retire the outstanding work and try again. + +**Dropping a `MemoryPlugin` is not a clean shutdown.** Because `try_unload` +consumes `self`, every early return, `?` and unwind reaches `Drop` instead. A +drop has no channel to refuse through, so when the gate is shut its only two +options are to unmap a module whose code a live object may still enter, or to +keep it mapped forever. It keeps it mapped, and counts the event in +`MemoryPlugin::forced_module_leaks`. + +Do not rely on the platform to save you here. Whether `dlclose` actually +unmaps is an accident of the loader: glibc unmaps a refcount-zero DSO that is +not marked `DF_1_NODELETE`, and Rust `cdylib`s are not; macOS commonly declines +to unmap. "It did not crash on my machine" is not a safety property. ## Release outcomes @@ -175,7 +205,9 @@ ABA race cannot make one allocation impersonate another. ## Unload gating `NxmemQueryUnloadReadiness` reports live allocators, allocations, views, -capabilities, and queued releases. The host refuses or defers unload while any +capabilities, and queued releases. A *view* is a shared prefix committed into a +live allocation — a plugin-owned window whose life is bounded by the allocation +it looks into, and which the host has no way to count for itself. The host refuses or defers unload while any count is non-zero, and it checks its own counters too, so a rejection carries both sides' tallies and a misreporting plugin is visible rather than fatal. @@ -187,7 +219,7 @@ invites the host to unmap code that is about to run. `onnx-runtime-memory-testplugin` is a `cdylib` loaded at runtime by the ABI tests — out-of-tree in the way that matters (`dlopen`, no workspace linking), in-tree in the way that does not (it lives in the repo so it is built and -linted with everything else). It publishes eight named mechanisms rather than +linted with everything else). It publishes twelve named mechanisms rather than switching on environment variables or globals, so the tests select behaviour by name: @@ -195,7 +227,11 @@ name: | --- | --- | | `eager` | the minimal conforming mechanism: required slots only | | `lazy` | virtual backing, shared mapping, deferred release, structured release | -| `short-struct` | a vtable claiming fewer bytes than the baseline prefix | +| `short-struct` | a vtable claiming fewer bytes than the baseline prefix, in an allocation that really is that short | +| `poisoned-tail` | a minor-0 vtable whose declared prefix is followed by populated garbage inside the same allocation | +| `missing-slot` | returns `Ok` with real state, then a vtable with a required slot null: tests the release path for a vtable the host cannot even parse | +| `bad-tier` | returns `Ok` with real state, then publishes a device tier no host knows: tests the host's post-`Ok` release obligation | +| `self-retaining` | keeps a reference to its own allocator state after the host lets go, so only the plugin's report can refuse unload | | `callback-probe` | calls `request_reclaim` and fails cleanly when the host refuses | | `legacy-1-0` | built to minor 0: an older participant under a newer host | | `quarantining` | keeps residual ownership on release |