Skip to main content

ax_task/sched/system/task_system/scheduling/
selection.rs

1//! Selection under the owning scheduler transaction.
2
3use super::*;
4
5impl TaskSystem {
6    /// Selects the next thread according to strict class precedence.
7    ///
8    /// `current` is the architecture-published task identity used only to
9    /// acquire task-owned scheduler state before the runqueue transaction.
10    /// `None` is valid only for an initial dispatch with no `rq->curr`.
11    pub fn schedule(
12        &self,
13        cpu: Pin<&mut CpuLocal>,
14        current: Option<&ThreadHandle>,
15    ) -> Result<ScheduleDecision, TaskError> {
16        self.schedule_owner(
17            cpu,
18            current.map(|thread| thread.runtime_core_arc().as_ref()),
19            OwnerRqEntry::IrqSave,
20        )
21    }
22
23    pub(super) fn schedule_owner(
24        &self,
25        mut cpu: Pin<&mut CpuLocal>,
26        current: Option<&ThreadCore>,
27        rq_entry: OwnerRqEntry,
28    ) -> Result<ScheduleDecision, TaskError> {
29        let validate_owner = rq_entry.requires_owner_context_validation();
30        if validate_owner {
31            self.ensure_owner_cpu_context(&cpu)?;
32        }
33        // SAFETY: the owner borrow pins the CpuLocal and its immutable remote
34        // endpoint while this scheduling transaction and switch tail are live.
35        let remote = unsafe { cpu.as_ref().get_ref().remote_for_owner() };
36        let initial_request = remote.claim_scheduler_request(SchedulerRequestScope::All);
37        self.drain_owner_work(cpu.as_mut())?;
38        if validate_owner {
39            self.ensure_owner_cpu_registration_online(&cpu)?;
40        }
41        let previous_core_hint = current;
42        let mut previous_sched = previous_core_hint.map(|core| {
43            // SAFETY: propagated from the selected entry contract.
44            unsafe { rq_entry.lock_thread_sched(core.sched()) }
45        });
46        // SAFETY: the public task entry chooses irqsave; the scheduler-frame
47        // entry is exposed only by its unsafe wrapper below.
48        let mut transaction = unsafe { rq_entry.begin(self, remote) };
49        let now_ns = transaction.clock().wall().as_nanos();
50        transaction.adopt_scheduler_request(initial_request);
51        transaction.merge_scheduler_request(SchedulerRequestScope::All);
52        let dispatch_commit = self.settle_owner_current_dispatch_in_rq(&mut transaction);
53        // Runtime accounting is part of this unconditional scheduling
54        // decision, exactly like Linux update_curr() preceding pick_next.
55        transaction.merge_scheduler_request(SchedulerRequestScope::All);
56        let previous = transaction.current_thread();
57        let previous_core = transaction.current_core();
58        let previous_endpoint = transaction.current_switch_endpoint();
59        let previous_urgency = transaction.current_scheduling_urgency();
60        if previous_core.as_deref().map(core::ptr::from_ref)
61            != previous_core_hint.map(core::ptr::from_ref)
62        {
63            task_runtime::fatal_invariant(0x5343_1201, cpu.owner().as_u32() as usize);
64        }
65        let mut migration = None;
66        if let Some(core) = previous_core.as_ref() {
67            let schedule_out = self.schedule_out_owner_running_in_rq(
68                cpu.as_mut(),
69                &mut transaction,
70                Arc::clone(core),
71                previous_sched.as_deref_mut().unwrap_or_else(|| {
72                    task_runtime::fatal_invariant(0x5343_1202, core.id().as_u64() as usize)
73                }),
74                now_ns,
75                EnqueueReason::Preempted,
76            );
77            migration = schedule_out.migration;
78        }
79        let next =
80            self.pick_owner_next_after_preemption_in_rq(cpu.as_mut(), &mut transaction, previous);
81        let OwnerNext {
82            core: next_core,
83            policy: next_policy_ref,
84            urgency: next_urgency,
85        } = next;
86        let next_endpoint = transaction.current_switch_endpoint().unwrap_or_else(|| {
87            task_runtime::fatal_invariant(0x5343_1203, next_core.as_ref().id().as_u64() as usize)
88        });
89        let migrated = migration.is_some();
90        let handoff = Self::prepare_switch_handoff(
91            previous,
92            previous_core.map(PreviousSwitchOwnership::retained),
93            next_core,
94            next_policy_ref,
95            PreviousSwitchDisposition::Live,
96            migration,
97        );
98        let reason = if migrated {
99            SwitchReason::Migrated
100        } else {
101            SwitchReason::Preempted
102        };
103        let deadline_rq_observation =
104            transaction.scheduler_deadline_rq_observation(cpu.as_ref().get_ref());
105        self.commit_owner_switch_selection(
106            cpu.as_mut(),
107            transaction,
108            handoff,
109            !migrated && !dispatch_commit.has_deferred_task_lock_work(),
110        );
111        drop(previous_sched);
112        let decision = Self::owner_switch_plan(previous_endpoint, next_endpoint, reason, now_ns);
113        self.finish_owner_dispatch_commit(dispatch_commit);
114        self.finish_owner_selection(
115            cpu.as_mut(),
116            decision.previous(),
117            decision.next(),
118            previous_urgency,
119            next_urgency,
120            OwnerSchedulerDeadline::Reevaluate(deadline_rq_observation),
121        );
122        Ok(decision)
123    }
124}