retty-io 0.0.0

Metal IO (mio) based Utilities, like Channel, Timer, UDP Socket with ECN, etc
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::{
    io,
    sync::{Arc, Mutex},
};

use mio::{event, Token, Waker};

/// Create a pair of the [`Sender`] and the [`Receiver`].
///
/// The [`Receiver`] implements the [`event::Source`] so that it can be registered
/// with the [`mio::poll::Poll`], while the [`Sender`] doesn't.
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
    let (tx, rx) = crossbeam::channel::unbounded();

    let waker = Arc::new(Mutex::new(HashMap::new()));

    (
        Sender {
            waker: waker.clone(),
            tx,
        },
        Receiver {
            waker,
            rx,
            id: 0,
            next_id: Arc::new(AtomicUsize::new(1)),
        },
    )
}

/// A wrapper of the [`crossbeam::channel::Receiver`].
///
/// It implements the [`event::Source`] so that it can be registered with the [`mio::poll::Poll`].
/// It ignores the [`mio::Interest`] and always cause readable events.
pub struct Receiver<T> {
    waker: Arc<Mutex<HashMap<usize, Option<Waker>>>>,
    rx: crossbeam::channel::Receiver<T>,
    id: usize,
    next_id: Arc<AtomicUsize>,
}

impl<T> Clone for Receiver<T> {
    fn clone(&self) -> Self {
        let next_id = self.next_id.clone();
        let id = next_id.fetch_add(1, Ordering::Relaxed);
        Self {
            waker: self.waker.clone(),
            rx: self.rx.clone(),
            id,
            next_id,
        }
    }
}

impl<T> Receiver<T> {
    /// Try to receive a value. It works just like [`crossbeam::channel::Receiver::try_recv`].
    pub fn try_recv(&self) -> Result<T, crossbeam::channel::TryRecvError> {
        self.rx.try_recv()
    }
}

impl<T> event::Source for Receiver<T> {
    fn register(
        &mut self,
        registry: &mio::Registry,
        token: Token,
        _: mio::Interest,
    ) -> io::Result<()> {
        let mut waker_map = self.waker.lock().unwrap();
        if let Some(waker) = waker_map.get_mut(&self.id) {
            if waker.is_none() {
                *waker = Some(Waker::new(registry, token)?);
            }
        } else {
            waker_map.insert(self.id, Some(Waker::new(registry, token)?));
        }

        Ok(())
    }

    fn reregister(
        &mut self,
        registry: &mio::Registry,
        token: Token,
        _: mio::Interest,
    ) -> io::Result<()> {
        let mut waker_map = self.waker.lock().unwrap();
        if let Some(waker) = waker_map.get_mut(&self.id) {
            *waker = Some(Waker::new(registry, token)?);
        } else {
            waker_map.insert(self.id, Some(Waker::new(registry, token)?));
        }

        Ok(())
    }

    fn deregister(&mut self, _: &mio::Registry) -> io::Result<()> {
        let mut waker_map = self.waker.lock().unwrap();
        if let Some(waker) = waker_map.get_mut(&self.id) {
            *waker = None;
        }

        Ok(())
    }
}

/// A wrapper of the [`crossbeam::channel::Sender`].
pub struct Sender<T> {
    waker: Arc<Mutex<HashMap<usize, Option<Waker>>>>,
    tx: crossbeam::channel::Sender<T>,
}

impl<T> Sender<T> {
    /// Try to send a value. It works just like [`crossbeam::channel::Sender::send`].
    /// After sending it, it's waking up the [`mio::poll::Poll`].
    ///
    /// Note that it does not return any I/O error even if it occurs
    /// when waking up the [`mio::poll::Poll`].
    pub fn send(&self, t: T) -> Result<(), crossbeam::channel::SendError<T>> {
        self.tx.send(t)?;

        let mut waker_map = self.waker.lock().unwrap();
        for waker in waker_map.values_mut().flatten() {
            let _ = waker.wake();
        }

        Ok(())
    }
}