Rust の組込み用非同期フレームワーク Embassy (7)

自分のためだけに Embassy コードリーディングシリーズです。 nRF52840 向けの Waker を wake している実装を見ていきます。

GPIO 割り込みから wake() するのが、おそらく最も単純な実装でしょう、ということで GPIO の example を探します。 examples/nrf52840/src/bin/gpiote_port.rs

中身は下のような感じです。 Task Pool が 4 になっていて、同じタスクを複数 spawn できるようになっているようですが、そこは一旦無視しましょう。 pin.wait_for_low() を見るのが良さそうです。

#[embassy_executor::task(pool_size = 4)]
async fn button_task(n: usize, mut pin: Input<'static, AnyPin>) {
    loop {
        pin.wait_for_low().await;
        info!("Button {:?} pressed!", n);
        pin.wait_for_high().await;
        info!("Button {:?} released!", n);
    }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    let p = embassy_nrf::init(Default::default());
    info!("Starting!");

    let btn1 = Input::new(p.P0_11.degrade(), Pull::Up);
    let btn2 = Input::new(p.P0_12.degrade(), Pull::Up);
    let btn3 = Input::new(p.P0_24.degrade(), Pull::Up);
    let btn4 = Input::new(p.P0_25.degrade(), Pull::Up);

    unwrap!(spawner.spawn(button_task(1, btn1)));
    unwrap!(spawner.spawn(button_task(2, btn2)));
    unwrap!(spawner.spawn(button_task(3, btn3)));
    unwrap!(spawner.spawn(button_task(4, btn4)));
}

Input は Flex のラッパーになっています。 wait_for_low() を呼び出すと、Flex の wait_for_low() を呼びます。

/// GPIO input driver.
pub struct Input<'d, T: Pin> {
    pub(crate) pin: Flex<'d, T>,
}

impl<'d, T: GpioPin> Input<'d, T> {
    /// Wait until the pin is low. If it is already low, return immediately.
    pub async fn wait_for_low(&mut self) {
        self.pin.wait_for_low().await
    }

ちなみに Flex は Flexible pin のことで、組込み Rust では GPIO ピンの状態を型として表現することが多いのですが、どのような状態も取れるピンとして定義されています。

Flex::wait_for_low() では、GPIO 入力が low になったら割り込みが入るように設定して、PortInputFuture のインスタンスを作って await しています。

impl<'d, T: GpioPin> Flex<'d, T> {
    /// Wait until the pin is low. If it is already low, return immediately.
    pub async fn wait_for_low(&mut self) {
        self.pin.conf().modify(|_, w| w.sense().low());
        PortInputFuture::new(&mut self.pin).await
    }

重要なのは Future トレイトを実装している PortInputFuture ですね。 poll の実装は次のようになっています。 PORT ごとに Waker を持っていて、register() で登録するようです。 最後は、GPIO 割り込みが検出されて SENSE が無効になっていたら、Poll::Ready(()) を返しています。

impl<'a> Future for PortInputFuture<'a> {
    type Output = ();

    fn poll(self: core::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        PORT_WAKERS[self.pin.pin_port() as usize].register(cx.waker());

        if self.pin.conf().read().sense().is_disabled() {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

register() の部分ですが、PORT_WAKERS は AtomicWaker の static な配列になっており、AtomicWaker がこうなので、割り込み待ちしている GPIO ピンの Waker を wake したら、このときに登録したコンテキストが run queue に入る、ということになります。

/// Utility struct to register and wake a waker.
pub struct AtomicWaker {
    waker: Mutex<CriticalSectionRawMutex, Cell<Option<Waker>>>,
}

impl AtomicWaker {
    /// Register a waker. Overwrites the previous waker, if any.
    pub fn register(&self, w: &Waker) {
        critical_section::with(|cs| {
            let cell = self.waker.borrow(cs);
            cell.set(match cell.replace(None) {
                Some(w2) if (w2.will_wake(w)) => Some(w2),
                _ => Some(w.clone()),
            })
        })
    }

    /// Wake the registered waker, if any.
    pub fn wake(&self) {
        critical_section::with(|cs| {
            let cell = self.waker.borrow(cs);
            if let Some(w) = cell.replace(None) {
                w.wake_by_ref();
                cell.set(Some(w));
            }
        })
    }
}

では続いて、GPIO の割り込みハンドラ embassy-nrf/src/gpiote.rs です。 重要なところはコード中にコメントで補足した、wake() しているところです。 この GPIO 割り込みの中で、wait_for_low() している実行コンテキストを run queue に入れて、WFE から復帰すると、無事 wait_for_low() の続きからプログラムが実行されることになります。

// 抜粋
unsafe fn handle_gpiote_interrupt() {
    let g: &pac::gpiote::RegisterBlock = regs();

    if g.events_port.read().bits() != 0 {
        g.events_port.write(|w| w);
        let ports = &[&*pac::P0::ptr(), &*pac::P1::ptr()];

        for (port, &p) in ports.iter().enumerate() {
            let bits = p.latch.read().bits();
            for pin in BitIter(bits) {
                // ここで Waker を wake() している
                p.pin_cnf[pin as usize].modify(|_, w| w.sense().disabled());
                PORT_WAKERS[port * 32 + pin as usize].wake();
            }
            p.latch.write(|w| w.bits(bits));
        }
    }
}

ということで、 embassy 自体には他にも機能がありそうですが、最もベースとなる部分の仕組みは大体わかった気になれましたね。