onnx_runtime_memory_api/
context_pin.rs

1//! Keeping a provider/context alive across an outstanding deferred release.
2//!
3//! A deferred release is a promise that someone will free device memory later.
4//! Whoever performs that free needs the provider/context it was allocated
5//! against to still exist: a stream-ordered free retiring after its CUDA
6//! context has been torn down unmaps handles the teardown already released,
7//! which is undefined behaviour at the driver level rather than a Rust error.
8//!
9//! The host side of a plugin boundary cannot depend on the governor, so the
10//! pin is expressed here as two object-safe traits. The governor implements
11//! them over its provider-context records; a mechanism holds
12//! [`ProviderContextPin`] values and knows nothing about how the pin is
13//! counted.
14//!
15//! The direction of the guarantee matters. A pin does not ask a context to
16//! stay alive; it makes teardown observe the outstanding work. Acquiring a pin
17//! against a context that is already retiring must fail rather than succeed,
18//! because a release that cannot be pinned is a release that must not be
19//! queued.
20
21use core::fmt::Debug;
22
23use crate::binding::ProviderContextIdentity;
24
25/// Why a provider/context pin could not be acquired.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum ProviderContextPinError {
28    /// The context is retiring, lost, or terminated, so no new work may be
29    /// attached to it. The caller must not queue the release.
30    ContextUnavailable(ProviderContextIdentity),
31    /// The implementation's outstanding-work counter cannot represent another
32    /// pin. Treated as a refusal rather than a wrap.
33    PinCountOverflow(ProviderContextIdentity),
34}
35
36impl core::fmt::Display for ProviderContextPinError {
37    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
38        match self {
39            Self::ContextUnavailable(identity) => write!(
40                formatter,
41                "provider context {identity:?} is no longer accepting work"
42            ),
43            Self::PinCountOverflow(identity) => write!(
44                formatter,
45                "provider context {identity:?} has too many outstanding pins"
46            ),
47        }
48    }
49}
50
51impl std::error::Error for ProviderContextPinError {}
52
53/// A live claim on a provider/context.
54///
55/// Dropping the value releases the claim, so implementations carry the
56/// bookkeeping in `Drop` rather than in an explicit method. Nothing about the
57/// pin is inspectable beyond the identity it belongs to: a holder must not be
58/// able to extend, transfer, or interrogate the context through it.
59pub trait ProviderContextPin: Send + Sync + Debug {
60    /// Which context this pin holds.
61    fn context(&self) -> ProviderContextIdentity;
62}
63
64/// Something that can hand out [`ProviderContextPin`]s for one context.
65pub trait ProviderContextPinSource: Send + Sync + Debug {
66    /// Which context this source pins.
67    fn context(&self) -> ProviderContextIdentity;
68
69    /// Claim the context until the returned pin is dropped.
70    ///
71    /// Fails when the context is no longer accepting work. Callers must treat
72    /// a failure as a refusal to queue the deferred release, not as a reason
73    /// to proceed unpinned.
74    fn pin(&self) -> Result<Box<dyn ProviderContextPin>, ProviderContextPinError>;
75}
onnx_runtime_memory_api/context_pin.rs

onnx_runtime_memory_api/
context_pin.rs
