ax_task/sched/
affinity.rs

1//! Completion ownership for asynchronous remote affinity changes.
2
3use alloc::sync::Arc;
4use core::sync::atomic::{AtomicU64, Ordering};
5
6use crate::{
7    sync::WaitQueue,
8    thread::{TaskError, ThreadCore, ThreadHandle, ThreadState},
9};
10
11/// Per-thread completion sequence shared by concurrent affinity setters.
12#[derive(Debug)]
13pub(crate) struct ThreadAffinityCompletion {
14    completed_generation: AtomicU64,
15    waiters: WaitQueue,
16}
17
18impl ThreadAffinityCompletion {
19    pub(crate) const fn new(completed_generation: u64) -> Self {
20        Self {
21            completed_generation: AtomicU64::new(completed_generation),
22            waiters: WaitQueue::new(),
23        }
24    }
25
26    pub(crate) fn publish(&self, generation: u64) -> bool {
27        let mut completed = self.completed_generation.load(Ordering::Acquire);
28        loop {
29            if completed >= generation {
30                return false;
31            }
32            match self.completed_generation.compare_exchange_weak(
33                completed,
34                generation,
35                Ordering::Release,
36                Ordering::Acquire,
37            ) {
38                Ok(_) => return true,
39                Err(observed) => completed = observed,
40            }
41        }
42    }
43
44    pub(crate) fn completed_generation(&self) -> u64 {
45        self.completed_generation.load(Ordering::Acquire)
46    }
47
48    pub(crate) fn notify_waiters(&self) {
49        self.waiters.notify_all();
50    }
51
52    fn wait_for(&self, request: &ThreadAffinityChange) -> Result<(), TaskError> {
53        self.waiters
54            .try_wait_until(|| request.try_result().is_some())?;
55        request
56            .try_result()
57            .expect("affinity wait predicate resolved the request")
58    }
59}
60
61/// Move-only completion for one generation of a remote affinity change.
62#[derive(Debug)]
63#[must_use = "dropping the change leaves the affinity update asynchronous"]
64pub struct ThreadAffinityChange {
65    thread: ThreadHandle,
66    generation: u64,
67}
68
69impl ThreadAffinityChange {
70    pub(crate) fn new(core: Arc<ThreadCore>, generation: u64) -> Self {
71        Self {
72            thread: ThreadHandle::from_core(core),
73            generation,
74        }
75    }
76
77    /// Returns the generation assigned to this affinity change.
78    pub const fn generation(&self) -> u64 {
79        self.generation
80    }
81
82    /// Observes completion, target exit, or a still-pending owner transition.
83    pub fn try_result(&self) -> Option<Result<(), TaskError>> {
84        if self
85            .thread
86            .core
87            .affinity_completion()
88            .completed_generation()
89            >= self.generation
90        {
91            Some(Ok(()))
92        } else if self.thread.state() == ThreadState::Exited {
93            Some(Err(TaskError::StaleThreadId))
94        } else {
95            None
96        }
97    }
98
99    /// Sleeps until the owner runqueue orders this generation.
100    ///
101    /// A later concurrent affinity update may supersede this request. In that
102    /// case both requests complete after the owner has installed the latest
103    /// placement, matching Linux's shared pending-affinity completion.
104    ///
105    /// # Errors
106    ///
107    /// Returns [`TaskError::StaleThreadId`] if the target exits before the
108    /// generation completes, and propagates scheduler blocking failures.
109    pub fn wait(self) -> Result<(), TaskError> {
110        self.thread.core.affinity_completion().wait_for(&self)
111    }
112}
ax_task/sched/affinity.rs

ax_task/sched/
affinity.rs
