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

1use super::*;
2
3const CPU_LIFECYCLE_OFFLINE: usize = 1 << (usize::BITS - 1);
4const CPU_LIFECYCLE_INACTIVE: usize = 1 << (usize::BITS - 2);
5const CPU_LIFECYCLE_DRAINING: usize = CPU_LIFECYCLE_OFFLINE | CPU_LIFECYCLE_INACTIVE;
6const CPU_LIFECYCLE_MASK: usize = CPU_LIFECYCLE_DRAINING;
7const CPU_PUBLICATION_COUNT_MASK: usize = !CPU_LIFECYCLE_MASK;
8const CPU_PUBLICATION_OVERFLOW_INVARIANT: u32 = 0x4350_5542;
9const CPU_PUBLICATION_RELEASE_INVARIANT: u32 = 0x4350_5544;
10
11pub(super) const INITIAL_CPU_LIFECYCLE_STATE: usize = CPU_LIFECYCLE_OFFLINE;
12
13#[derive(Clone, Copy)]
14enum CpuPublicationClass {
15    Placement,
16    OwnerControl,
17}
18
19impl CpuPublicationClass {
20    const fn accepts(self, state: usize) -> bool {
21        match self {
22            Self::Placement => state & CPU_LIFECYCLE_MASK == 0,
23            Self::OwnerControl => state & CPU_LIFECYCLE_OFFLINE == 0,
24        }
25    }
26}
27
28#[derive(Debug)]
29pub(super) struct CpuPublicationState {
30    state: AtomicUsize,
31}
32
33impl CpuPublicationState {
34    pub(super) const fn new() -> Self {
35        Self {
36            state: AtomicUsize::new(INITIAL_CPU_LIFECYCLE_STATE),
37        }
38    }
39}
40
41/// Placement and remote-publication state of one logical CPU.
42#[derive(Clone, Copy, Debug, Eq, PartialEq)]
43pub enum CpuLifecycleState {
44    /// The CPU accepts placement and remote scheduler publications.
45    Online,
46    /// New placement is closed while owner-directed control delivery may finish.
47    Inactive,
48    /// Every remote publication is closed while the owner proves work is gone.
49    Draining,
50    /// The CPU owns no schedulable work and is absent from the root domain.
51    Offline,
52}
53
54impl CpuRemote {
55    /// Returns the CPU's placement and publication lifecycle.
56    pub fn lifecycle_state(&self) -> CpuLifecycleState {
57        match self.publication.state.load(Ordering::Acquire) & CPU_LIFECYCLE_MASK {
58            0 => CpuLifecycleState::Online,
59            CPU_LIFECYCLE_INACTIVE => CpuLifecycleState::Inactive,
60            CPU_LIFECYCLE_DRAINING => CpuLifecycleState::Draining,
61            CPU_LIFECYCLE_OFFLINE => CpuLifecycleState::Offline,
62            _ => unreachable!("CPU lifecycle mask has four encoded states"),
63        }
64    }
65
66    /// Returns whether owner initialization and online publication completed.
67    pub fn is_online(&self) -> bool {
68        matches!(
69            self.lifecycle_state(),
70            CpuLifecycleState::Online | CpuLifecycleState::Inactive
71        )
72    }
73
74    /// Returns whether new runnable placement may target this CPU.
75    pub(crate) fn accepts_placement(&self) -> bool {
76        self.lifecycle_state() == CpuLifecycleState::Online
77    }
78
79    pub(crate) fn mark_online(&self) -> bool {
80        self.publication
81            .state
82            .compare_exchange(
83                CPU_LIFECYCLE_OFFLINE,
84                0,
85                Ordering::Release,
86                Ordering::Acquire,
87            )
88            .is_ok()
89    }
90
91    pub(crate) fn try_deactivate(&self) -> bool {
92        // Linux clears cpu_active before synchronizing prior placement readers.
93        // Preserve their leases while closing admission to new placement.
94        let inactive = self
95            .publication
96            .state
97            .try_update(Ordering::AcqRel, Ordering::Acquire, |state| {
98                (state & CPU_LIFECYCLE_MASK == 0).then_some(state | CPU_LIFECYCLE_INACTIVE)
99            })
100            .is_ok();
101        if inactive {
102            self.cancel_idle_pull_if_uncommitted();
103        }
104        inactive
105    }
106
107    pub(crate) fn resume_owner_drain(&self) {
108        self.publication
109            .state
110            .compare_exchange(
111                CPU_LIFECYCLE_DRAINING,
112                CPU_LIFECYCLE_INACTIVE,
113                Ordering::Release,
114                Ordering::Acquire,
115            )
116            .expect("owner drain resumes only from a closed publication gate");
117    }
118
119    pub(crate) fn cancel_deactivation(&self) {
120        let mut current = self.publication.state.load(Ordering::Acquire);
121        loop {
122            if current & CPU_LIFECYCLE_MASK != CPU_LIFECYCLE_INACTIVE {
123                task_runtime::fatal_invariant(
124                    CPU_PUBLICATION_RELEASE_INVARIANT,
125                    self.owner.as_u32() as usize,
126                );
127            }
128            let online = current & !CPU_LIFECYCLE_INACTIVE;
129            match self.publication.state.compare_exchange_weak(
130                current,
131                online,
132                Ordering::Release,
133                Ordering::Acquire,
134            ) {
135                Ok(_) => return,
136                Err(actual) => current = actual,
137            }
138        }
139    }
140
141    pub(crate) fn try_begin_draining(&self) -> bool {
142        // An exact inactive state also proves that every owner-directed control
143        // delivery has completed its publication and doorbell transaction.
144        self.publication
145            .state
146            .compare_exchange(
147                CPU_LIFECYCLE_INACTIVE,
148                CPU_LIFECYCLE_DRAINING,
149                Ordering::AcqRel,
150                Ordering::Acquire,
151            )
152            .is_ok()
153    }
154
155    pub(crate) fn cancel_draining(&self) {
156        if self
157            .publication
158            .state
159            .compare_exchange(
160                CPU_LIFECYCLE_DRAINING,
161                0,
162                Ordering::Release,
163                Ordering::Acquire,
164            )
165            .is_err()
166        {
167            task_runtime::fatal_invariant(
168                CPU_PUBLICATION_RELEASE_INVARIANT,
169                self.owner.as_u32() as usize,
170            );
171        }
172    }
173
174    pub(crate) fn finish_offline(&self) {
175        self.reset_scheduler_for_offline();
176        if self
177            .publication
178            .state
179            .compare_exchange(
180                CPU_LIFECYCLE_DRAINING,
181                CPU_LIFECYCLE_OFFLINE,
182                Ordering::Release,
183                Ordering::Acquire,
184            )
185            .is_err()
186        {
187            task_runtime::fatal_invariant(
188                CPU_PUBLICATION_RELEASE_INVARIANT,
189                self.owner.as_u32() as usize,
190            );
191        }
192    }
193
194    pub(crate) fn begin_publication(&self) -> Option<CpuRemotePublication<'_>> {
195        self.try_acquire_publication(CpuPublicationClass::Placement)
196            .then(|| CpuRemotePublication { remote: self })
197    }
198
199    fn try_acquire_publication(&self, class: CpuPublicationClass) -> bool {
200        let mut current = self.publication.state.load(Ordering::Acquire);
201        loop {
202            if !class.accepts(current) {
203                return false;
204            }
205            let count = current & CPU_PUBLICATION_COUNT_MASK;
206            if count == CPU_PUBLICATION_COUNT_MASK {
207                task_runtime::fatal_invariant(
208                    CPU_PUBLICATION_OVERFLOW_INVARIANT,
209                    self.owner.as_u32() as usize,
210                );
211            }
212            match self.publication.state.compare_exchange_weak(
213                current,
214                current + 1,
215                Ordering::AcqRel,
216                Ordering::Acquire,
217            ) {
218                Ok(_) => {
219                    #[cfg(feature = "qperf-metrics")]
220                    match class {
221                        CpuPublicationClass::Placement => {
222                            crate::diagnostics::counters::record_cpu_placement_publication_acquire(
223                            );
224                        }
225                        CpuPublicationClass::OwnerControl => {
226                            crate::diagnostics::counters::record_cpu_owner_control_publication_acquire();
227                        }
228                    }
229                    return true;
230                }
231                Err(actual) => current = actual,
232            }
233        }
234    }
235
236    /// Pins an online placement target across an owner context switch.
237    ///
238    /// Linux holds the CPU-hotplug/read-side ownership that makes the selected
239    /// destination runqueue stable before it commits `TASK_ON_RQ_MIGRATING`.
240    /// The owned form carries the same lifetime proof through switch tail,
241    /// where a borrowed runqueue guard cannot survive the architecture switch.
242    pub(crate) fn begin_owned_publication(self: &Arc<Self>) -> Option<OwnedCpuRemotePublication> {
243        self.try_acquire_publication(CpuPublicationClass::Placement)
244            .then(|| OwnedCpuRemotePublication {
245                remote: Arc::clone(self),
246            })
247    }
248
249    pub(crate) fn begin_owner_delivery(&self) -> Option<CpuRemotePublication<'_>> {
250        self.try_acquire_publication(CpuPublicationClass::OwnerControl)
251            .then(|| CpuRemotePublication { remote: self })
252    }
253
254    pub(crate) fn is_quiescent_for_offline(&self) -> bool {
255        self.publication.state.load(Ordering::Acquire) == CPU_LIFECYCLE_DRAINING
256            && self.ktimer_is_quiescent_for_offline()
257            && !self.needs_reschedule()
258            && !self.has_remote_work()
259            && !self.is_idle_polling()
260            && self.idle_pull_is_quiescent()
261    }
262}
263
264pub(crate) struct CpuRemotePublication<'remote> {
265    remote: &'remote CpuRemote,
266}
267
268#[derive(Debug)]
269pub(crate) struct OwnedCpuRemotePublication {
270    remote: Arc<CpuRemote>,
271}
272
273impl OwnedCpuRemotePublication {
274    pub(crate) fn publish_owner_control(
275        self,
276        node: Pin<&'static InboxNode>,
277        message: InboxMessage,
278    ) -> PublishResult {
279        self.remote.publish_owner_control_owned(node, message)
280    }
281}
282
283impl CpuRemotePublication<'_> {
284    pub(crate) fn publish_owner_control(
285        self,
286        node: Pin<&'static InboxNode>,
287        message: InboxMessage,
288    ) -> PublishResult {
289        self.remote.publish_owner_control_owned(node, message)
290    }
291}
292
293impl Drop for CpuRemotePublication<'_> {
294    fn drop(&mut self) {
295        release_publication(self.remote);
296    }
297}
298
299impl Drop for OwnedCpuRemotePublication {
300    fn drop(&mut self) {
301        release_publication(&self.remote);
302    }
303}
304
305fn release_publication(remote: &CpuRemote) {
306    let mut current = remote.publication.state.load(Ordering::Acquire);
307    loop {
308        if current & CPU_LIFECYCLE_OFFLINE != 0 || current & CPU_PUBLICATION_COUNT_MASK == 0 {
309            task_runtime::fatal_invariant(
310                CPU_PUBLICATION_RELEASE_INVARIANT,
311                remote.owner.as_u32() as usize,
312            );
313        }
314        match remote.publication.state.compare_exchange_weak(
315            current,
316            current - 1,
317            Ordering::Release,
318            Ordering::Acquire,
319        ) {
320            Ok(_) => return,
321            Err(actual) => current = actual,
322        }
323    }
324}
ax_task/sched/system/cpu/remote/lifecycle.rs

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