mill_io/

reactor.rs

use crate::{
    error::Result,
    poll::PollHandle,
    thread_pool::{ComputePoolMetrics, ComputeThreadPool, TaskPriority, ThreadPool},
};
use mio::{event::Event, Events};
use std::{
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, RwLock,
    },
    time::Duration,
};

pub const DEFAULT_EVENTS_CAPACITY: usize = 1024;
pub const DEFAULT_POLL_TIMEOUT_MS: u64 = 150;

pub struct ReactorOptions {
    pub direct_dispatch: bool,
}

pub struct Reactor {
    pub(crate) poll_handle: PollHandle,
    events: Arc<RwLock<Events>>,
    pool: ThreadPool,
    compute_pool: ComputeThreadPool,
    running: AtomicBool,
    poll_timeout_ms: u64,
    options: ReactorOptions,
}

impl Default for Reactor {
    fn default() -> Self {
        Self {
            poll_handle: PollHandle::new().unwrap(),
            events: Arc::new(RwLock::new(Events::with_capacity(DEFAULT_EVENTS_CAPACITY))),
            pool: ThreadPool::default(),
            compute_pool: ComputeThreadPool::new(
                std::thread::available_parallelism()
                    .map(|n| n.get())
                    .unwrap_or(4),
            ),
            running: AtomicBool::new(false),
            poll_timeout_ms: DEFAULT_POLL_TIMEOUT_MS,
            options: ReactorOptions {
                direct_dispatch: false,
            },
        }
    }
}

impl Reactor {
    pub fn new(pool_size: usize, events_capacity: usize, poll_timeout_ms: u64) -> Result<Self> {
        Ok(Self {
            poll_handle: PollHandle::new()?,
            events: Arc::new(RwLock::new(Events::with_capacity(events_capacity))),
            pool: ThreadPool::new(pool_size),
            compute_pool: ComputeThreadPool::default(),
            running: AtomicBool::new(false),
            poll_timeout_ms,
            options: ReactorOptions {
                direct_dispatch: false,
            },
        })
    }

    pub fn new_with_options(
        pool_size: usize,
        events_capacity: usize,
        poll_timeout_ms: u64,
        options: ReactorOptions,
    ) -> Result<Self> {
        Ok(Self {
            poll_handle: PollHandle::new()?,
            events: Arc::new(RwLock::new(Events::with_capacity(events_capacity))),
            pool: ThreadPool::new(pool_size),
            compute_pool: ComputeThreadPool::default(),
            running: AtomicBool::new(false),
            poll_timeout_ms,
            options,
        })
    }

    pub fn run(&self) -> Result<()> {
        self.running.store(true, Ordering::SeqCst);

        while self.running.load(Ordering::SeqCst) {
            let _ = self.poll_handle.poll(
                &mut self.events.write().unwrap(),
                Some(Duration::from_millis(self.poll_timeout_ms)),
            )?;

            for event in self.events.read().unwrap().iter() {
                self.dispatch_event(event.clone())?;
            }
        }
        Ok(())
    }

    pub fn get_shutdown_handle(&self) -> ShutdownHandle<'_> {
        ShutdownHandle {
            running: &self.running,
            poll_handle: &self.poll_handle,
        }
    }

    pub fn dispatch_event(&self, event: Event) -> Result<()> {
        let token = event.token();
        let is_readable = event.is_readable();
        let is_writable = event.is_writable();

        let registry = self.poll_handle.get_registery();

        if self.options.direct_dispatch {
            // FAST PATH: Direct execution on reactor thread
            if let Some(entry) = registry.get(&token) {
                let interest = entry.1.interest;
                let handler = entry.1.handler.as_ref();
                if (interest.is_readable() && is_readable)
                    || (interest.is_writable() && is_writable)
                {
                    handler.handle_event(&event);
                }
            }
            Ok(())
        } else {
            // SLOW PATH: Dispatch to thread pool
            self.pool.exec(move || {
                let entry = registry.get(&token);
                if let Some(entry) = entry {
                    let interest = entry.1.interest;
                    let handler = entry.1.handler.as_ref();
                    if (interest.is_readable() && is_readable)
                        || (interest.is_writable() && is_writable)
                    {
                        handler.handle_event(&event);
                    }
                }
            })
        }
    }

    pub fn spawn_compute<F>(&self, task: F, priority: TaskPriority)
    where
        F: FnOnce() + Send + 'static,
    {
        self.compute_pool.spawn(task, priority);
    }

    pub fn get_compute_metrics(&self) -> Arc<ComputePoolMetrics> {
        self.compute_pool.metrics()
    }

    pub fn get_events(&self) -> Arc<RwLock<Events>> {
        self.events.clone()
    }
}

