Skip to main content

net_io/
id.rs

1//! Connection, listener, and worker identities.
2
3use std::fmt;
4use std::num::NonZeroU64;
5
6use std::sync::atomic::{AtomicU64, Ordering};
7
8/// The identity of one connection, unique and never reused within a process.
9///
10/// Workers take blocks of 1,024 ids from one process-wide counter, so ids are
11/// ordered within a worker, not globally.
12///
13/// # Examples
14///
15/// ```
16/// let id = net_io::ConnectionId::new(42).unwrap();
17/// assert_eq!(id.get(), 42);
18/// assert!(net_io::ConnectionId::new(0).is_none());
19/// ```
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
21pub struct ConnectionId(NonZeroU64);
22
23impl ConnectionId {
24    /// The id as a number.
25    #[inline]
26    #[must_use]
27    pub const fn get(self) -> u64 {
28        self.0.get()
29    }
30
31    /// An id from a number, for tests and conversions; `None` for 0.
32    ///
33    /// An id built this way is not reserved: use it only to name a connection
34    /// net-io issued, or in tests.
35    #[must_use]
36    pub const fn new(id: u64) -> Option<ConnectionId> {
37        match NonZeroU64::new(id) {
38            Some(n) => Some(ConnectionId(n)),
39            None => None,
40        }
41    }
42}
43
44impl fmt::Display for ConnectionId {
45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
46        write!(f, "#{}", self.0)
47    }
48}
49
50/// The index of a listener within one server.
51///
52/// ```
53/// # fn main() -> Result<(), net_io::NetError> {
54/// let server = net_io::serve("127.0.0.1:0", |_conn| Ok(()))?;
55/// assert_eq!(server.listeners()[0].id.get(), 0);
56/// # server.handle().request_stop(); let _ = server.wait_for_stop(); Ok(()) }
57/// ```
58#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
59pub struct ListenerId(pub(crate) u16);
60
61impl ListenerId {
62    /// The index as a number.
63    #[must_use]
64    pub const fn get(self) -> u16 {
65        self.0
66    }
67}
68
69/// The index of a worker within one server.
70///
71/// ```
72/// # fn main() -> Result<(), net_io::NetError> {
73/// let server = net_io::serve("127.0.0.1:0", |_conn| Ok(()))?;
74/// assert!(server.handle().completer(net_io::WorkerId::first()).is_some());
75/// # server.handle().request_stop(); let _ = server.wait_for_stop(); Ok(()) }
76/// ```
77#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
78pub struct WorkerId(pub(crate) u16);
79
80impl WorkerId {
81    /// The index as a number.
82    #[must_use]
83    pub const fn get(self) -> u16 {
84        self.0
85    }
86
87    /// Worker 0.
88    #[must_use]
89    pub const fn first() -> WorkerId {
90        WorkerId(0)
91    }
92
93    /// Worker `index`, for tests of per-worker state (a server never hands out a
94    /// worker it does not have).
95    ///
96    /// ```
97    /// assert_eq!(net_io::WorkerId::for_testing(3).get(), 3);
98    /// ```
99    #[must_use]
100    pub const fn for_testing(index: u16) -> WorkerId {
101        WorkerId(index)
102    }
103}
104
105/// Ids handed out per block.
106pub(crate) const BLOCK: u64 = 1024;
107
108/// The process-wide counter of ids. Starts at 1 so id 0 is never issued; at one
109/// million ids per second it would take about 584,000 years to wrap.
110static NEXT: AtomicU64 = AtomicU64::new(1);
111
112/// A block of ids owned by one worker.
113#[derive(Debug)]
114pub(crate) struct IdBlock {
115    next: u64,
116    end: u64,
117}
118
119impl IdBlock {
120    /// An empty block; the first `issue` takes a block from the counter.
121    pub(crate) const fn new() -> IdBlock {
122        IdBlock { next: 0, end: 0 }
123    }
124
125    /// The next id; `take` returns the start of a fresh block when this one is used up.
126    pub(crate) fn issue_with(&mut self, take: impl FnOnce() -> u64) -> ConnectionId {
127        if self.next >= self.end {
128            let start = take();
129            self.next = start;
130            self.end = start.saturating_add(BLOCK);
131        }
132        let id = self.next;
133        self.next += 1;
134        to_id(id)
135    }
136
137    /// The next id from the process-wide counter.
138    pub(crate) fn issue(&mut self) -> ConnectionId {
139        self.issue_with(|| NEXT.fetch_add(BLOCK, Ordering::Relaxed))
140    }
141}
142
143fn to_id(n: u64) -> ConnectionId {
144    match NonZeroU64::new(n) {
145        Some(n) => ConnectionId(n),
146        // The counter starts at 1 and only grows, so this arm is not reached.
147        None => ConnectionId(NonZeroU64::MIN),
148    }
149}
150
151/// One id for a caller without a block (standalone connections, tests).
152pub(crate) fn issue_single() -> ConnectionId {
153    to_id(NEXT.fetch_add(1, Ordering::Relaxed))
154}
155
156#[cfg(all(test, not(net_io_loom)))]
157mod tests {
158    use super::*;
159    use std::collections::HashSet;
160
161    #[test]
162    fn test_id_zero_never_issued() {
163        let counter = AtomicU64::new(1);
164        let mut block = IdBlock::new();
165        let first = block.issue_with(|| counter.fetch_add(BLOCK, Ordering::Relaxed));
166        assert_eq!(first.get(), 1);
167    }
168
169    #[test]
170    fn test_id_block_takes_new_block_after_1024() {
171        let counter = AtomicU64::new(1);
172        let mut block = IdBlock::new();
173        for _ in 0..BLOCK {
174            block.issue_with(|| counter.fetch_add(BLOCK, Ordering::Relaxed));
175        }
176        let next = block.issue_with(|| counter.fetch_add(BLOCK, Ordering::Relaxed));
177        assert_eq!(next.get(), 1 + BLOCK);
178        assert_eq!(counter.load(Ordering::Relaxed), 1 + 2 * BLOCK);
179    }
180
181    #[test]
182    fn test_ids_unique_across_threads() {
183        let handles: Vec<_> = (0..8)
184            .map(|_| {
185                std::thread::spawn(|| {
186                    let mut block = IdBlock::new();
187                    (0..5000).map(|_| block.issue().get()).collect::<Vec<_>>()
188                })
189            })
190            .collect();
191        let mut all = HashSet::new();
192        for h in handles {
193            for id in h.join().unwrap() {
194                assert!(all.insert(id), "id {id} issued twice");
195            }
196        }
197        for _ in 0..100 {
198            assert!(all.insert(issue_single().get()));
199        }
200    }
201
202    #[test]
203    fn test_connection_id_new_and_display() {
204        assert!(ConnectionId::new(0).is_none());
205        assert_eq!(ConnectionId::new(9).unwrap().to_string(), "#9");
206    }
207}
208
209#[cfg(all(test, net_io_loom))]
210mod loom_tests {
211    use super::*;
212    use loom::sync::Arc;
213
214    #[test]
215    fn loom_id_blocks_never_overlap() {
216        loom::model(|| {
217            let counter = Arc::new(loom::sync::atomic::AtomicU64::new(1));
218            let c2 = counter.clone();
219            let t = loom::thread::spawn(move || {
220                let mut b = IdBlock::new();
221                b.issue_with(|| c2.fetch_add(BLOCK, loom::sync::atomic::Ordering::Relaxed))
222                    .get()
223            });
224            let mut b = IdBlock::new();
225            let a = b
226                .issue_with(|| counter.fetch_add(BLOCK, loom::sync::atomic::Ordering::Relaxed))
227                .get();
228            let other = t.join().unwrap();
229            assert_ne!(a, other);
230            assert!(a.abs_diff(other) >= BLOCK);
231        });
232    }
233}