lustre-executor 0.2.0

//! Mio-based reactor for I/O event polling.

use mio::Poll as MioPoll;
use mio::{Events, Interest, Token};
use std::collections::HashMap;
use std::io;
use std::task::{Context, Poll as TaskPoll, Waker};
use std::time::Duration;

/// Reactor for handling I/O readiness events.
pub struct Reactor {
    poll: MioPoll,
    events: Events,
    wakers: HashMap<Token, Waker>,
    next_token: usize,
}

impl Reactor {
    /// Creates a new reactor.
    pub fn new() -> io::Result<Self> {
        Ok(Self {
            poll: MioPoll::new()?,
            events: Events::with_capacity(1024),
            wakers: HashMap::new(),
            next_token: 0,
        })
    }

    /// Registers an I/O source with the reactor.
    pub fn register(
        &mut self,
        source: &mut impl mio::event::Source,
        interest: Interest,
        waker: Waker,
    ) -> Token {
        let token = Token(self.next_token);
        self.next_token += 1;
        self.poll
            .registry()
            .register(source, token, interest)
            .unwrap();
        self.wakers.insert(token, waker);
        token
    }

    /// Polls for I/O events and wakes associated tasks.
    pub fn poll(&mut self, _cx: &mut Context, timeout: Option<Duration>) -> TaskPoll<()> {
        self.poll.poll(&mut self.events, timeout).unwrap();
        for event in &self.events {
            if let Some(waker) = self.wakers.get(&event.token()) {
                waker.wake_by_ref();
            }
        }
        if self.events.is_empty() {
            TaskPoll::Pending
        } else {
            TaskPoll::Ready(())
        }
    }
}