ax_task/sched/system/task_system/
thread_api.rs

1//! Generation-checked thread inspection and policy updates.
2
3use super::*;
4
5impl TaskSystem {
6    /// Returns the current state of a live registry entry.
7    pub fn thread_state(&self, thread: ThreadId) -> Result<ThreadState, TaskError> {
8        Ok(self
9            .state
10            .lock()
11            .thread_record(thread)?
12            .sched
13            .lock()
14            .lifecycle
15            .state())
16    }
17
18    /// Returns cumulative charged CPU runtime.
19    ///
20    /// Like Linux `task_sched_runtime()`, a running thread is sampled only
21    /// after locking its assigned runqueue and updating that runqueue's clock.
22    /// A stopped thread returns its already charged value without inventing a
23    /// scheduler timestamp.
24    pub fn thread_runtime(&self, thread: ThreadId) -> Result<ThreadRuntimeSnapshot, TaskError> {
25        let (core, sched_cell) = {
26            let state = self.state.lock();
27            let record = state.thread_record(thread)?;
28            (Arc::clone(&record.core), Arc::clone(&record.sched))
29        };
30        let sched = sched_cell.lock();
31        let snapshot = if let Some(cpu) = sched.placement.assigned_cpu() {
32            let remote = self
33                .cpu_remotes
34                .get(cpu.as_usize())
35                .ok_or(TaskError::InvalidCpu(cpu.as_u32()))?;
36            let transaction = OwnerRqTxn::begin(self, remote);
37            let rq_state = transaction.task_state(thread, &sched.placement);
38            let running_interval_ns = if rq_state.is_current() {
39                let dispatch = transaction
40                    .current()
41                    .filter(|dispatch| dispatch.thread() == thread);
42                Some(
43                    dispatch
44                        .unwrap_or_else(|| {
45                            task_runtime::fatal_invariant(0x5251_1210, thread.as_u64() as usize)
46                        })
47                        .runtime_interval_ns(transaction.clock().task().as_nanos()),
48                )
49            } else {
50                None
51            };
52            let snapshot = core.runtime_snapshot(running_interval_ns);
53            transaction.commit();
54            snapshot
55        } else {
56            core.runtime_snapshot(None)
57        };
58        Ok(snapshot)
59    }
60
61    /// Replaces the current running thread's opaque address-space token.
62    ///
63    /// The caller must hold the owner CPU's IRQ-off scheduler-safe window. This
64    /// operation updates only scheduler metadata; installing the hardware page
65    /// table and invalidating translations remain runtime responsibilities.
66    pub fn replace_current_address_space(
67        &self,
68        cpu: Pin<&mut CpuLocal>,
69        address_space: &mut crate::runtime::resource::AddressSpaceToken,
70    ) -> Result<crate::runtime::resource::AddressSpaceToken, TaskError> {
71        self.ensure_owner_cpu_context(&cpu)?;
72        if address_space.is_none() {
73            return Err(TaskError::InvalidConfiguration);
74        }
75        let mut state = self.state.lock();
76        state.ensure_cpu_online(&cpu)?;
77        let owner = cpu.owner();
78        let current = cpu.current().ok_or(TaskError::NoRunnableThread)?;
79        let record = state.thread_record_mut(current)?;
80        let mut sched = record.sched.lock();
81        if sched.lifecycle.state() != ThreadState::Running
82            || sched.placement.queued_cpu() != Some(owner)
83            || sched.placement.on_cpu() != Some(owner)
84        {
85            return Err(TaskError::InvalidConfiguration);
86        }
87        let next_handle = address_space.handle();
88        let next_membarrier_state = task_runtime::address_space_membarrier_state(next_handle);
89        let binding =
90            crate::runtime::switch::ThreadRuntimeBinding::new(sched.runtime.context, next_handle);
91        let remote = Arc::clone(cpu.remote());
92        let mut transaction = OwnerRqTxn::begin(self, &remote);
93        transaction.update_current_runtime_binding(current, binding, next_membarrier_state);
94        let next = core::mem::replace(
95            address_space,
96            crate::runtime::resource::AddressSpaceToken::NONE,
97        );
98        transaction.commit();
99        let previous = record.resources.replace_address_space(next);
100        sched.runtime.address_space = next_handle;
101        Ok(previous)
102    }
103
104    /// Detaches the current running thread from its user address space.
105    ///
106    /// The caller must enter the runtime's lazy kernel address-space state in
107    /// the same IRQ-off transaction before releasing the returned token.
108    pub fn detach_current_address_space(
109        &self,
110        cpu: Pin<&mut CpuLocal>,
111    ) -> Result<crate::runtime::resource::AddressSpaceToken, TaskError> {
112        self.ensure_owner_cpu_context(&cpu)?;
113        let mut state = self.state.lock();
114        state.ensure_cpu_online(&cpu)?;
115        let owner = cpu.owner();
116        let current = cpu.current().ok_or(TaskError::NoRunnableThread)?;
117        let record = state.thread_record_mut(current)?;
118        let mut sched = record.sched.lock();
119        if sched.lifecycle.state() != ThreadState::Running
120            || sched.placement.queued_cpu() != Some(owner)
121            || sched.placement.on_cpu() != Some(owner)
122            || record.resources.address_space().is_none()
123        {
124            return Err(TaskError::InvalidConfiguration);
125        }
126        let binding = crate::runtime::switch::ThreadRuntimeBinding::new(
127            sched.runtime.context,
128            crate::runtime::resource::AddressSpaceHandle::NONE,
129        );
130        let remote = Arc::clone(cpu.remote());
131        let mut transaction = OwnerRqTxn::begin(self, &remote);
132        transaction.update_current_runtime_binding(
133            current,
134            binding,
135            crate::runtime::resource::AddressSpaceMembarrierState::NONE,
136        );
137        let previous = record.resources.take_address_space();
138        transaction.commit();
139        sched.runtime.address_space = crate::runtime::resource::AddressSpaceHandle::NONE;
140        Ok(previous)
141    }
142
143    /// Acquires a strong handle for a generation-valid registry entry.
144    pub fn thread_handle(&self, thread: ThreadId) -> Result<ThreadHandle, TaskError> {
145        let state = self.state.lock();
146        let record = state.thread_record(thread)?;
147        Ok(ThreadHandle::from_core(Arc::clone(&record.core)))
148    }
149
150    /// Borrows the opaque OS extension through a generation-valid strong handle.
151    ///
152    /// The borrow cannot outlive `handle`, which prevents the registry reaper
153    /// from releasing the extension data while a caller interprets it.
154    pub fn thread_extension<'thread>(
155        &self,
156        handle: &'thread ThreadHandle,
157    ) -> Result<Option<ThreadExtensionBorrow<'thread>>, TaskError> {
158        let view = self.thread_extension_view(handle)?;
159        Ok(view.map(|view| ThreadExtensionBorrow::new(view, handle)))
160    }
161
162    /// Acquires an owned lease for callers that looked up a temporary handle.
163    pub fn thread_extension_lease(
164        &self,
165        handle: ThreadHandle,
166    ) -> Result<Option<ThreadExtensionLease>, TaskError> {
167        let view = self.thread_extension_view(&handle)?;
168        Ok(view.map(|view| ThreadExtensionLease::new(view, handle)))
169    }
170
171    fn thread_extension_view(
172        &self,
173        handle: &ThreadHandle,
174    ) -> Result<Option<ThreadExtensionView>, TaskError> {
175        let state = self.state.lock();
176        let record = state.thread_record(handle.id())?;
177        if !Arc::ptr_eq(&record.core, &handle.core) {
178            return Err(TaskError::StaleThreadId);
179        }
180        Ok(handle.extension_view())
181    }
182
183    /// Replaces a task's base policy in one synchronous owner-rq transaction.
184    pub fn set_thread_policy(
185        &self,
186        thread: ThreadId,
187        policy: SchedulePolicy,
188    ) -> Result<(), TaskError> {
189        policy.validate()?;
190        // Allocate the affinity snapshot before entering IRQ-disabled cold
191        // domains. Copying into this fixed-topology buffer is allocation-free.
192        let mut affinity = CpuSet::empty(self.config.cpu_count());
193        let core = {
194            let state = self.state.lock();
195            Arc::clone(&state.thread_record(thread)?.core)
196        };
197        // Serialize the complete policy/admission/rq transaction against
198        // current-thread exit. Linux holds the task's PI lifetime lock before
199        // task_rq_lock(); a policy writer that loses this edge must not mutate
200        // base policy or root-domain bandwidth for an exiting task.
201        let _activity = core.try_scheduler_activity().ok_or(TaskError::NotReady)?;
202        let state = self.state.lock();
203        let mut root_domain = self.root_domain.lock();
204        let record = state.thread_record(thread)?;
205        if !Arc::ptr_eq(&record.core, &core) {
206            return Err(TaskError::StaleThreadId);
207        }
208        let sched_cell = Arc::clone(&record.sched);
209        let mut sched = sched_cell.lock();
210        if sched.lifecycle.state() == ThreadState::Exited {
211            return Err(TaskError::NotReady);
212        }
213        affinity.copy_from_set(&sched.affinity.affinity)?;
214        let applied_reservation = sched.deadline.bandwidth.reservation_scaled();
215        let pending_reservation = sched
216            .policy
217            .pending_update()
218            .map_or(0, |pending| pending.reservation_scaled);
219        // Linux serializes sched_setscheduler() against PI and wakeup by
220        // retaining p->pi_lock through task_rq_lock() and the class change.
221        // Keeping this scheduler guard across the owner-rq transaction also
222        // serializes concurrent policy writers; no generation can consume a
223        // later writer's pending value.
224        let owner = sched
225            .placement
226            .assigned_cpu()
227            .ok_or(TaskError::InvalidPiState)?;
228        let reservation_owner = sched.deadline.bandwidth.reservation_owner();
229        if let Some(reservation_owner) = reservation_owner
230            && owner != reservation_owner
231        {
232            task_runtime::fatal_invariant(0x444c_1201, core.id().as_u64() as usize);
233        }
234        let remote = self
235            .cpu_remotes
236            .get(owner.as_usize())
237            .ok_or(TaskError::InvalidCpu(owner.as_u32()))?;
238        let reservation = root_domain.deadline_reservation_for(policy, &affinity)?;
239        let pending = sched.policy.prepare_update(policy, reservation)?;
240        let old_held = applied_reservation.max(pending_reservation);
241        let new_held = applied_reservation.max(reservation);
242        root_domain.replace_deadline_utilization(old_held, new_held)?;
243        sched.policy.publish_update(pending);
244        drop(state);
245
246        let applied = self
247            .apply_owner_policy_update_locked(remote, &core, &mut sched, pending.generation)
248            .unwrap_or_else(|_| {
249                task_runtime::fatal_invariant(0x5251_1208, core.id().as_u64() as usize)
250            });
251        Self::finish_policy_admission_locked(&mut root_domain, &core, applied.commit);
252        drop(root_domain);
253        drop(sched);
254        Self::notify_policy_generation(&core, applied.commit);
255        self.recompute_pi_after_policy_update(core.id())
256            .unwrap_or_else(|_| {
257                task_runtime::fatal_invariant(0x5049_1216, core.id().as_u64() as usize)
258            });
259
260        let owner_work_required =
261            applied.scheduler_deadline_refresh_required || applied.rt_period_started;
262        match (applied.reschedule, owner_work_required) {
263            (Some(kind), true) => {
264                // Preemption and owner-deadline facts belong to one rq transaction.
265                // Publish both logical reasons before a single physical edge.
266                remote.request_remote_reschedule_with_scheduler_work(kind);
267            }
268            (Some(kind), false) => {
269                remote.request_remote_reschedule(kind);
270            }
271            (None, true) => {
272                // Scheduler deadlines are pinned to the rq owner. Ask that
273                // owner to derive its physical timer; a remote setter must
274                // not program another CPU's comparator directly.
275                remote.kick_scheduler_work();
276            }
277            (None, false) => {}
278        }
279        Ok(())
280    }
281
282    /// Returns a copy of the thread CPU affinity mask.
283    pub fn thread_affinity(&self, thread: ThreadId) -> Result<CpuSet, TaskError> {
284        Ok(self
285            .state
286            .lock()
287            .thread_record(thread)?
288            .sched
289            .lock()
290            .affinity
291            .affinity
292            .as_ref()
293            .clone())
294    }
295}
ax_task/sched/system/task_system/thread_api.rs

ax_task/sched/system/task_system/
thread_api.rs
