1use std::fmt;
4use std::num::NonZeroU64;
5
6use std::sync::atomic::{AtomicU64, Ordering};
7
8#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
21pub struct ConnectionId(NonZeroU64);
22
23impl ConnectionId {
24 #[inline]
26 #[must_use]
27 pub const fn get(self) -> u64 {
28 self.0.get()
29 }
30
31 #[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#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
59pub struct ListenerId(pub(crate) u16);
60
61impl ListenerId {
62 #[must_use]
64 pub const fn get(self) -> u16 {
65 self.0
66 }
67}
68
69#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
78pub struct WorkerId(pub(crate) u16);
79
80impl WorkerId {
81 #[must_use]
83 pub const fn get(self) -> u16 {
84 self.0
85 }
86
87 #[must_use]
89 pub const fn first() -> WorkerId {
90 WorkerId(0)
91 }
92
93 #[must_use]
100 pub const fn for_testing(index: u16) -> WorkerId {
101 WorkerId(index)
102 }
103}
104
105pub(crate) const BLOCK: u64 = 1024;
107
108static NEXT: AtomicU64 = AtomicU64::new(1);
111
112#[derive(Debug)]
114pub(crate) struct IdBlock {
115 next: u64,
116 end: u64,
117}
118
119impl IdBlock {
120 pub(crate) const fn new() -> IdBlock {
122 IdBlock { next: 0, end: 0 }
123 }
124
125 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 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 None => ConnectionId(NonZeroU64::MIN),
148 }
149}
150
151pub(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}