ax_task/sched/system/cpu/remote/
owner.rs

1use super::*;
2
3#[derive(Debug)]
4pub(super) struct OwnerState {
5    claimed: AtomicBool,
6    idle_thread: AtomicU64,
7    busy_runtime_ns: AtomicU64,
8}
9
10impl OwnerState {
11    pub(super) const fn new() -> Self {
12        Self {
13            claimed: AtomicBool::new(false),
14            idle_thread: AtomicU64::new(0),
15            busy_runtime_ns: AtomicU64::new(0),
16        }
17    }
18}
19
20impl CpuRemote {
21    /// Returns the CPU that owns the corresponding runqueue.
22    pub const fn owner(&self) -> CpuId {
23        self.owner
24    }
25
26    /// Claims exclusive access to the corresponding owner-only scheduler object.
27    ///
28    /// # Safety
29    ///
30    /// `cpu` must identify the pinned, live [`CpuLocal`] associated with this
31    /// endpoint. After runtime publication, every access that can overlap this
32    /// claim must use the same endpoint rather than retaining an ungated borrow.
33    pub unsafe fn claim_local(
34        &self,
35        cpu: *mut CpuLocal,
36    ) -> Result<CpuLocalOwnerBorrow<'_>, TaskError> {
37        let cpu = NonNull::new(cpu).ok_or(TaskError::InvalidRuntimeHandle)?;
38        self.owner_state
39            .claimed
40            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
41            .map_err(|_| TaskError::CpuOwnerBorrowed)?;
42
43        // SAFETY: the caller guarantees that this is the live pinned CpuLocal
44        // paired with this endpoint. The successful gate claim excludes every
45        // other runtime-derived reference while the identity is checked.
46        let actual = unsafe { cpu.as_ref() }.owner();
47        if actual != self.owner {
48            self.owner_state.claimed.store(false, Ordering::Release);
49            return Err(TaskError::CpuOwnerMismatch {
50                expected: self.owner.as_u32(),
51                actual: actual.as_u32(),
52            });
53        }
54        #[cfg(feature = "qperf-metrics")]
55        crate::diagnostics::counters::record_runtime_cpu_owner_claim();
56        Ok(CpuLocalOwnerBorrow {
57            remote: self,
58            cpu,
59            release_claim: true,
60            _not_send_or_sync: PhantomData,
61        })
62    }
63
64    /// Borrows the owner-only scheduler state under an existing scheduler baton.
65    ///
66    /// # Safety
67    ///
68    /// `cpu` must identify this endpoint's pinned [`CpuLocal`]. The caller must
69    /// own the CPU's IRQ-off scheduler frame for the complete returned borrow,
70    /// and no dynamically claimed owner borrow may overlap it.
71    pub unsafe fn borrow_local_in_scheduler_frame(
72        &self,
73        cpu: NonNull<CpuLocal>,
74    ) -> CpuLocalOwnerBorrow<'_> {
75        CpuLocalOwnerBorrow {
76            remote: self,
77            cpu,
78            release_claim: false,
79            _not_send_or_sync: PhantomData,
80        }
81    }
82
83    /// Returns `rq->curr` under the authoritative runqueue lock.
84    pub fn current_thread(&self) -> Option<ThreadId> {
85        self.lock_run_queue(RunQueueGuardSource::OwnerCurrentThreadObservation)
86            .current_thread()
87    }
88
89    /// Returns the configured idle-thread snapshot.
90    pub fn idle_thread(&self) -> Option<ThreadId> {
91        decode_thread_id(self.owner_state.idle_thread.load(Ordering::Acquire))
92    }
93
94    pub(in crate::sched::system::cpu) fn publish_idle_thread(&self, idle: ThreadId) {
95        self.owner_state
96            .idle_thread
97            .store(idle.as_u64(), Ordering::Release);
98    }
99
100    /// Returns cumulative time this CPU has executed non-idle scheduler threads.
101    pub fn busy_runtime_ns(&self) -> u64 {
102        self.owner_state.busy_runtime_ns.load(Ordering::Relaxed)
103    }
104
105    pub(in crate::sched::system::cpu) fn charge_busy_runtime(&self, runtime_ns: u64) {
106        // Runtime charging is serialized by this CPU's owner rq lock. Other
107        // CPUs only sample the counter, so an atomic read/write publication is
108        // sufficient and avoids pretending there are concurrent writers.
109        let committed = self.owner_state.busy_runtime_ns.load(Ordering::Relaxed);
110        self.owner_state
111            .busy_runtime_ns
112            .store(committed.saturating_add(runtime_ns), Ordering::Relaxed);
113    }
114}
115
116/// Exclusive owner borrow of one pinned [`CpuLocal`].
117///
118/// Ordinary callers acquire the dynamic gate in the separately allocated
119/// [`CpuRemote`] endpoint. A live IRQ-off scheduler frame may instead lend its
120/// stronger CPU-owner baton to the same borrow type without another atomic
121/// ownership transaction.
122pub struct CpuLocalOwnerBorrow<'remote> {
123    remote: &'remote CpuRemote,
124    cpu: NonNull<CpuLocal>,
125    release_claim: bool,
126    _not_send_or_sync: PhantomData<*mut ()>,
127}
128
129impl CpuLocalOwnerBorrow<'_> {
130    /// Borrows the pinned owner state mutably for one audited call scope.
131    pub fn as_pin_mut(&mut self) -> Pin<&mut CpuLocal> {
132        // SAFETY: construction claimed the unique runtime owner gate, the
133        // pointer remains pinned, and the returned lifetime is bounded by the
134        // mutable borrow of this gate-owning wrapper.
135        unsafe { Pin::new_unchecked(self.cpu.as_mut()) }
136    }
137}
138
139impl Deref for CpuLocalOwnerBorrow<'_> {
140    type Target = CpuLocal;
141
142    fn deref(&self) -> &Self::Target {
143        // SAFETY: the wrapper owns the endpoint's exclusive claim and its
144        // lifetime is bounded by that claim.
145        unsafe { self.cpu.as_ref() }
146    }
147}
148
149impl Drop for CpuLocalOwnerBorrow<'_> {
150    fn drop(&mut self) {
151        if self.release_claim {
152            self.remote
153                .owner_state
154                .claimed
155                .store(false, Ordering::Release);
156        }
157    }
158}
159
160fn decode_thread_id(raw: u64) -> Option<ThreadId> {
161    (raw != 0).then(|| ThreadId::from_parts(raw as u32, (raw >> 32) as u32))
162}
ax_task/sched/system/cpu/remote/owner.rs

ax_task/sched/system/cpu/remote/
owner.rs
