Skip to main content

ax_task/sync/
migration.rs

1//! Migration exclusion for preemptible task-context critical sections.
2
3use alloc::sync::Arc;
4use core::marker::PhantomData;
5
6use crate::{
7    runtime::{context::runtime_task_system, lock::PreemptScope, task_runtime},
8    sched::CpuId,
9    thread::{TaskError, ThreadCore, current::current_thread_core_arc},
10};
11
12/// Pins the current task to its CPU without disabling preemption or interrupts.
13///
14/// Pins nest. Remote affinity requests remain pending until the outermost pin
15/// is released. The guard cannot be transferred to another task.
16#[must_use]
17pub struct MigrationGuard {
18    current: Arc<ThreadCore>,
19    cpu: CpuId,
20    _not_send: PhantomData<*mut ()>,
21}
22
23impl MigrationGuard {
24    /// Acquires a task migration pin; hard IRQ contexts cannot own one.
25    pub fn new() -> Result<Self, TaskError> {
26        if task_runtime::in_hard_irq() {
27            return Err(TaskError::UnsafeContext);
28        }
29        let _preempt = PreemptScope::enter();
30        let current = current_thread_core_arc()?;
31        let cpu = runtime_task_system()?.disable_current_migration(&current)?;
32        Ok(Self {
33            current,
34            cpu,
35            _not_send: PhantomData,
36        })
37    }
38
39    /// Returns the CPU retained by this guard, including across preemption.
40    pub const fn cpu(&self) -> CpuId {
41        self.cpu
42    }
43}
44
45impl Drop for MigrationGuard {
46    fn drop(&mut self) {
47        let _preempt = PreemptScope::enter();
48        runtime_task_system()
49            .and_then(|system| system.enable_current_migration(&self.current, self.cpu))
50            .expect("migration guard must release on its owning task and CPU");
51    }
52}