ArceOS Tutorial Book

第三节 触发抢占的外部条件与触发点

前面提到，在内核运行的某些阶段中，需要暂时禁用抢占，而且这种禁用有可能叠加。

为了处理叠加的情况，我们为任务增加一个计数值记录抢占禁用的情况。

pub struct Task {
	... ...
    preempt_disable_count: AtomicUsize,
    ... ...
}

impl Task {
    fn new_common(id: TaskId, name: String) -> Self {
        Self {
            ... ...
            preempt_disable_count: AtomicUsize::new(0),
            ... ...
        }
    }

    #[inline]
    pub(crate) fn can_preempt(&self, current_disable_count: usize) -> bool {
        self.preempt_disable_count.load(Ordering::Acquire) == current_disable_count
    }

    #[inline]
    pub(crate) fn disable_preempt(&self) {
        self.preempt_disable_count.fetch_add(1, Ordering::Relaxed);
    }

    #[inline]
    pub(crate) fn enable_preempt(&self, resched: bool) {
        if self.preempt_disable_count.fetch_sub(1, Ordering::Relaxed) == 1 && resched {
            // If current task is pending to be preempted, do rescheduling.
            Self::current_check_preempt_pending();
        }
    }
}

关注一下 enable_preempt/disable_preempt，禁用抢占时递增，启用时递减。典型情况下，计数值是 0 时可触发抢占。

内核中直接调用 enable_preempt/disable_preempt 很容易出错，所以我们来封装一个 Guard 类型 NoPreempt 来实现对抢占的控制。

// spinlock/src/lib.rs
pub struct NoPreempt;

impl NoPreempt {
    fn acquire() {
        crate_interface::call_interface!(KernelGuardIf::disable_preempt);
    }
    fn release() {
        crate_interface::call_interface!(KernelGuardIf::enable_preempt);
    }

    pub fn new() -> Self {
        Self::acquire();
        Self
    }
}

impl Drop for NoPreempt {
    fn drop(&mut self) {
        Self::release();
    }
}

// axtask/src/lib.rs
struct KernelGuardIfImpl;

#[crate_interface::impl_interface]
impl kernel_guard::KernelGuardIf for KernelGuardIfImpl {
    fn disable_preempt() {
        if let Some(curr) = current_may_uninit() {
            curr.disable_preempt();
        }
    }

    fn enable_preempt() {
        if let Some(curr) = current_may_uninit() {
            curr.enable_preempt(true);
        }
    }
}

pub fn current_may_uninit() -> Option<CurrentTask> {
    CurrentTask::try_get()
}

最后，重构一下自旋锁的实现，确保加锁时抢占是禁用状态，防止此期间任务被调度出去破坏原子性。

pub struct SpinNoIrqGuard<T> {
    irq_state: usize,
    data: *mut T,
}

impl<T> SpinNoIrq<T> {
    #[inline(always)]
    pub fn lock(&self) -> SpinNoIrqGuard<T> {
        let irq_state = NoPreemptIrqSave::acquire();
        SpinNoIrqGuard {
            irq_state,
            data: unsafe { &mut *self.data.get() },
        }
    }
}
impl<T> Drop for SpinNoIrqGuard<T> {
    #[inline(always)]
    fn drop(&mut self) {
        NoPreemptIrqSave::release(self.irq_state);
    }
}

加锁时，调用 NoPreemptIrqSave::acquire 关抢占；解锁时，调用 NoPreemptIrqSave::release 开抢占。

这样就把抢占的控制与当前任务关联起来，下面来看一个示例：

XXX