pub struct ShutdownHandle<'a> {
    running: &'a AtomicBool,
    poll_handle: &'a PollHandle,
}

impl ShutdownHandle<'_> {
    pub fn shutdown(&self) {
        self.running.store(false, Ordering::SeqCst);
        self.poll_handle.wake().unwrap();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::handler::*;
    use mio::{Interest, Token};
    use std::sync::{Arc, Condvar, Mutex};
    use std::time::Duration;

    #[derive(Clone)]
    struct TestHandler {
        counter: Arc<Mutex<usize>>,
        condition: Arc<Condvar>,
    }

    impl EventHandler for TestHandler {
        fn handle_event(&self, _event: &Event) {
            let mut count = self.counter.lock().unwrap();
            *count += 1;
            self.condition.notify_one();
        }
    }

    #[cfg(unix)]
    #[test]
    fn test_reactor_start_stop() {
        let reactor = Arc::new(Reactor::default());
        let shutdown_handle = reactor.get_shutdown_handle();

        let reactor_clone = Arc::clone(&reactor);
        let handle = std::thread::spawn(move || {
            reactor_clone.run().unwrap();
        });

        std::thread::sleep(Duration::from_millis(100));

        shutdown_handle.shutdown();

        handle.join().unwrap();
    }

    #[cfg(unix)]
    #[test]
    fn test_with_pipe() -> std::io::Result<()> {
        use mio::net::UnixStream;

        let reactor =
            Arc::new(Reactor::new(2, DEFAULT_EVENTS_CAPACITY, DEFAULT_POLL_TIMEOUT_MS).unwrap());
        let counter = Arc::new(Mutex::new(0));
        let condition = Arc::new(Condvar::new());

        let (mut stream1, mut stream2) = UnixStream::pair()?;

        let handler = TestHandler {
            counter: Arc::clone(&counter),
            condition: Arc::clone(&condition),
        };

        let token = Token(1);

        reactor
            .poll_handle
            .register(&mut stream1, token, Interest::READABLE, handler)
            .unwrap();

        let reactor_clone = Arc::clone(&reactor);
        let handle = std::thread::spawn(move || {
            // Poll once
            let events_result = {
                let mut events = reactor_clone.events.write().unwrap();
                reactor_clone
                    .poll_handle
                    .poll(&mut events, Some(Duration::from_millis(100)))
            };

            if events_result.is_ok() {
                let events = reactor_clone.events.read().unwrap();
                for event in events.iter() {
                    let _ = reactor_clone.dispatch_event(event.clone());
                }
            }
        });

        std::io::Write::write_all(&mut stream2, b"test data")?;

        handle.join().unwrap();

        let count = counter.lock().unwrap();
        let result = condition
            .wait_timeout(count, Duration::from_millis(500))
            .unwrap();

        if !result.1.timed_out() {
            assert_eq!(*result.0, 1);
        }

        Ok(())
    }

    #[test]
    fn test_with_tcp() -> std::io::Result<()> {
        use mio::net::{TcpListener, TcpStream};
        use std::net::SocketAddr;

        let reactor =
            Arc::new(Reactor::new(2, DEFAULT_EVENTS_CAPACITY, DEFAULT_POLL_TIMEOUT_MS).unwrap());
        let counter = Arc::new(Mutex::new(0));
        let condition = Arc::new(Condvar::new());

        // Create a TCP listener on localhost
        let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
        let mut listener = TcpListener::bind(addr)?;
        let listener_addr = listener.local_addr()?;

        let handler = TestHandler {
            counter: Arc::clone(&counter),
            condition: Arc::clone(&condition),
        };

        let token = Token(1);

        reactor
            .poll_handle
            .register(&mut listener, token, Interest::READABLE, handler)
            .unwrap();

        let reactor_clone = Arc::clone(&reactor);
        let handle = std::thread::spawn(move || {
            // Poll once
            let events_result = {
                let mut events = reactor_clone.events.write().unwrap();
                reactor_clone
                    .poll_handle
                    .poll(&mut events, Some(Duration::from_millis(100)))
            };

            if events_result.is_ok() {
                let events = reactor_clone.events.read().unwrap();
                for event in events.iter() {
                    let _ = reactor_clone.dispatch_event(event.clone());
                }
            }
        });

        // Connect to trigger the event
        let _stream = TcpStream::connect(listener_addr)?;

        handle.join().unwrap();

        let count = counter.lock().unwrap();
        let result = condition
            .wait_timeout(count, Duration::from_millis(500))
            .unwrap();

        if !result.1.timed_out() {
            assert_eq!(*result.0, 1);
        }

        Ok(())
    }
}