compio_net/socket/
linux.rs

1use std::sync::atomic::Ordering;
2
3use compio_driver::{AsSocketState, Extra};
4#[cfg(feature = "sync")]
5use synchrony::sync::atomic::AtomicU8;
6#[cfg(not(feature = "sync"))]
7use synchrony::unsync::atomic::AtomicU8;
8
9use super::SocketWithState;
10
11const RECV_OFFSET: usize = 0;
12const ACCEPT_OFFSET: usize = 2;
13
14const UNSET: u8 = 0;
15const EMPTY: u8 = 1;
16const NON_EMPTY: u8 = 2;
17
18#[derive(Debug)]
19pub(super) struct SocketState {
20    state: AtomicU8,
21}
22
23impl Default for SocketState {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl SocketState {
30    pub(super) fn new() -> Self {
31        Self {
32            state: AtomicU8::new(0),
33        }
34    }
35
36    fn get_bit(&self, offset: usize) -> Option<bool> {
37        let state = self.state.load(Ordering::Relaxed);
38        match (state >> offset) & 0b11 {
39            UNSET => None,
40            EMPTY => Some(false),
41            NON_EMPTY => Some(true),
42            _ => unreachable!(),
43        }
44    }
45
46    fn set_bit(&self, offset: usize, value: bool) {
47        let bits = if value { NON_EMPTY } else { EMPTY } << offset;
48        self.state
49            .update(Ordering::Relaxed, Ordering::Relaxed, |state| {
50                (state & !(0b11 << offset)) | bits
51            });
52    }
53
54    pub(super) fn set_recv(&self, extra: &Extra) {
55        if let Ok(n) = extra.sock_nonempty() {
56            self.set_bit(RECV_OFFSET, n);
57        }
58    }
59
60    /// Whether the socket is known to be empty, so that a receive should arm
61    /// poll before its first attempt.
62    pub(super) fn recv_poll_first(&self) -> bool {
63        self.get_bit(RECV_OFFSET) == Some(false)
64    }
65
66    pub(super) fn set_accept(&self, extra: &Extra) {
67        if let Ok(n) = extra.sock_nonempty() {
68            self.set_bit(ACCEPT_OFFSET, n);
69        }
70    }
71
72    /// Same as [`SocketState::recv_poll_first`], for accept.
73    pub(super) fn accept_poll_first(&self) -> bool {
74        self.get_bit(ACCEPT_OFFSET) == Some(false)
75    }
76}
77
78impl AsSocketState for SocketWithState {
79    fn recv_poll_first(&self) -> bool {
80        self.state.recv_poll_first()
81    }
82
83    fn accept_poll_first(&self) -> bool {
84        self.state.accept_poll_first()
85    }
86
87    fn set_recv(&self, extra: &Extra) {
88        self.state.set_recv(extra);
89    }
90
91    fn set_accept(&self, extra: &Extra) {
92        self.state.set_accept(extra);
93    }
94}
compio_net/socket/linux.rs

compio_net/socket/
linux.rs
