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
use std::{
collections::HashMap,
hash::Hash,
io,
os::fd::{AsRawFd, RawFd},
};
use crate::cutils::cerr;
use libc::{c_short, pollfd, POLLIN, POLLOUT};
/// A set of indexed file descriptors to be polled using the [`poll`](https://manpage.me/?q=poll) system call.
pub struct PollSet<K> {
fds: HashMap<K, (RawFd, c_short)>,
}
impl<K: Eq + PartialEq + Hash + Clone> Default for PollSet<K> {
fn default() -> Self {
Self::new()
}
}
impl<K: Eq + PartialEq + Hash + Clone> PollSet<K> {
/// Create an empty set of file descriptors.
pub fn new() -> Self {
Self {
fds: HashMap::new(),
}
}
/// Add a file descriptor under the provided key. This descriptor will be checked for read events and return a unique identifier
/// for the descriptor inside the set.
///
/// If the provided key is already in the set, calling this function will overwrite the file
/// descriptor for that key.
pub fn add_fd_read<F: AsRawFd>(&mut self, key: K, fd: &F) {
self.add_fd(key, fd, POLLIN)
}
/// Add a file descriptor under the provided key. This descriptor will be checked for write events and return a unique identifier
/// for the descriptor inside the set.
///
/// If the provided key is already in the set, calling this function will overwrite the file
/// descriptor for that key.
pub fn add_fd_write<F: AsRawFd>(&mut self, key: K, fd: &F) {
self.add_fd(key, fd, POLLOUT)
}
fn add_fd<F: AsRawFd>(&mut self, key: K, fd: &F, events: c_short) {
self.fds.insert(key, (fd.as_raw_fd(), events));
}
/// Poll the set of file descriptors and return the key of the descriptors that are ready to be
/// read or written.
///
/// Calling this function will block until one of the file descriptors in the set is ready.
pub fn poll(&mut self) -> io::Result<Vec<K>> {
let mut fds: Vec<pollfd> = self
.fds
.values()
.map(|&(fd, events)| pollfd {
fd,
events,
revents: 0,
})
.collect();
// FIXME: we should set either a timeout or use ppoll when available.
let n = cerr(unsafe { libc::poll(fds.as_mut_ptr(), fds.len() as _, -1) })?;
let mut keys = Vec::with_capacity(n as usize);
for (key, fd) in self.fds.keys().zip(fds) {
let events = fd.events & fd.revents;
if (events & POLLIN != 0) || (events & POLLOUT != 0) {
keys.push(key.clone());
}
}
Ok(keys)
}
}