net_io/
completer.rs

1//! The per-worker completion queue: other threads hand results back to a parked connection.
2
3use std::fmt;
4use std::sync::Arc;
5
6#[cfg(net_io_loom)]
7use loom::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
8#[cfg(not(net_io_loom))]
9use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
10
11use crate::id::{ConnectionId, WorkerId};
12
13/// Why a post was refused.
14///
15/// ```
16/// assert_eq!(net_io::PostError::Full.to_string(), "the worker's completion queue is full");
17/// ```
18#[non_exhaustive]
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub enum PostError {
21    /// The ring is full; the poster decides what to do (for example answer 503).
22    Full,
23    /// The connection belongs to another worker.
24    WrongWorker,
25    /// The server has stopped.
26    ServerStopped,
27}
28
29impl fmt::Display for PostError {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        f.write_str(match self {
32            PostError::Full => "the worker's completion queue is full",
33            PostError::WrongWorker => "the connection belongs to another worker",
34            PostError::ServerStopped => "the server has stopped",
35        })
36    }
37}
38
39impl std::error::Error for PostError {}
40
41/// One posted completion.
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43pub(crate) struct Posted {
44    pub connection: ConnectionId,
45    pub token: u32,
46    pub value: u64,
47}
48
49struct Slot {
50    seq: AtomicUsize,
51    id: AtomicU64,
52    token: AtomicU64,
53    value: AtomicU64,
54}
55
56/// A bounded multi-producer, single-consumer ring (Vyukov's bounded queue):
57/// one sequence number per slot, no allocation after creation, never blocks.
58pub(crate) struct Ring {
59    slots: Box<[Slot]>,
60    mask: usize,
61    head: AtomicUsize,
62    tail: AtomicUsize,
63}
64
65impl Ring {
66    /// A ring of `capacity` rounded up to a power of two (at least 1).
67    pub(crate) fn new(capacity: usize) -> Ring {
68        let cap = capacity.max(1).next_power_of_two();
69        let slots = (0..cap)
70            .map(|i| Slot {
71                seq: AtomicUsize::new(i),
72                id: AtomicU64::new(0),
73                token: AtomicU64::new(0),
74                value: AtomicU64::new(0),
75            })
76            .collect::<Vec<_>>()
77            .into_boxed_slice();
78        Ring {
79            slots,
80            mask: cap - 1,
81            head: AtomicUsize::new(0),
82            tail: AtomicUsize::new(0),
83        }
84    }
85
86    pub(crate) fn capacity(&self) -> usize {
87        self.mask + 1
88    }
89
90    /// Pushes from any thread; `false` when full.
91    pub(crate) fn push(&self, item: Posted) -> bool {
92        let mut pos = self.tail.load(Ordering::Relaxed);
93        loop {
94            let slot = &self.slots[pos & self.mask];
95            let seq = slot.seq.load(Ordering::Acquire);
96            let diff = seq as isize - pos as isize;
97            match diff.cmp(&0) {
98                std::cmp::Ordering::Equal => match self.tail.compare_exchange_weak(
99                    pos,
100                    pos.wrapping_add(1),
101                    Ordering::Relaxed,
102                    Ordering::Relaxed,
103                ) {
104                    Ok(_) => {
105                        slot.id.store(item.connection.get(), Ordering::Relaxed);
106                        slot.token.store(u64::from(item.token), Ordering::Relaxed);
107                        slot.value.store(item.value, Ordering::Relaxed);
108                        slot.seq.store(pos.wrapping_add(1), Ordering::Release);
109                        return true;
110                    }
111                    Err(actual) => pos = actual,
112                },
113                // The slot still holds an item the consumer has not taken: full.
114                std::cmp::Ordering::Less => return false,
115                std::cmp::Ordering::Greater => pos = self.tail.load(Ordering::Relaxed),
116            }
117        }
118    }
119
120    /// Pops on the owning thread only.
121    pub(crate) fn pop(&self) -> Option<Posted> {
122        let pos = self.head.load(Ordering::Relaxed);
123        let slot = &self.slots[pos & self.mask];
124        let seq = slot.seq.load(Ordering::Acquire);
125        if seq != pos.wrapping_add(1) {
126            return None;
127        }
128        let item = Posted {
129            connection: ConnectionId::new(slot.id.load(Ordering::Relaxed))?,
130            token: slot.token.load(Ordering::Relaxed) as u32,
131            value: slot.value.load(Ordering::Relaxed),
132        };
133        slot.seq
134            .store(pos.wrapping_add(self.mask + 1), Ordering::Release);
135        self.head.store(pos.wrapping_add(1), Ordering::Relaxed);
136        Some(item)
137    }
138
139    #[cfg(all(test, not(net_io_loom)))]
140    pub(crate) fn is_empty(&self) -> bool {
141        let pos = self.head.load(Ordering::Relaxed);
142        self.slots[pos & self.mask].seq.load(Ordering::Acquire) != pos.wrapping_add(1)
143    }
144}
145
146/// What a completer needs from the worker that owns it.
147pub(crate) trait CompleterTarget: Send + Sync {
148    /// Whether `id` is (or was) a connection of this worker.
149    fn owns(&self, id: ConnectionId) -> bool;
150    /// Wakes the worker.
151    fn wake(&self);
152    /// Whether the server stopped.
153    fn stopped(&self) -> bool;
154}
155
156pub(crate) struct CompleterShared {
157    pub(crate) ring: Ring,
158    pub(crate) wake_pending: AtomicBool,
159    pub(crate) target: Arc<dyn CompleterTarget>,
160    pub(crate) worker: WorkerId,
161}
162
163impl CompleterShared {
164    pub(crate) fn new(
165        capacity: usize,
166        target: Arc<dyn CompleterTarget>,
167        worker: WorkerId,
168    ) -> CompleterShared {
169        CompleterShared {
170            ring: Ring::new(capacity),
171            wake_pending: AtomicBool::new(false),
172            target,
173            worker,
174        }
175    }
176}
177
178/// A worker's completion queue, for other threads to hand results back to a
179/// connection that returned `Step::Park`.
180///
181/// # Examples
182///
183/// ```no_run
184/// use net_io::{Input, Session, SessionContext, Step};
185///
186/// struct Slow;
187/// impl Session for Slow {
188///     fn on_data(&mut self, cx: &mut SessionContext<'_>, input: &mut Input) -> Step {
189///         let n = input.bytes().len();
190///         input.consume(n);
191///         let (completer, id) = (cx.completer(), cx.info().id);
192///         std::thread::spawn(move || {
193///             let _ = completer.post(id, 1, 42);
194///         });
195///         Step::Park { until: None }
196///     }
197///     fn on_wake(&mut self, cx: &mut SessionContext<'_>, _token: u32, value: u64) -> Step {
198///         let _ = cx.output().push(format!("{value}\n").as_bytes());
199///         Step::Write
200///     }
201/// }
202/// ```
203#[derive(Clone)]
204pub struct Completer {
205    pub(crate) shared: Arc<CompleterShared>,
206}
207
208impl Completer {
209    /// Posts a completion for `connection`. Never blocks.
210    ///
211    /// # Errors
212    ///
213    /// [`PostError::Full`], [`PostError::WrongWorker`], [`PostError::ServerStopped`].
214    pub fn post(&self, connection: ConnectionId, token: u32, value: u64) -> Result<(), PostError> {
215        let s = &self.shared;
216        if s.target.stopped() {
217            return Err(PostError::ServerStopped);
218        }
219        if !s.target.owns(connection) {
220            return Err(PostError::WrongWorker);
221        }
222        if !s.ring.push(Posted {
223            connection,
224            token,
225            value,
226        }) {
227            return Err(PostError::Full);
228        }
229        // One wake per burst: only the post that finds no wake pending wakes the worker.
230        if !s.wake_pending.swap(true, Ordering::AcqRel) {
231            s.target.wake();
232        }
233        Ok(())
234    }
235
236    /// The worker this queue belongs to.
237    #[must_use]
238    pub fn worker(&self) -> WorkerId {
239        self.shared.worker
240    }
241
242    /// The ring's capacity.
243    #[must_use]
244    pub fn capacity(&self) -> usize {
245        self.shared.ring.capacity()
246    }
247}
248
249impl fmt::Debug for Completer {
250    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
251        f.debug_struct("Completer")
252            .field("worker", &self.shared.worker)
253            .field("capacity", &self.capacity())
254            .finish()
255    }
256}
257
258#[cfg(all(test, not(net_io_loom)))]
259mod tests {
260    use super::*;
261
262    fn post(id: u64) -> Posted {
263        Posted {
264            connection: ConnectionId::new(id).unwrap(),
265            token: 1,
266            value: id * 10,
267        }
268    }
269
270    #[test]
271    fn test_ring_fifo_and_full() {
272        let r = Ring::new(3);
273        assert_eq!(r.capacity(), 4);
274        for i in 1..=4 {
275            assert!(r.push(post(i)));
276        }
277        assert!(!r.push(post(5)));
278        for i in 1..=4 {
279            assert_eq!(r.pop(), Some(post(i)));
280        }
281        assert_eq!(r.pop(), None);
282        assert!(r.is_empty());
283        assert!(r.push(post(6)));
284        assert_eq!(r.pop(), Some(post(6)));
285    }
286
287    #[test]
288    fn test_ring_many_producers_never_exceeds_capacity() {
289        let r = Arc::new(Ring::new(64));
290        let handles: Vec<_> = (0..8u64)
291            .map(|t| {
292                let r = r.clone();
293                std::thread::spawn(move || {
294                    (0..1000u64)
295                        .filter(|i| r.push(post(t * 1000 + i + 1)))
296                        .count()
297                })
298            })
299            .collect();
300        let pushed: usize = handles.into_iter().map(|h| h.join().unwrap()).sum();
301        let mut popped = 0;
302        while r.pop().is_some() {
303            popped += 1;
304        }
305        assert_eq!(pushed, popped);
306        assert!(popped <= 64);
307    }
308}
309
310#[cfg(all(test, net_io_loom))]
311mod loom_tests {
312    use super::*;
313
314    #[test]
315    fn loom_ring_two_producers_one_consumer() {
316        loom::model(|| {
317            let r = loom::sync::Arc::new(Ring::new(2));
318            let mk = |id: u64| Posted {
319                connection: ConnectionId::new(id).unwrap(),
320                token: 0,
321                value: id,
322            };
323            let r1 = r.clone();
324            let t1 = loom::thread::spawn(move || r1.push(mk(1)));
325            let r2 = r.clone();
326            let t2 = loom::thread::spawn(move || r2.push(mk(2)));
327            let mut got = Vec::new();
328            for _ in 0..3 {
329                if let Some(p) = r.pop() {
330                    got.push(p.value);
331                }
332            }
333            let a = t1.join().unwrap();
334            let b = t2.join().unwrap();
335            while let Some(p) = r.pop() {
336                got.push(p.value);
337            }
338            assert!(a && b);
339            got.sort_unstable();
340            assert_eq!(got, vec![1, 2]);
341        });
342    }
343}
net_io/completer.rs

net_io/
completer.rs
