ax_runtime/thread/lifecycle/
publication.rs

1//! Prepared-thread placement and staged activation transaction.
2
3use super::*;
4
5impl PreparedThread {
6    /// Returns a strong handle for binding OS-owned identity before publication.
7    pub fn thread_handle(&self) -> ThreadHandle {
8        self.handle
9            .as_ref()
10            .expect("prepared thread was already consumed")
11            .clone()
12    }
13
14    /// Places and immediately activates a thread with no external publication
15    /// transaction.
16    pub fn publish(self) -> Result<ThreadHandle, TaskError> {
17        Ok(self.stage()?.activate())
18    }
19
20    /// Completes the fallible scheduler placement phase without entering the
21    /// caller-owned thread entry point.
22    ///
23    /// The scheduler may select the staged thread, but its runtime trampoline
24    /// remains blocked on an internal start gate. This lets an OS complete its
25    /// public identity transaction before [`StagedThread::activate`] provides
26    /// the final infallible release, matching Linux's `wake_up_new_task`
27    /// boundary.
28    pub fn stage(mut self) -> Result<StagedThread, TaskError> {
29        let handle = self
30            .handle
31            .take()
32            .expect("prepared thread was already consumed");
33        publish_prepared_thread(self.system, handle).map(|handle| StagedThread {
34            handle: Some(handle),
35            start: Arc::clone(&self.start),
36        })
37    }
38
39    pub(in crate::thread) fn new(
40        system: &'static TaskSystem,
41        handle: ThreadHandle,
42        start: Arc<RuntimeThreadStart>,
43    ) -> Self {
44        Self {
45            system,
46            handle: Some(handle),
47            start,
48        }
49    }
50}
51
52impl StagedThread {
53    /// Returns a strong handle for the OS publication transaction.
54    pub fn thread_handle(&self) -> ThreadHandle {
55        self.handle
56            .as_ref()
57            .expect("staged thread was already consumed")
58            .clone()
59    }
60
61    /// Releases the staged thread to execute its caller-owned entry point.
62    pub fn activate(mut self) -> ThreadHandle {
63        let handle = self
64            .handle
65            .take()
66            .expect("staged thread was already consumed");
67        self.start.activate();
68        handle
69    }
70}
71
72impl Drop for StagedThread {
73    fn drop(&mut self) {
74        self.start.abort();
75    }
76}
77
78impl Drop for PreparedThread {
79    fn drop(&mut self) {
80        if let Some(handle) = self.handle.take() {
81            cleanup_failed_thread(self.system, handle);
82        }
83    }
84}
85
86impl RuntimeThreadStart {
87    pub(in crate::thread) const fn new() -> Self {
88        Self {
89            state: AtomicU8::new(THREAD_START_PENDING),
90            wait: WaitQueue::new(),
91        }
92    }
93
94    fn activate(&self) {
95        self.state
96            .compare_exchange(
97                THREAD_START_PENDING,
98                THREAD_START_ACTIVE,
99                Ordering::AcqRel,
100                Ordering::Acquire,
101            )
102            .unwrap_or_else(|state| panic!("invalid staged-thread activation state: {state}"));
103        self.wait.notify_all();
104    }
105
106    fn abort(&self) {
107        if self
108            .state
109            .compare_exchange(
110                THREAD_START_PENDING,
111                THREAD_START_ABORTED,
112                Ordering::AcqRel,
113                Ordering::Acquire,
114            )
115            .is_ok()
116        {
117            self.wait.notify_all();
118        }
119    }
120
121    pub(in crate::thread) fn wait_for_activation(&self) -> bool {
122        match self.state.load(Ordering::Acquire) {
123            THREAD_START_ACTIVE => return true,
124            THREAD_START_ABORTED => return false,
125            THREAD_START_PENDING => {}
126            state => panic!("invalid runtime thread-start state: {state}"),
127        }
128        self.wait
129            .wait_until(|| self.state.load(Ordering::Acquire) != THREAD_START_PENDING);
130        match self.state.load(Ordering::Acquire) {
131            THREAD_START_ACTIVE => true,
132            THREAD_START_ABORTED => false,
133            state => panic!("invalid completed thread-start state: {state}"),
134        }
135    }
136}
137
138fn publish_prepared_thread(
139    system: &'static TaskSystem,
140    handle: ThreadHandle,
141) -> Result<ThreadHandle, TaskError> {
142    let result = with_current_cpu_local_mut_owner(|cpu| system.start_thread(cpu, handle.id()));
143    if let Err(error) = result {
144        cleanup_failed_thread(system, handle);
145        return Err(error);
146    }
147    Ok(handle)
148}
149
150fn cleanup_failed_thread(system: &TaskSystem, handle: ThreadHandle) {
151    let thread = handle.id();
152    let _ = system.mark_exited(thread);
153    drop(handle);
154    let _ = system.reap_thread(thread);
155}
ax_runtime/thread/lifecycle/publication.rs

ax_runtime/thread/lifecycle/
publication.rs
