Skip to main content

ax_runtime/thread/
mm_activation.rs

1//! Allocation-free ownership transferred at the runtime's root-switch boundary.
2
3use alloc::sync::Arc;
4
5use ax_hal::context::InstalledAddressSpace;
6
7use super::TaskError;
8
9/// Evidence that this CPU has replaced its previous hardware root.
10#[derive(Debug)]
11pub struct AddressSpaceSwitchProof {
12    cpu: usize,
13}
14
15impl AddressSpaceSwitchProof {
16    #[cfg(feature = "uspace")]
17    pub(super) const fn new(cpu: usize) -> Self {
18        Self { cpu }
19    }
20
21    /// Returns the CPU whose old translations are no longer installed.
22    pub const fn cpu(&self) -> usize {
23        self.cpu
24    }
25}
26
27/// Preallocated MM accounting retained by one CPU activation.
28///
29/// These callbacks run with IRQs disabled. They must not allocate, sleep, or
30/// destroy the last lifetime anchor for the page tables or their backing store.
31pub trait SchedulerAddressSpaceOwner: Send + Sync {
32    /// Releases accounting after the runtime completed a hardware root switch.
33    fn release_after_root_switch(self: Arc<Self>, proof: AddressSpaceSwitchProof);
34    /// Cancels a reservation which the runtime never installed in hardware.
35    fn cancel_before_install(self: Arc<Self>, cpu: usize);
36    /// Retains an installed root for which no retirement proof was obtained.
37    fn abandon(self: Arc<Self>, cpu: usize);
38}
39
40/// OS ownership held by a runtime task token until its last CPU lease drains.
41///
42/// # Safety
43/// Every prepared activation must own the declared root and publish its CPU in
44/// the MM's TLB target set before returning. The root must remain valid until
45/// the activation is released. Detaching a task must retain storage needed by
46/// lazy CPUs; activation callbacks must not perform a final storage release.
47pub unsafe trait UserAddressSpaceOwner: Send + Sync {
48    /// Acquires a CPU activation without allocation, blocking, or hardware I/O.
49    fn prepare_activation(&self, cpu: usize) -> Result<SchedulerAddressSpaceActivation, TaskError>;
50    /// Drops task-scoped ownership in ordinary task context, exactly once.
51    fn detach_from_task(&self);
52}
53
54/// An inline reservation, subsequently owned by the CPU that installs it.
55///
56/// The owner Arc already exists; unsizing or cloning it does not allocate.
57pub struct SchedulerAddressSpaceActivation {
58    installed: InstalledAddressSpace,
59    cpu: usize,
60    owner: Option<Arc<dyn SchedulerAddressSpaceOwner>>,
61    committed: bool,
62}
63
64impl SchedulerAddressSpaceActivation {
65    /// Transfers an acquired, not yet installed MM activation to the runtime.
66    pub fn new(
67        installed: InstalledAddressSpace,
68        cpu: usize,
69        owner: Arc<dyn SchedulerAddressSpaceOwner>,
70    ) -> Self {
71        Self {
72            installed,
73            cpu,
74            owner: Some(owner),
75            committed: false,
76        }
77    }
78
79    /// Returns the complete root, hardware tag, generation and epoch identity.
80    pub const fn installed(&self) -> InstalledAddressSpace {
81        self.installed
82    }
83
84    #[cfg(feature = "uspace")]
85    pub(super) fn commit(&mut self, cpu: usize) {
86        assert_eq!(self.cpu, cpu);
87        assert!(!self.committed);
88        self.committed = true;
89    }
90
91    #[cfg(feature = "uspace")]
92    pub(super) fn release(mut self, proof: AddressSpaceSwitchProof) {
93        assert_eq!(self.cpu, proof.cpu());
94        assert!(self.committed);
95        self.owner
96            .take()
97            .expect("activation is consumed once")
98            .release_after_root_switch(proof);
99    }
100}
101
102impl Drop for SchedulerAddressSpaceActivation {
103    fn drop(&mut self) {
104        if let Some(owner) = self.owner.take() {
105            if self.committed {
106                owner.abandon(self.cpu);
107            } else {
108                owner.cancel_before_install(self.cpu);
109            }
110        }
111    }
112}