ax_task/thread/
spec.rs

1//! Thread construction data kept independent from an operating system.
2
3use alloc::{sync::Arc, vec, vec::Vec};
4
5use crate::{
6    runtime::{
7        resource::{
8            AddressSpaceHandle, AddressSpaceToken, ExecutionContextHandle, StackHandle, TlsHandle,
9        },
10        service::{SchedulerTickCpuTime, SchedulerTickGate, SchedulerTickTaskWork},
11        task_runtime,
12    },
13    sched::{CpuId, SchedulePolicy},
14    thread::{SchedulerTickWork, TaskError, ThreadHandle, ThreadId},
15};
16
17/// Runtime-owned resources whose lifetime follows one thread.
18#[repr(C)]
19#[derive(Debug, Eq, PartialEq)]
20pub struct ThreadResources {
21    context: ExecutionContextHandle,
22    stack: StackHandle,
23    tls: TlsHandle,
24    address_space: AddressSpaceToken,
25}
26
27impl ThreadResources {
28    /// Empty resources for pure scheduler models.
29    pub const NONE: Self = Self {
30        context: ExecutionContextHandle::NONE,
31        stack: StackHandle::NONE,
32        tls: TlsHandle::NONE,
33        address_space: AddressSpaceToken::NONE,
34    };
35
36    /// Creates a complete runtime resource bundle from uniquely owned handles.
37    ///
38    /// # Safety
39    ///
40    /// Every non-empty handle must be live, belong to the currently installed
41    /// [`crate::runtime::TaskRuntime`], and have its unique destruction right
42    /// transferred into this bundle. The caller must not construct another
43    /// owning bundle from the same scalar handles.
44    pub const unsafe fn new(
45        context: ExecutionContextHandle,
46        stack: StackHandle,
47        tls: TlsHandle,
48        address_space: AddressSpaceToken,
49    ) -> Self {
50        Self {
51            context,
52            stack,
53            tls,
54            address_space,
55        }
56    }
57
58    /// Returns the execution context.
59    pub const fn context(&self) -> ExecutionContextHandle {
60        self.context
61    }
62    /// Returns the guarded stack allocation.
63    pub const fn stack(&self) -> StackHandle {
64        self.stack
65    }
66    /// Returns the TLS allocation.
67    pub const fn tls(&self) -> TlsHandle {
68        self.tls
69    }
70    /// Returns the address-space handle.
71    pub const fn address_space(&self) -> AddressSpaceHandle {
72        self.address_space.handle()
73    }
74
75    pub(crate) fn replace_address_space(
76        &mut self,
77        address_space: AddressSpaceToken,
78    ) -> AddressSpaceToken {
79        core::mem::replace(&mut self.address_space, address_space)
80    }
81
82    pub(crate) fn take_address_space(&mut self) -> AddressSpaceToken {
83        core::mem::replace(&mut self.address_space, AddressSpaceToken::NONE)
84    }
85
86    /// Releases thread-private resources and returns the independent active-mm
87    /// ownership token.
88    ///
89    /// The registry calls this only after switch tail has cleared physical CPU
90    /// ownership. Context, TLS, and stack destruction are consequently
91    /// one-way operations with no retry state. The address-space token has a
92    /// separate active-CPU lifetime and is handed to that reclaim protocol
93    /// instead of retaining already-dead thread resources.
94    pub(crate) fn release(mut self) -> AddressSpaceToken {
95        if !self.context.is_none() {
96            task_runtime::destroy_context(self.context);
97            self.context = ExecutionContextHandle::NONE;
98        }
99
100        if !self.tls.is_none() {
101            task_runtime::deallocate_tls(self.tls);
102            self.tls = TlsHandle::NONE;
103        }
104
105        if !self.stack.is_none() {
106            task_runtime::deallocate_stack(self.stack);
107            self.stack = StackHandle::NONE;
108        }
109
110        core::mem::replace(&mut self.address_space, AddressSpaceToken::NONE)
111    }
112}
113
114/// Why a running thread relinquished its execution context.
115///
116/// The value crosses the OS extension callback boundary, so its numeric layout
117/// is stable and may also be written directly to allocation-free trace records.
118#[repr(u32)]
119#[derive(Clone, Copy, Debug, Eq, PartialEq)]
120pub enum SwitchReason {
121    /// A scheduler request selected a more urgent or otherwise eligible thread.
122    Preempted = 1,
123    /// The thread voluntarily yielded its current service position.
124    Yield     = 2,
125    /// The thread committed a park or another blocking operation.
126    Blocked   = 3,
127    /// The thread terminated and will never become runnable again.
128    Exited    = 4,
129    /// CPU affinity or balancing moved the thread away from this CPU.
130    Migrated  = 5,
131}
132
133/// CPU affinity expressed against one [`crate::runtime::TaskSystem`] topology.
134#[derive(Clone, Debug, Eq, PartialEq)]
135pub struct CpuSet {
136    words: Vec<usize>,
137    topology_len: usize,
138    // Mirrors Linux task_struct::nr_cpus_allowed so scheduler class decisions
139    // do not repeatedly derive affinity cardinality from the mask.
140    allowed_count: usize,
141}
142
143impl CpuSet {
144    const BITS_PER_WORD: usize = usize::BITS as usize;
145
146    /// Creates a set that permits every CPU in a topology.
147    pub fn all(cpu_count: usize) -> Self {
148        let mut words = vec![usize::MAX; cpu_count.div_ceil(Self::BITS_PER_WORD)];
149        if let Some(last) = words.last_mut()
150            && !cpu_count.is_multiple_of(Self::BITS_PER_WORD)
151        {
152            *last = (1usize << (cpu_count % Self::BITS_PER_WORD)) - 1;
153        }
154        Self {
155            words,
156            topology_len: cpu_count,
157            allowed_count: cpu_count,
158        }
159    }
160
161    pub(crate) fn try_all(cpu_count: usize) -> Result<Self, super::TaskError> {
162        let mut words =
163            crate::thread::allocation::try_vec(cpu_count.div_ceil(Self::BITS_PER_WORD))?;
164        words.resize(cpu_count.div_ceil(Self::BITS_PER_WORD), usize::MAX);
165        if let Some(last) = words.last_mut()
166            && !cpu_count.is_multiple_of(Self::BITS_PER_WORD)
167        {
168            *last = (1usize << (cpu_count % Self::BITS_PER_WORD)) - 1;
169        }
170        Ok(Self {
171            words,
172            topology_len: cpu_count,
173            allowed_count: cpu_count,
174        })
175    }
176
177    /// Creates an empty CPU set for a topology.
178    pub fn empty(cpu_count: usize) -> Self {
179        Self {
180            words: vec![0; cpu_count.div_ceil(Self::BITS_PER_WORD)],
181            topology_len: cpu_count,
182            allowed_count: 0,
183        }
184    }
185
186    /// Enables one CPU if it is represented by this set.
187    pub fn insert(&mut self, cpu: CpuId) -> bool {
188        let index = cpu.as_usize();
189        if index >= self.topology_len {
190            return false;
191        }
192        let mask = 1usize << (index % Self::BITS_PER_WORD);
193        let word = &mut self.words[index / Self::BITS_PER_WORD];
194        let changed = *word & mask == 0;
195        *word |= mask;
196        if changed {
197            self.allowed_count += 1;
198        }
199        changed
200    }
201
202    /// Disables one CPU if it is represented by this set.
203    pub fn remove(&mut self, cpu: CpuId) -> bool {
204        let index = cpu.as_usize();
205        if index >= self.topology_len {
206            return false;
207        }
208        let mask = 1usize << (index % Self::BITS_PER_WORD);
209        let word = &mut self.words[index / Self::BITS_PER_WORD];
210        let changed = *word & mask != 0;
211        *word &= !mask;
212        if changed {
213            self.allowed_count -= 1;
214        }
215        changed
216    }
217
218    pub(crate) fn clear(&mut self) {
219        self.words.fill(0);
220        self.allowed_count = 0;
221    }
222
223    /// Tests whether a CPU is allowed.
224    pub fn contains(&self, cpu: CpuId) -> bool {
225        let index = cpu.as_usize();
226        index < self.topology_len
227            && self.words[index / Self::BITS_PER_WORD] & (1usize << (index % Self::BITS_PER_WORD))
228                != 0
229    }
230
231    /// Returns the number of CPUs represented by the set.
232    pub fn topology_len(&self) -> usize {
233        self.topology_len
234    }
235
236    /// Returns the number of CPUs selected by this set.
237    pub(crate) fn count(&self) -> usize {
238        self.allowed_count
239    }
240
241    /// Iterates selected CPUs in ascending logical-ID order.
242    pub fn iter(&self) -> impl Iterator<Item = CpuId> + '_ {
243        (0..self.topology_len)
244            .map(|index| CpuId::new(index as u32))
245            .filter(|cpu| self.contains(*cpu))
246    }
247
248    /// Returns the only allowed CPU when migration is impossible.
249    pub(crate) fn sole_cpu(&self) -> Option<CpuId> {
250        if self.allowed_count != 1 {
251            return None;
252        }
253        let (word_index, word) = self
254            .words
255            .iter()
256            .copied()
257            .enumerate()
258            .find(|(_, word)| *word != 0)?;
259        let index = word_index * Self::BITS_PER_WORD + word.trailing_zeros() as usize;
260        (index < self.topology_len).then_some(CpuId::new(index as u32))
261    }
262
263    /// Returns whether a runnable thread can leave its current allowed CPU.
264    pub(crate) fn is_migration_capable(&self) -> bool {
265        self.allowed_count > 1
266    }
267
268    /// Returns whether this set permits every CPU selected by `required`.
269    pub fn covers(&self, required: &Self) -> bool {
270        self.topology_len == required.topology_len
271            && self
272                .words
273                .iter()
274                .zip(&required.words)
275                .all(|(allowed, is_required)| allowed & is_required == *is_required)
276    }
277
278    pub(crate) fn copy_from_set(&mut self, source: &Self) -> Result<(), TaskError> {
279        if self.topology_len != source.topology_len {
280            return Err(TaskError::InvalidConfiguration);
281        }
282        self.words.copy_from_slice(&source.words);
283        self.allowed_count = source.allowed_count;
284        Ok(())
285    }
286
287    /// Returns the first CPU in the intersection that satisfies `accepts`.
288    ///
289    /// This is the `cpumask_any_and()` primitive used by cpupri/cpudl: the
290    /// intersection is formed a machine word at a time rather than scanning
291    /// every logical CPU.
292    pub(crate) fn first_intersection(
293        &self,
294        other: &Self,
295        mut accepts: impl FnMut(CpuId) -> bool,
296    ) -> Option<CpuId> {
297        if self.topology_len != other.topology_len {
298            return None;
299        }
300        for (word_index, (left, right)) in self.words.iter().zip(&other.words).enumerate() {
301            let mut candidates = left & right;
302            while candidates != 0 {
303                let bit = candidates.trailing_zeros() as usize;
304                candidates &= candidates - 1;
305                let index = word_index * Self::BITS_PER_WORD + bit;
306                if index >= self.topology_len {
307                    break;
308                }
309                let cpu = CpuId::new(index as u32);
310                if accepts(cpu) {
311                    return Some(cpu);
312                }
313            }
314        }
315        None
316    }
317
318    pub(crate) fn word(&self, word_index: usize) -> usize {
319        self.words.get(word_index).copied().unwrap_or(0)
320    }
321}
322
323/// OS-owned callbacks attached to a thread without exposing OS types.
324#[repr(C)]
325#[derive(Debug)]
326pub struct ThreadExtensionOps {
327    /// Invoked after the incoming thread becomes current. The runtime value
328    /// is the rq-charged total before its new execution interval, allowing OS
329    /// accounting to use the switch boundary without querying the registry.
330    pub on_switch_in: unsafe extern "Rust" fn(
331        data: usize,
332        thread: ThreadId,
333        policy: SchedulePolicy,
334        charged_runtime_ns: u64,
335    ),
336    /// Invoked after the thread stops being the current execution context.
337    pub on_switch_out: unsafe extern "Rust" fn(data: usize, thread: ThreadId, reason: SwitchReason),
338    /// Invoked in task context after the thread exits.
339    pub on_exit: unsafe extern "Rust" fn(data: usize, thread: ThreadId),
340    /// Invoked in task context for requested Deadline overrun notification.
341    pub on_deadline_overrun: unsafe extern "Rust" fn(data: usize, thread: ThreadId),
342    /// Releases the OS-owned extension data in task or reaper context.
343    pub drop: unsafe extern "Rust" fn(data: usize),
344}
345
346/// Bounded OS hook invoked when the owner changes a running thread's base policy.
347pub type RunningPolicyAppliedHook = unsafe extern "Rust" fn(
348    data: usize,
349    thread: ThreadId,
350    base_policy: SchedulePolicy,
351    observed_ns: u64,
352);
353
354/// Opaque OS-specific data attached to a thread.
355#[derive(Debug)]
356pub struct ThreadExtension {
357    data: usize,
358    ops: &'static ThreadExtensionOps,
359    running_policy_applied_hook: Option<RunningPolicyAppliedHook>,
360    scheduler_tick_cpu_time: Option<Arc<SchedulerTickCpuTime>>,
361    scheduler_tick_work: Option<SchedulerTickWork>,
362}
363
364impl ThreadExtension {
365    /// Creates an extension from opaque data and a static callback table.
366    ///
367    /// # Safety
368    ///
369    /// `data` must satisfy every callback contract in `ops`, and the owning OS
370    /// must ensure callbacks do not allocate, block, or re-enter the scheduler
371    /// when invoked as switch hooks. Task-context callbacks must return to the
372    /// dedicated service thread; abandoning that stack leaves their explicit
373    /// in-flight lifetime claim closed to prevent use-after-free.
374    pub const unsafe fn new(data: usize, ops: &'static ThreadExtensionOps) -> Self {
375        Self {
376            data,
377            ops,
378            running_policy_applied_hook: None,
379            scheduler_tick_cpu_time: None,
380            scheduler_tick_work: None,
381        }
382    }
383
384    /// Attaches IRQ-safe user/system CPU-time sampling to this thread.
385    ///
386    /// The scheduler retains the capability and charges it directly from each
387    /// periodic tick. No OS callback or deferred task work runs in hard IRQ.
388    pub fn with_scheduler_tick_cpu_time(mut self, accounting: Arc<SchedulerTickCpuTime>) -> Self {
389        self.scheduler_tick_cpu_time = Some(accounting);
390        self
391    }
392
393    /// Adds a bounded callback for base-policy changes applied to a running thread.
394    ///
395    /// The callback runs after the scheduler releases the thread-state lock.
396    /// The current CPU still owns the scheduler baton, so the callback is
397    /// serialized with switch hooks for the same thread. Queued and inactive
398    /// base-policy changes are observed through the policy snapshot passed to
399    /// the next switch-in instead. PI donation does not change this value.
400    ///
401    /// # Safety
402    ///
403    /// `callback` must interpret `data` according to this extension, remain
404    /// valid for its complete lifetime, and perform only bounded operations.
405    /// It must not allocate, block, or re-enter the scheduler.
406    pub unsafe fn with_running_policy_applied_hook(
407        mut self,
408        callback: RunningPolicyAppliedHook,
409    ) -> Self {
410        self.running_policy_applied_hook = Some(callback);
411        self
412    }
413
414    /// Adds task-context work gated by scheduler tick interest.
415    ///
416    /// The scheduler hard-IRQ path only publishes a typed deferred-work record.
417    /// The callback runs later on the dedicated task-work service thread.
418    ///
419    /// # Safety
420    ///
421    /// `callback` must interpret `data` according to this extension, remain
422    /// valid for its complete lifetime, and return normally to the task-work
423    /// service. The callback may use task-context synchronization but must not
424    /// retain the borrowed extension data after it returns. It may return
425    /// [`SchedulerTickWorkDisposition::Retry`] only after a transient conflict
426    /// and before publishing any accounting, timer, or signal state.
427    pub unsafe fn with_scheduler_tick_work(
428        mut self,
429        gate: Arc<SchedulerTickGate>,
430        callback: SchedulerTickTaskWork,
431    ) -> Self {
432        self.scheduler_tick_work = Some(SchedulerTickWork::new(gate, callback));
433        self
434    }
435
436    /// Returns the opaque OS-owned value.
437    pub const fn data(&self) -> usize {
438        self.data
439    }
440
441    /// Returns the callback table used as the extension type identity.
442    pub const fn ops(&self) -> &'static ThreadExtensionOps {
443        self.ops
444    }
445
446    /// Clones the IRQ-safe CPU-time sampling capability.
447    ///
448    /// Thread creation retains this capability alongside the extension.
449    pub fn scheduler_tick_cpu_time(&self) -> Option<Arc<SchedulerTickCpuTime>> {
450        self.scheduler_tick_cpu_time.as_ref().map(Arc::clone)
451    }
452
453    pub(crate) const fn as_view(&self) -> ThreadExtensionView {
454        ThreadExtensionView {
455            data: self.data,
456            ops: self.ops,
457            running_policy_applied_hook: self.running_policy_applied_hook,
458        }
459    }
460
461    pub(crate) fn scheduler_tick_work(&self) -> Option<SchedulerTickWork> {
462        self.scheduler_tick_work.clone()
463    }
464}
465
466impl Drop for ThreadExtension {
467    fn drop(&mut self) {
468        // SAFETY: construction transfers the unique callback-data destruction
469        // right into this non-cloneable owner.
470        unsafe { (self.ops.drop)(self.data) };
471    }
472}
473
474/// Copy-only borrowed identity for an installed OS extension.
475#[derive(Clone, Copy, Debug)]
476pub struct ThreadExtensionView {
477    data: usize,
478    ops: &'static ThreadExtensionOps,
479    running_policy_applied_hook: Option<RunningPolicyAppliedHook>,
480}
481
482/// Extension identity borrowed for exactly as long as a strong thread handle.
483///
484/// This wrapper deliberately does not expose its copyable internal view. The
485/// strong handle borrowed by the wrapper prevents the registry reaper from
486/// destroying the extension while its opaque data is being inspected.
487#[derive(Debug)]
488pub struct ThreadExtensionBorrow<'thread> {
489    view: ThreadExtensionView,
490    _thread: &'thread ThreadHandle,
491}
492
493impl<'thread> ThreadExtensionBorrow<'thread> {
494    pub(crate) const fn new(view: ThreadExtensionView, thread: &'thread ThreadHandle) -> Self {
495        Self {
496            view,
497            _thread: thread,
498        }
499    }
500
501    /// Returns the borrowed opaque data value.
502    pub const fn data(&self) -> usize {
503        self.view.data()
504    }
505
506    /// Returns the callback table used as the extension type identity.
507    pub const fn ops(&self) -> &'static ThreadExtensionOps {
508        self.view.ops()
509    }
510}
511
512/// Owned extension lease used when the caller has no pre-existing handle.
513///
514/// Keeping this value alive pins both the thread header and the registry record,
515/// so current-thread helpers cannot return data that becomes stale immediately
516/// after their temporary lookup handle is dropped.
517#[derive(Debug)]
518pub struct ThreadExtensionLease {
519    view: ThreadExtensionView,
520    thread: ThreadHandle,
521}
522
523impl ThreadExtensionLease {
524    pub(crate) const fn new(view: ThreadExtensionView, thread: ThreadHandle) -> Self {
525        Self { view, thread }
526    }
527
528    /// Returns the generation-bearing identity pinned by this lease.
529    pub fn thread_id(&self) -> ThreadId {
530        self.thread.id()
531    }
532
533    /// Returns the leased opaque data value.
534    pub const fn data(&self) -> usize {
535        self.view.data()
536    }
537
538    /// Returns the callback table used as the extension type identity.
539    pub const fn ops(&self) -> &'static ThreadExtensionOps {
540        self.view.ops()
541    }
542}
543
544impl ThreadExtensionView {
545    /// Returns the borrowed opaque data value.
546    pub const fn data(self) -> usize {
547        self.data
548    }
549
550    /// Returns the callback table used as the extension type identity.
551    pub const fn ops(self) -> &'static ThreadExtensionOps {
552        self.ops
553    }
554
555    pub(crate) unsafe fn notify_running_policy_applied(
556        self,
557        thread: ThreadId,
558        base_policy: SchedulePolicy,
559        observed_ns: u64,
560    ) {
561        if let Some(callback) = self.running_policy_applied_hook {
562            unsafe { callback(self.data, thread, base_policy, observed_ns) };
563        }
564    }
565}
566
567/// Validated inputs used to create a scheduler thread record.
568#[derive(Debug)]
569pub struct ThreadSpec {
570    pub(crate) execution: Option<Arc<crate::thread::execution::ThreadExecution>>,
571    policy: SchedulePolicy,
572    affinity: Option<CpuSet>,
573    // Runtime resources must be dropped before the extension that owns their
574    // address-space and entry metadata, including on fallback destruction.
575    resources: ThreadResources,
576    extension: Option<ThreadExtension>,
577}
578
579impl ThreadSpec {
580    /// Creates a thread specification with full topology affinity.
581    pub const fn new(policy: SchedulePolicy) -> Self {
582        Self {
583            execution: None,
584            policy,
585            affinity: None,
586            resources: ThreadResources::NONE,
587            extension: None,
588        }
589    }
590
591    /// Restricts the thread to an explicit CPU set.
592    pub fn with_affinity(mut self, affinity: CpuSet) -> Self {
593        self.affinity = Some(affinity);
594        self
595    }
596
597    /// Attaches OS-specific state.
598    pub fn with_extension(mut self, extension: ThreadExtension) -> Self {
599        self.extension = Some(extension);
600        self
601    }
602
603    /// Associates a complete runtime resource bundle with the thread.
604    ///
605    /// # Safety
606    ///
607    /// `resources` must satisfy [`ThreadResources::new`] and must be consumed by
608    /// exactly this specification and its eventual scheduler record.
609    pub unsafe fn with_resources(mut self, resources: ThreadResources) -> Self {
610        self.resources = resources;
611        self
612    }
613
614    /// Returns the base scheduling policy.
615    pub const fn policy(&self) -> SchedulePolicy {
616        self.policy
617    }
618
619    /// Returns explicit affinity, if one was supplied.
620    pub fn affinity(&self) -> Option<&CpuSet> {
621        self.affinity.as_ref()
622    }
623
624    pub(crate) fn take_affinity(&mut self) -> Option<CpuSet> {
625        self.affinity.take()
626    }
627    pub(crate) fn resources(&self) -> &ThreadResources {
628        &self.resources
629    }
630    pub(crate) fn extension(&self) -> Option<&ThreadExtension> {
631        self.extension.as_ref()
632    }
633
634    pub(crate) fn into_owned_parts(mut self) -> (Option<ThreadExtension>, ThreadResources) {
635        let extension = self.extension.take();
636        let resources = core::mem::replace(&mut self.resources, ThreadResources::NONE);
637        (extension, resources)
638    }
639}
ax_task/thread/spec.rs

ax_task/thread/
spec.rs
