virtual_mio/
guard.rs

1use std::{
2    io,
3    sync::{Arc, Weak},
4};
5
6use mio::Token;
7
8use crate::{InterestHandler, InterestType, InterestWakerMap, Selector};
9
10#[derive(Debug)]
11#[must_use = "Leaking token guards will break the IO subsystem"]
12pub struct InterestGuard {
13    selector: Weak<Selector>,
14    pub(crate) token: Token,
15}
16impl Drop for InterestGuard {
17    fn drop(&mut self) {
18        self.drop_internal();
19    }
20}
21impl InterestGuard {
22    pub fn new(
23        selector: &Arc<Selector>,
24        handler: Box<dyn InterestHandler + Send + Sync>,
25        source: &mut dyn mio::event::Source,
26        interest: mio::Interest,
27    ) -> io::Result<InterestGuard> {
28        let token = selector.add(handler, source, interest)?;
29        Ok(Self {
30            selector: Arc::downgrade(selector),
31            token,
32        })
33    }
34
35    pub fn unregister(&mut self, source: &mut dyn mio::event::Source) -> io::Result<()> {
36        if let Some(selector) = self.selector.upgrade() {
37            selector.remove(self.token, Some(source))?;
38        }
39        Ok(())
40    }
41
42    pub fn replace_handler(
43        &mut self,
44        handler: Box<dyn InterestHandler + Send + Sync>,
45    ) -> Result<(), Box<dyn InterestHandler + Send + Sync>> {
46        match self.selector.upgrade() {
47            Some(selector) => {
48                selector.replace(self.token, handler);
49                Ok(())
50            }
51            _ => Err(handler),
52        }
53    }
54
55    pub fn interest(&mut self, interest: InterestType) {
56        if let Some(selector) = self.selector.upgrade() {
57            selector.push_interest(self.token, interest);
58        }
59    }
60
61    fn drop_internal(&mut self) {
62        if let Some(selector) = self.selector.upgrade() {
63            selector.remove(self.token, None).ok();
64        }
65    }
66}
67
68#[derive(Debug)]
69pub enum HandlerGuardState {
70    None,
71    ExternalHandler(InterestGuard),
72    WakerMap(InterestGuard, InterestWakerMap),
73}
74
75pub fn state_as_waker_map<'a>(
76    state: &'a mut HandlerGuardState,
77    selector: &'_ Arc<Selector>,
78    source: &'_ mut dyn mio::event::Source,
79) -> io::Result<&'a mut InterestWakerMap> {
80    if !matches!(state, HandlerGuardState::WakerMap(_, _)) {
81        let waker_map = InterestWakerMap::default();
82        *state = HandlerGuardState::WakerMap(
83            InterestGuard::new(
84                selector,
85                Box::new(waker_map.clone()),
86                source,
87                mio::Interest::READABLE | mio::Interest::WRITABLE,
88            )?,
89            waker_map,
90        );
91    }
92    Ok(match state {
93        HandlerGuardState::WakerMap(_, map) => map,
94        _ => unreachable!(),
95    })
96}
virtual_mio/guard.rs

virtual_mio/
guard.rs
