From 5f1b249e21608c6cb2824711944ba16735a5c4d6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 10:48:05 -0700 Subject: [PATCH] fix(memory): preserve binding pins and selection Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/onnx-runtime-memory-api/src/binding.rs | 112 ++++++++++++- .../tests/binding_registry.rs | 153 ++++++++++++++++++ docs/memory/MEMORY_ARCHITECTURE.md | 9 +- 3 files changed, 268 insertions(+), 6 deletions(-) diff --git a/crates/onnx-runtime-memory-api/src/binding.rs b/crates/onnx-runtime-memory-api/src/binding.rs index 5bf73d7961..8e13522480 100644 --- a/crates/onnx-runtime-memory-api/src/binding.rs +++ b/crates/onnx-runtime-memory-api/src/binding.rs @@ -488,6 +488,16 @@ struct RegistryState { struct RegistryInner { identities: Arc, state: Mutex, + #[cfg(test)] + select_after_validation_hook: Mutex>, +} + +#[cfg(test)] +#[derive(Clone, Debug)] +struct SelectAfterValidationHook { + mechanism: MechanismIdentity, + entered: Arc, + resume: Arc, } /// A narrow registry for provider/context pins and binding identity. @@ -509,6 +519,8 @@ impl BindingRegistry { inner: Arc::new(RegistryInner { identities: Arc::new(IdentitySource::new()?), state: Mutex::new(RegistryState::default()), + #[cfg(test)] + select_after_validation_hook: Mutex::new(None), }), }) } @@ -675,14 +687,21 @@ impl BindingRegistry { operation: "selecting a mechanism", }); } - self.lock_state("publishing mechanism selection")? + #[cfg(test)] + self.wait_after_select_validation(mechanism.identity); + let prior = self + .lock_state("publishing mechanism selection")? .selected .insert(mechanism.device, mechanism.identity); let published = entry.snapshot()?; if published.lifecycle != MechanismLifecycle::Active { let mut state = self.lock_state("withdrawing inactive mechanism selection")?; if state.selected.get(&mechanism.device) == Some(&mechanism.identity) { - state.selected.remove(&mechanism.device); + if let Some(prior) = prior { + state.selected.insert(mechanism.device, prior); + } else { + state.selected.remove(&mechanism.device); + } } return Err(BindingError::InactiveMechanism { mechanism: mechanism.identity, @@ -693,6 +712,20 @@ impl BindingRegistry { Ok(()) } + #[cfg(test)] + fn wait_after_select_validation(&self, mechanism: MechanismIdentity) { + let hook = self + .inner + .select_after_validation_hook + .lock() + .expect("select test hook lock") + .clone(); + if let Some(hook) = hook.filter(|hook| hook.mechanism == mechanism) { + hook.entered.wait(); + hook.resume.wait(); + } + } + pub fn bind(&self, device: DeviceKey) -> Result { let entry = { let state = self.lock_state("looking up the selected mechanism")?; @@ -1069,6 +1102,11 @@ impl MemoryBinding { /// On validation/device-loss failure the allocation is returned to the caller /// inside [`ExplicitReleaseError`] so its metadata and lifetime pins are not /// silently discarded. + /// + /// This Phase-3 path calls canonical [`DeviceAllocator::deallocate`]. It does + /// not expose mapped-attribution refunds from + /// [`DeviceAllocator::deallocate_with_unmapped`]; EP adoption must wait for + /// Phase-4 accounting reconciliation. pub fn release(&self, allocation: BoundAllocation) -> Result<(), ExplicitReleaseError> { let operation = match self.mechanism.begin_release(self.identity, &allocation) { Ok(operation) => operation, @@ -1518,8 +1556,8 @@ impl BoundSharedMapping { .create_shared_prefix(bytes)?; drop(active); Ok(BoundSharedPrefix { - binding: self.binding.clone(), prefix, + binding: self.binding.clone(), }) } @@ -1577,10 +1615,15 @@ impl BoundSharedMapping { } /// Shared physical-prefix handle pinned to one binding. +/// +/// Physical teardown remains driven by dropping this handle, not by +/// [`BindingRegistry`] lifecycle transitions. The prefix is dropped before its +/// binding pin so its provider context is still alive during teardown. Phase 4 +/// owns stream ordering and partial-failure-safe physical release. #[derive(Debug)] pub struct BoundSharedPrefix { - binding: MemoryBinding, prefix: Box, + binding: MemoryBinding, } impl BoundSharedPrefix { @@ -1617,3 +1660,64 @@ fn validate_binding_identity( } Ok(()) } + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Barrier}; + use std::thread; + + use crate::{BindingResource, HostAllocator}; + + use super::*; + + #[test] + fn failed_select_restores_the_prior_healthy_selection() { + let registry = BindingRegistry::new().unwrap(); + let context = registry + .register_provider_context(DeviceKey::HOST, Arc::new(()) as Arc) + .unwrap(); + let authority = registry + .register_authority(DeviceKey::HOST, Arc::new(()) as Arc) + .unwrap(); + let prior = registry + .register_allocator( + context, + authority, + Arc::new(HostAllocator) as Arc, + ) + .unwrap(); + let candidate = registry + .register_allocator( + context, + authority, + Arc::new(HostAllocator) as Arc, + ) + .unwrap(); + let entered = Arc::new(Barrier::new(2)); + let resume = Arc::new(Barrier::new(2)); + *registry.inner.select_after_validation_hook.lock().unwrap() = + Some(SelectAfterValidationHook { + mechanism: candidate.identity, + entered: Arc::clone(&entered), + resume: Arc::clone(&resume), + }); + + let selecting_registry = registry.clone(); + let selecting = thread::spawn(move || selecting_registry.select(candidate)); + entered.wait(); + registry.retire(candidate).unwrap(); + resume.wait(); + + let error = selecting.join().unwrap().unwrap_err(); + assert!(matches!( + error, + BindingError::InactiveMechanism { + mechanism, + lifecycle: MechanismLifecycle::Retired, + operation: "selecting a mechanism", + } if mechanism == candidate.identity + )); + let selected = registry.bind(DeviceKey::HOST).unwrap(); + assert_eq!(selected.identity().mechanism(), prior.identity); + } +} diff --git a/crates/onnx-runtime-memory-api/tests/binding_registry.rs b/crates/onnx-runtime-memory-api/tests/binding_registry.rs index 1893b1b1d2..5033fee73b 100644 --- a/crates/onnx-runtime-memory-api/tests/binding_registry.rs +++ b/crates/onnx-runtime-memory-api/tests/binding_registry.rs @@ -50,6 +50,111 @@ impl SharedDevicePrefix for TestPrefix { } } +#[derive(Debug)] +struct ContextDropFlag { + released: Arc, +} + +impl Drop for ContextDropFlag { + fn drop(&mut self) { + self.released.store(true, Ordering::SeqCst); + } +} + +#[derive(Debug)] +struct ObservedPrefix { + bytes: usize, + context_released: Arc, + dropped: Arc, +} + +impl Drop for ObservedPrefix { + fn drop(&mut self) { + assert!( + !self.context_released.load(Ordering::SeqCst), + "shared-prefix teardown ran after its provider context pin was released" + ); + self.dropped.store(true, Ordering::SeqCst); + } +} + +impl SharedDevicePrefix for ObservedPrefix { + fn device_ptr(&self) -> u64 { + 0x2000 + } + + fn committed_physical_bytes(&self) -> u64 { + self.bytes as u64 + } + + fn mapped_bytes(&self) -> usize { + self.bytes + } + + fn requested_bytes(&self) -> usize { + self.bytes + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +#[derive(Debug)] +struct DropOrderMechanism { + context_released: Arc, + prefix_dropped: Arc, +} + +impl DeviceAllocator for DropOrderMechanism { + fn allocate(&self, bytes: usize, align: usize) -> Result, MemoryError> { + HostAllocator.allocate(bytes, align) + } + + unsafe fn deallocate(&self, ptr: NonNull, bytes: usize, align: usize) { + // SAFETY: forwarded unchanged from this method's contract. + unsafe { HostAllocator.deallocate(ptr, bytes, align) }; + } + + fn device(&self) -> DeviceKey { + DeviceKey::HOST + } + + fn as_shared_mapping(&self) -> Option<&dyn SharedMapping> { + Some(self) + } +} + +impl SharedMapping for DropOrderMechanism { + fn create_shared_prefix( + &self, + bytes: usize, + ) -> Result, MemoryError> { + Ok(Box::new(ObservedPrefix { + bytes, + context_released: Arc::clone(&self.context_released), + dropped: Arc::clone(&self.prefix_dropped), + })) + } + + fn incremental_owned_bytes_for_shared_prefix( + &self, + prefix: &dyn SharedDevicePrefix, + ) -> Result { + Ok(prefix.committed_physical_bytes()) + } + + fn commit_shared_prefix( + &self, + _prefix: &dyn SharedDevicePrefix, + _ptr: NonNull, + _allocation_bytes: usize, + _byte_offset: usize, + ) -> Result { + Ok(SharedPrefixCommitInfo::default()) + } +} + #[derive(Debug)] struct TestMechanism { device: DeviceKey, @@ -578,6 +683,54 @@ fn shared_prefix_and_allocation_must_have_the_same_binding() { first.release(allocation).unwrap(); } +#[test] +fn shared_prefix_drops_before_context_pin_after_terminal_registry_teardown() { + let context_released = Arc::new(AtomicBool::new(false)); + let prefix_dropped = Arc::new(AtomicBool::new(false)); + let registry = BindingRegistry::new().unwrap(); + let context = registry + .register_provider_context( + DeviceKey::HOST, + Arc::new(ContextDropFlag { + released: Arc::clone(&context_released), + }) as Arc, + ) + .unwrap(); + let authority = registry + .register_authority(DeviceKey::HOST, Arc::new(()) as Arc) + .unwrap(); + let mechanism = registry + .register_allocator( + context, + authority, + Arc::new(DropOrderMechanism { + context_released: Arc::clone(&context_released), + prefix_dropped: Arc::clone(&prefix_dropped), + }) as Arc, + ) + .unwrap(); + let binding = registry.bind(DeviceKey::HOST).unwrap(); + let mapping = binding.shared_mapping().unwrap().unwrap(); + let prefix = mapping.create_shared_prefix(64).unwrap(); + drop(mapping); + drop(binding); + + registry + .invalidate_device(DeviceKey::HOST, "drop-order regression") + .unwrap(); + registry.confirm_context_terminated(context).unwrap(); + registry.remove(mechanism).unwrap(); + registry.remove_provider_context(context).unwrap(); + registry.remove_authority(authority).unwrap(); + drop(registry); + assert!(!prefix_dropped.load(Ordering::SeqCst)); + assert!(!context_released.load(Ordering::SeqCst)); + + drop(prefix); + assert!(prefix_dropped.load(Ordering::SeqCst)); + assert!(context_released.load(Ordering::SeqCst)); +} + #[test] fn device_loss_invalidates_without_free_or_resource_refund() { let context_drops = Arc::new(AtomicUsize::new(0)); diff --git a/docs/memory/MEMORY_ARCHITECTURE.md b/docs/memory/MEMORY_ARCHITECTURE.md index 56eadac22a..98e4633727 100644 --- a/docs/memory/MEMORY_ARCHITECTURE.md +++ b/docs/memory/MEMORY_ARCHITECTURE.md @@ -1697,7 +1697,9 @@ capability handles clone that same binding. Dropping a session/EP front-end or the registry therefore cannot destroy the context while any bound metadata or active operation still references it. `BoundAllocation` has no `Drop` release; whole-allocation release remains explicit through `MemoryBinding::release` and -the original `DeviceAllocator`. +the original `DeviceAllocator`. Phase 3 calls canonical `deallocate` and does +not expose the EP's mapped-attribution refund from `deallocate_with_unmapped`; +EP adoption must wait for Phase-4 accounting reconciliation. Changing the selected mechanism affects only later `bind(device)` calls. An old binding remains pinned to its original mechanism and explicitly releases there. @@ -1734,7 +1736,10 @@ step after mechanism registrations are quiescent and removed. Phase 3 adds no deferred-free queue, fence/event scheduling, owning allocation RAII, physical-release completeness state, partial-unmap recovery, quarantine, -or pointer-only retry API. Those remain Phase 4 work. +or pointer-only retry API. `BoundSharedPrefix` teardown is still drop-driven, +not lifecycle-gated by the registry; field order only guarantees that its +provider-context/resource pin outlives physical-prefix destruction. True +stream-ordered, partial-failure-safe teardown remains Phase 4 work. ### 1.2 Two directions, both backends