fix(memory): preserve Phase 3 binding teardown and selection

crates/onnx-runtime-memory-api/src/binding.rs

@@ -488,6 +488,16 @@ struct RegistryState {
 struct RegistryInner {
     identities: Arc<IdentitySource>,
     state: Mutex<RegistryState>,
+    #[cfg(test)]
+    select_after_validation_hook: Mutex<Option<SelectAfterValidationHook>>,
+}
+
+#[cfg(test)]
+#[derive(Clone, Debug)]
+struct SelectAfterValidationHook {
+    mechanism: MechanismIdentity,
+    entered: Arc<std::sync::Barrier>,
+    resume: Arc<std::sync::Barrier>,
 }
 
 /// 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<MemoryBinding, BindingError> {
         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<dyn SharedDevicePrefix>,
+    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<dyn BindingResource>)
+            .unwrap();
+        let authority = registry
+            .register_authority(DeviceKey::HOST, Arc::new(()) as Arc<dyn BindingResource>)
+            .unwrap();
+        let prior = registry
+            .register_allocator(
+                context,
+                authority,
+                Arc::new(HostAllocator) as Arc<dyn DeviceAllocator>,
+            )
+            .unwrap();
+        let candidate = registry
+            .register_allocator(
+                context,
+                authority,
+                Arc::new(HostAllocator) as Arc<dyn DeviceAllocator>,
+            )
+            .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);
+    }
+}

crates/onnx-runtime-memory-api/tests/binding_registry.rs

@@ -50,6 +50,111 @@ impl SharedDevicePrefix for TestPrefix {
     }
 }
 
+#[derive(Debug)]
+struct ContextDropFlag {
+    released: Arc<AtomicBool>,
+}
+
+impl Drop for ContextDropFlag {
+    fn drop(&mut self) {
+        self.released.store(true, Ordering::SeqCst);
+    }
+}
+
+#[derive(Debug)]
+struct ObservedPrefix {
+    bytes: usize,
+    context_released: Arc<AtomicBool>,
+    dropped: Arc<AtomicBool>,
+}
+
+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<AtomicBool>,
+    prefix_dropped: Arc<AtomicBool>,
+}
+
+impl DeviceAllocator for DropOrderMechanism {
+    fn allocate(&self, bytes: usize, align: usize) -> Result<NonNull<u8>, MemoryError> {
+        HostAllocator.allocate(bytes, align)
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, 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<Box<dyn SharedDevicePrefix>, 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<u64, MemoryError> {
+        Ok(prefix.committed_physical_bytes())
+    }
+
+    fn commit_shared_prefix(
+        &self,
+        _prefix: &dyn SharedDevicePrefix,
+        _ptr: NonNull<u8>,
+        _allocation_bytes: usize,
+        _byte_offset: usize,
+    ) -> Result<SharedPrefixCommitInfo, MemoryError> {
+        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<dyn BindingResource>,
+        )
+        .unwrap();
+    let authority = registry
+        .register_authority(DeviceKey::HOST, Arc::new(()) as Arc<dyn BindingResource>)
+        .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<dyn DeviceAllocator>,
+        )
+        .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));

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