ax_task/sched/system/task_system/
thread_creation.rs

1//! Transactional thread creation and initial CPU binding.
2
3use super::*;
4
5#[derive(Clone, Copy)]
6enum ThreadCreationContext {
7    Runtime,
8    OfflineBootstrap,
9}
10
11impl TaskSystem {
12    /// Creates a thread in the [`ThreadState::New`] state.
13    ///
14    /// Deadline threads are admitted immediately and therefore must cover the
15    /// complete online root domain.
16    pub fn create_thread(&self, spec: ThreadSpec) -> Result<ThreadHandle, TaskError> {
17        // SAFETY: the runtime publishes the calling CPU identity before task
18        // creation is enabled. Like Linux fork, this establishes task_cpu()
19        // before the new task can participate in PI or become runnable.
20        let initial_cpu = CpuId::new(unsafe { task_runtime::current_cpu_id() }.as_u32());
21        self.create_thread_on_cpu(spec, initial_cpu, ThreadCreationContext::Runtime)
22    }
23
24    /// Builds an unpublished task with an explicit initial `task_cpu`.
25    ///
26    /// Ordinary fork uses the calling CPU. Per-CPU bootstrap and idle tasks
27    /// instead mirror Linux `init_idle()` and bind the target rq before the
28    /// task can be observed by PI, policy, or hotplug code.
29    fn create_thread_on_cpu(
30        &self,
31        spec: ThreadSpec,
32        initial_cpu: CpuId,
33        context: ThreadCreationContext,
34    ) -> Result<ThreadHandle, TaskError> {
35        if initial_cpu.as_usize() >= self.config.cpu_count() {
36            return Err(TaskError::InvalidCpu(initial_cpu.as_u32()));
37        }
38        let policy = spec.policy();
39        let affinity = spec
40            .affinity()
41            .cloned()
42            .unwrap_or_else(|| CpuSet::all(self.config.cpu_count()));
43        let unpublished = UnpublishedThreadGuard::new(self, spec);
44        policy.validate()?;
45        validate_affinity(&affinity, self.config.cpu_count())?;
46        let (slot, generation, reservation) = {
47            let mut state = self.state.lock();
48            let mut root_domain = self.root_domain.lock();
49            let reservation = root_domain.reserve_deadline(policy, &affinity)?;
50            let (slot, generation) = match state.allocate_thread_slot(self.config.thread_capacity())
51            {
52                Ok(identity) => identity,
53                Err(error) => {
54                    root_domain.release_deadline(reservation);
55                    return Err(error);
56                }
57            };
58            state.slots[slot as usize].pending_deadline_reservation = reservation;
59            (slot, generation, reservation)
60        };
61        let id = ThreadId::from_parts(slot, generation);
62
63        // Linux embeds class nodes in task_struct before publication. Prepare
64        // the Rust class-node indexes at the same cold construction boundary,
65        // so a first wake or cross-CPU migration cannot allocate under rq
66        // irqsave locks.
67        for remote in &self.cpu_remotes {
68            let mut run_queue = match context {
69                ThreadCreationContext::Runtime => {
70                    remote.lock_run_queue(RunQueueGuardSource::Lifecycle)
71                }
72                ThreadCreationContext::OfflineBootstrap => {
73                    // SAFETY: per-CPU bootstrap retains raw IRQ exclusion and
74                    // PREEMPT_DISABLED until the complete rq/current/idle
75                    // owner is published.
76                    unsafe { remote.lock_run_queue_irq_disabled() }
77                }
78            };
79            run_queue.prepare_thread_slot(slot as usize);
80        }
81
82        // Runtime construction may allocate, fault, or call into platform
83        // code. Keep it outside the IRQ-disabled registry domain. The removed
84        // slot is a private reservation until the short commit below.
85        let deadline_server = DeadlineServer::unbound();
86        let entity = SchedulingEntity::new_with_deadline_server(
87            policy,
88            self.config.fair_slice_ns(),
89            0,
90            deadline_server.clone(),
91        );
92        let (extension, resources) = unpublished.into_owned_parts();
93        let switch_extension = extension.as_ref().map(ThreadExtension::as_view);
94        let scheduler_tick_cpu_time = extension
95            .as_ref()
96            .and_then(ThreadExtension::scheduler_tick_cpu_time);
97        let scheduler_tick_work = extension
98            .as_ref()
99            .and_then(ThreadExtension::scheduler_tick_work);
100        let address_space = resources.address_space();
101        let membarrier_identity = if address_space.is_none() {
102            crate::runtime::resource::AddressSpaceMembarrierId::NONE
103        } else {
104            task_runtime::address_space_membarrier_state(address_space).identity()
105        };
106        let sched = Arc::new(ThreadSchedCell::new(
107            id,
108            ThreadSchedInit {
109                policy: ThreadPolicyInit { policy, entity },
110                placement: ThreadPlacementInit {
111                    initial_cpu,
112                    affinity: affinity.clone(),
113                },
114                deadline: ThreadDeadlineInit {
115                    server: deadline_server,
116                    reservation_scaled: reservation,
117                },
118                runtime: ThreadRuntimeInit {
119                    context: resources.context(),
120                    address_space,
121                },
122            },
123        ));
124        let core = Arc::new(ThreadCore::new(ThreadCoreInit {
125            id,
126            policy,
127            sched: Arc::clone(&sched),
128            extension: switch_extension,
129            scheduler_tick_cpu_time,
130            scheduler_tick_work,
131            membarrier_identity,
132            task_work: Some(Arc::clone(&self.task_work)),
133        }));
134        let record = ThreadRecord {
135            core: Arc::clone(&core),
136            sched,
137            resources,
138            extension,
139            callbacks: ThreadCallbackState::new(),
140        };
141        let context = record.resources.context();
142        if !context.is_none() {
143            let status = task_runtime::bind_context_thread(ContextThreadBinding {
144                context,
145                publication: CurrentThreadPublication::from_core(id, &core),
146            });
147            if status != RuntimeStatus::Success {
148                {
149                    let mut state = self.state.lock();
150                    let mut root_domain = self.root_domain.lock();
151                    let failed_slot = &mut state.slots[slot as usize];
152                    debug_assert_eq!(failed_slot.generation, generation);
153                    debug_assert!(failed_slot.record.is_none());
154                    debug_assert_eq!(failed_slot.pending_deadline_reservation, reservation);
155                    failed_slot.pending_deadline_reservation = 0;
156                    if advance_thread_slot_generation(failed_slot) {
157                        state.free_slots.push(slot);
158                    }
159                    root_domain.release_deadline(reservation);
160                }
161                drop(core);
162                self.release_thread_record(record);
163                return Err(TaskError::RuntimeFailure(status as u32));
164            }
165        }
166
167        let mut record = Some(record);
168        let commit_error = {
169            let mut state = self.state.lock();
170            let mut root_domain = self.root_domain.lock();
171            let is_deadline = matches!(policy, SchedulePolicy::Deadline(_));
172            let topology_rejects_deadline = is_deadline && !affinity.covers(&root_domain.online);
173            let admission_overcommitted = is_deadline && root_domain.admission_overcommitted();
174            if topology_rejects_deadline || admission_overcommitted {
175                let failed_slot = &mut state.slots[slot as usize];
176                debug_assert_eq!(failed_slot.generation, generation);
177                debug_assert!(failed_slot.record.is_none());
178                debug_assert_eq!(failed_slot.pending_deadline_reservation, reservation);
179                failed_slot.pending_deadline_reservation = 0;
180                if advance_thread_slot_generation(failed_slot) {
181                    state.free_slots.push(slot);
182                }
183                root_domain.release_deadline(reservation);
184                Some(if topology_rejects_deadline {
185                    TaskError::DeadlineAffinity
186                } else {
187                    TaskError::DeadlineAdmission
188                })
189            } else {
190                let reserved_slot = &mut state.slots[slot as usize];
191                debug_assert_eq!(reserved_slot.generation, generation);
192                debug_assert!(reserved_slot.record.is_none());
193                debug_assert_eq!(reserved_slot.pending_deadline_reservation, reservation);
194                reserved_slot.pending_deadline_reservation = 0;
195                reserved_slot.record = record.take();
196                None
197            }
198        };
199        if let Some(error) = commit_error {
200            drop(core);
201            self.release_thread_record(
202                record.expect("rejected thread commit must retain its resource record"),
203            );
204            return Err(error);
205        }
206        Ok(ThreadHandle::from_core(core))
207    }
208
209    /// Performs the initial runnable transition before the owner CPU is online.
210    ///
211    /// # Safety
212    ///
213    /// The caller must retain the boot CPU's raw IRQ exclusion and
214    /// `PREEMPT_DISABLED` ownership.
215    unsafe fn make_ready_bootstrap(&self, thread: ThreadId) -> Result<(), TaskError> {
216        let state = self.state.lock();
217        let record = state.thread_record(thread)?;
218        // SAFETY: forwarded from this method's offline boot-owner contract.
219        let mut sched = unsafe { record.sched.lock_bootstrap() };
220        sched.transition(&record.core, ThreadState::Running)
221    }
222
223    /// Installs the CPU's already-running bootstrap execution context.
224    ///
225    /// This operation is used before a CPU is published online and performs no
226    /// context switch. The runtime must call it exactly once with an empty
227    /// `CpuLocal` current slot.
228    pub fn install_bootstrap_thread(
229        &self,
230        mut cpu: Pin<&mut CpuLocal>,
231        spec: ThreadSpec,
232    ) -> Result<ThreadHandle, TaskError> {
233        let unpublished = UnpublishedThreadGuard::new(self, spec);
234        self.ensure_owner_cpu_context(&cpu)?;
235        if !matches!(
236            unpublished.spec().policy(),
237            SchedulePolicy::Fair {
238                mode: FairMode::Normal | FairMode::Batch,
239                ..
240            }
241        ) {
242            return Err(TaskError::InvalidConfiguration);
243        }
244        {
245            let state = self.state.lock();
246            let registration = state.cpu_registration(cpu.owner())?;
247            if !Arc::ptr_eq(&registration.remote, cpu.remote()) {
248                return Err(TaskError::InvalidRuntimeHandle);
249            }
250            // SAFETY: install_bootstrap_thread is an offline owner operation;
251            // its caller retains the boot CPU's raw IRQ exclusion.
252            if unsafe { cpu.remote().lock_run_queue_irq_disabled() }
253                .current_thread()
254                .is_some()
255            {
256                return Err(TaskError::InvalidConfiguration);
257            }
258        }
259
260        let thread = self.create_thread_on_cpu(
261            unpublished.into_spec(),
262            cpu.owner(),
263            ThreadCreationContext::OfflineBootstrap,
264        )?;
265        let setup = (|| {
266            let core = {
267                let state = self.state.lock();
268                Arc::clone(&state.thread_record(thread.id())?.core)
269            };
270            // SAFETY: the CPU is still offline under the boot owner's raw IRQ
271            // exclusion.
272            let mut sched = unsafe { core.sched().lock_bootstrap() };
273            sched.transition(&core, ThreadState::Running)?;
274            let remote = Arc::clone(cpu.remote());
275            // SAFETY: the CPU is still offline under the boot owner's raw IRQ
276            // exclusion and cannot enter the runtime IRQ-exit service.
277            let mut transaction = unsafe { OwnerRqTxn::begin_bootstrap(self, &remote) };
278            let _enqueue_consumed_by_immediate_bootstrap_pick = self
279                .link_owner_ready_thread_locked(
280                    cpu.owner(),
281                    &mut transaction,
282                    &core,
283                    &mut sched,
284                    EnqueueReason::Wake,
285                );
286            let next = self.pick_owner_bootstrap_in_rq(cpu.as_mut(), &mut transaction);
287            if !core::ptr::eq(next.core.as_ref(), Arc::as_ref(&core)) {
288                task_runtime::fatal_invariant(0x4254_0001, core.id().as_u64() as usize);
289            }
290            transaction.commit_bootstrap();
291            Ok(())
292        })();
293        if let Err(error) = setup {
294            return match self.discard_unpublished_thread(thread) {
295                Ok(()) => Err(error),
296                Err(cleanup_error) => Err(cleanup_error),
297            };
298        }
299        Ok(thread)
300    }
301
302    /// Creates and registers a dedicated CPU idle thread before online publish.
303    pub fn register_idle_thread(
304        &self,
305        mut cpu: Pin<&mut CpuLocal>,
306        spec: ThreadSpec,
307    ) -> Result<ThreadHandle, TaskError> {
308        let unpublished = UnpublishedThreadGuard::new(self, spec);
309        self.ensure_owner_cpu_context(&cpu)?;
310        if !matches!(
311            unpublished.spec().policy(),
312            SchedulePolicy::Fair {
313                mode: crate::sched::FairMode::Idle,
314                ..
315            }
316        ) {
317            return Err(TaskError::InvalidConfiguration);
318        }
319        {
320            let state = self.state.lock();
321            let registration = state.cpu_registration(cpu.owner())?;
322            if !Arc::ptr_eq(&registration.remote, cpu.remote()) {
323                return Err(TaskError::InvalidRuntimeHandle);
324            }
325            // SAFETY: register_idle_thread runs in the same offline bootstrap
326            // owner transaction as install_bootstrap_thread.
327            if unsafe { cpu.remote().lock_run_queue_irq_disabled() }
328                .idle()
329                .is_some()
330            {
331                return Err(TaskError::InvalidConfiguration);
332            }
333        }
334
335        let thread = self.create_thread_on_cpu(
336            unpublished.into_spec(),
337            cpu.owner(),
338            ThreadCreationContext::OfflineBootstrap,
339        )?;
340        // SAFETY: the target CPU remains offline and boot-owned until idle is
341        // installed and the complete runtime endpoint is published.
342        let setup = unsafe { self.make_ready_bootstrap(thread.id()) }.and_then(|()| {
343            let state = self.state.lock();
344            let core = Arc::clone(&state.thread_record(thread.id())?.core);
345            drop(state);
346            self.install_idle_core(cpu.as_mut(), core)
347        });
348        if let Err(error) = setup {
349            return match self.discard_unpublished_thread(thread) {
350                Ok(()) => Err(error),
351                Err(cleanup_error) => Err(cleanup_error),
352            };
353        }
354        Ok(thread)
355    }
356
357    /// Installs the dedicated idle task directly into its owner rq, matching
358    /// Linux `init_idle()` rather than passing idle through a scheduling-class
359    /// enqueue/dequeue cycle.
360    pub(super) fn install_idle_core(
361        &self,
362        mut cpu: Pin<&mut CpuLocal>,
363        core: Arc<ThreadCore>,
364    ) -> Result<(), TaskError> {
365        let owner = cpu.owner();
366        // SAFETY: idle installation precedes CPU online publication and the
367        // boot owner retains local IRQ exclusion.
368        if unsafe { cpu.remote().lock_run_queue_irq_disabled() }
369            .idle()
370            .is_some()
371        {
372            return Err(TaskError::InvalidConfiguration);
373        }
374        // SAFETY: install_idle_core is reached only from the offline bootstrap
375        // transaction above.
376        let mut sched = unsafe { core.sched().lock_bootstrap() };
377        let policy = core.sched().active(&sched).policy();
378        if sched.lifecycle.state() != ThreadState::Running
379            || !matches!(
380                policy,
381                SchedulePolicy::Fair {
382                    mode: crate::sched::FairMode::Idle,
383                    ..
384                }
385            )
386            || !sched.affinity.affinity.contains(owner)
387            || sched.placement.assigned_cpu() != Some(owner)
388            || sched.placement.on_cpu().is_some()
389            || sched.placement.requested_migration().is_some()
390        {
391            return Err(TaskError::InvalidConfiguration);
392        }
393        let metadata = sched.rq_task_metadata()?;
394        let rt_quota_exempt = sched.is_pi_boosted_rt_owner_for(policy);
395        let active = core.sched().take_active(&mut sched);
396        // SAFETY: the CPU remains offline and boot-owned through this direct
397        // init_idle-style rq transaction.
398        unsafe {
399            cpu.as_mut().install_idle_bootstrap(
400                self,
401                core.id(),
402                Arc::clone(&core),
403                active,
404                metadata,
405                rt_quota_exempt,
406            )
407        };
408        core.set_wake_cpu_hint(owner);
409        Ok(())
410    }
411
412    fn discard_unpublished_thread(&self, handle: ThreadHandle) -> Result<(), TaskError> {
413        let record = {
414            let mut state = self.state.lock();
415            let mut root_domain = self.root_domain.lock();
416            let (record, released) = state.remove_unpublished_thread_with_handle(&handle)?;
417            root_domain.release_deadline(released);
418            record
419        };
420        drop(handle);
421        self.release_thread_record(record);
422        Ok(())
423    }
424}
ax_task/sched/system/task_system/thread_creation.rs

ax_task/sched/system/task_system/
thread_creation.rs
