netstack3_tcp/socket/
accept_queue.rs

1// Copyright 2023 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Defines a thread safe accept queue to be shared between listener and
6//! connected sockets.
7
8use alloc::collections::VecDeque;
9use alloc::sync::Arc;
10#[cfg(test)]
11use alloc::vec::Vec;
12use core::fmt::Debug;
13use core::hash::Hash;
14use core::ops::DerefMut;
15
16use assert_matches::assert_matches;
17use derivative::Derivative;
18use netstack3_base::Inspector;
19use netstack3_base::sync::{BindingsSyncTypes, Mutex};
20use netstack3_hashmap::HashMap;
21
22/// A change that occurred on an [`AcceptQueue`].
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum AcceptQueueChange {
25    /// A pending connection completed the handshake and was added to the queue.
26    ConnectionAdded {
27        /// Whether the accept queue was empty before adding this connection.
28        was_empty: bool,
29    },
30    /// A connection was removed from the queue (via `accept` or reset/close).
31    ConnectionRemoved {
32        /// Whether the queue is now empty after removing this connection.
33        is_empty: bool,
34    },
35}
36
37/// A notifier used to tell Bindings about ready connections for a single
38/// listening socket.
39pub trait ListenerNotifier {
40    /// Called when the accept queue has changed.
41    fn accept_queue_changed(&mut self, change: AcceptQueueChange);
42}
43
44/// A thread-safe shareable accept queue.
45///
46/// [`AcceptQueue`] implements all necessary locking internally and can be
47/// safely shared between a listener and the connected sockets generated from
48/// it.
49#[derive(Derivative)]
50#[derivative(Clone(bound = ""), Debug(bound = "S: Debug, R: Debug, N: Debug"))]
51pub struct AcceptQueue<S, R, N, BT: BindingsSyncTypes>(
52    Arc<Mutex<AcceptQueueInner<S, R, N>, BT::Mutex>>,
53);
54
55#[cfg(test)]
56impl<S, R, N, BT: BindingsSyncTypes> PartialEq for AcceptQueue<S, R, N, BT>
57where
58    AcceptQueueInner<S, R, N>: PartialEq,
59{
60    fn eq(&self, Self(other): &Self) -> bool {
61        let Self(inner) = self;
62        if Arc::ptr_eq(other, inner) {
63            return true;
64        }
65        let guard = inner.lock();
66        let other_guard = other.lock();
67        (&*guard).eq(&*other_guard)
68    }
69}
70
71#[cfg(test)]
72impl<S, R, N, BT: BindingsSyncTypes> Eq for AcceptQueue<S, R, N, BT> where Self: PartialEq {}
73
74#[derive(Debug)]
75#[cfg_attr(test, derive(Eq, PartialEq))]
76enum EntryState {
77    Pending,
78    Ready,
79}
80
81#[derive(Debug)]
82#[cfg_attr(test, derive(Derivative))]
83#[cfg_attr(
84    test,
85    derivative(PartialEq(bound = "S: Hash + Clone + Eq + PartialEq, R: PartialEq, N: PartialEq"))
86)]
87struct AcceptQueueInner<S, R, N> {
88    ready_queue: VecDeque<(S, R)>,
89    all_sockets: HashMap<S, EntryState>,
90    notifier: Option<N>,
91}
92
93impl<S, R, N, BT: BindingsSyncTypes> AcceptQueue<S, R, N, BT>
94where
95    S: Hash + Clone + Eq + PartialEq + Debug,
96    N: ListenerNotifier,
97{
98    /// Creates a new [`AcceptQueue`] with `notifier`.
99    pub(crate) fn new(notifier: N) -> Self {
100        Self(Arc::new(Mutex::new(AcceptQueueInner::new(notifier))))
101    }
102
103    fn lock(&self) -> impl DerefMut<Target = AcceptQueueInner<S, R, N>> + '_ {
104        let Self(inner) = self;
105        inner.lock()
106    }
107
108    /// Pops a ready entry from the queue.
109    ///
110    /// Returns the socket and its associated ready state if the queue was not
111    /// empty.
112    ///
113    /// `pop_ready` also hits the notifier.
114    pub(crate) fn pop_ready(&self) -> Option<(S, R)> {
115        self.lock().pop_ready()
116    }
117
118    /// Collects all currently pending sockets.
119    #[cfg(test)]
120    pub(crate) fn collect_pending(&self) -> Vec<S> {
121        self.lock().collect_pending()
122    }
123
124    /// Pushes a new `pending` connection.
125    ///
126    /// Panics if `pending` is already in the queue or if the queue is already
127    /// closed.
128    pub(crate) fn push_pending(&self, pending: S) {
129        self.lock().push_pending(pending)
130    }
131
132    /// Returns the total number of entries in the queue (pending and ready).
133    pub(crate) fn len(&self) -> usize {
134        self.lock().len()
135    }
136
137    /// Returns the number of ready entries in the queue.
138    #[cfg(test)]
139    pub(crate) fn ready_len(&self) -> usize {
140        self.lock().ready_len()
141    }
142
143    /// Returns the number of pending entries in the queue.
144    #[cfg(test)]
145    pub(crate) fn pending_len(&self) -> usize {
146        self.lock().pending_len()
147    }
148
149    /// Notifies that [`newly_ready`] is now ready to be accepted.
150    ///
151    /// Does nothing if the queue is closed.
152    ///
153    /// Panics if the queue is not closed and `newly_ready` was not in the queue
154    /// as a pending entry.
155    pub(crate) fn notify_ready(&self, newly_ready: &S, ready_state: R) {
156        self.lock().notify_ready(newly_ready, ready_state)
157    }
158
159    /// Removes an entry from the queue.
160    ///
161    /// No-op if the queue is closed or if the entry was not in the queue.
162    pub(crate) fn remove(&self, entry: &S) {
163        self.lock().remove(entry)
164    }
165
166    /// Closes the queue and returns all entries that were in the queue and the
167    /// notifier.
168    ///
169    /// Panics if the queue is already closed.
170    pub(crate) fn close(&self) -> (impl Iterator<Item = S> + use<S, R, N, BT>, N) {
171        self.lock().close()
172    }
173
174    /// Returns true if the queue is closed.
175    pub(crate) fn is_closed(&self) -> bool {
176        self.lock().is_closed()
177    }
178
179    /// Inspects the queue.
180    pub(crate) fn inspect<I: Inspector>(&self, inspector: &mut I) {
181        let inner = self.lock();
182        inspector.record_usize("NumReady", inner.ready_len());
183        inspector.record_usize("NumPending", inner.pending_len());
184        inspector.record_debug("Contents", &inner.all_sockets);
185    }
186}
187
188impl<S, R, N> AcceptQueueInner<S, R, N>
189where
190    S: Hash + Clone + Eq + PartialEq + Debug,
191    N: ListenerNotifier,
192{
193    fn new(notifier: N) -> Self {
194        Self {
195            ready_queue: Default::default(),
196            all_sockets: Default::default(),
197            notifier: Some(notifier),
198        }
199    }
200
201    fn pop_ready(&mut self) -> Option<(S, R)> {
202        let AcceptQueueInner { ready_queue, all_sockets, notifier } = self;
203        let (socket, ready_state) = ready_queue.pop_front()?;
204        // Assert internal consistency on state.
205        assert_matches!(all_sockets.remove(&socket), Some(EntryState::Ready));
206        // Invariant: queue must be empty if notifier is none.
207        let notifier = notifier.as_mut().unwrap();
208        notifier.accept_queue_changed(AcceptQueueChange::ConnectionRemoved {
209            is_empty: ready_queue.is_empty(),
210        });
211        Some((socket, ready_state))
212    }
213
214    #[cfg(test)]
215    pub(crate) fn collect_pending(&self) -> Vec<S> {
216        let AcceptQueueInner { all_sockets, .. } = self;
217        all_sockets
218            .iter()
219            .filter_map(|(socket, state)| match state {
220                EntryState::Ready => None,
221                EntryState::Pending => Some(socket.clone()),
222            })
223            .collect()
224    }
225
226    fn push_pending(&mut self, pending: S) {
227        let AcceptQueueInner { all_sockets, notifier, .. } = self;
228        // Protect against the listener attempting to push pending
229        // connections if the queue is closed.
230        assert!(notifier.is_some());
231        assert_matches!(all_sockets.insert(pending, EntryState::Pending), None);
232    }
233
234    fn len(&self) -> usize {
235        let AcceptQueueInner { all_sockets, .. } = self;
236        all_sockets.len()
237    }
238
239    fn ready_len(&self) -> usize {
240        let AcceptQueueInner { ready_queue, .. } = self;
241        ready_queue.len()
242    }
243
244    fn pending_len(&self) -> usize {
245        let AcceptQueueInner { ready_queue, all_sockets, .. } = self;
246        all_sockets.len() - ready_queue.len()
247    }
248
249    fn notify_ready(&mut self, newly_ready: &S, ready_state: R) {
250        let AcceptQueueInner { ready_queue, all_sockets, notifier } = self;
251        let notifier = match notifier {
252            Some(notifier) => notifier,
253
254            None => {
255                // Queue is closed, listener must be shutting down.
256                // Maintain the invariant that a closed queue has cleared all
257                // its entries.
258                debug_assert!(ready_queue.is_empty());
259                debug_assert!(all_sockets.is_empty());
260                return;
261            }
262        };
263        let entry = all_sockets
264            .get_mut(newly_ready)
265            .expect("attempted to notify ready entry that was not in queue");
266        let prev_state = core::mem::replace(entry, EntryState::Ready);
267        assert_matches!(prev_state, EntryState::Pending);
268        let was_empty = ready_queue.is_empty();
269        ready_queue.push_back((newly_ready.clone(), ready_state));
270        notifier.accept_queue_changed(AcceptQueueChange::ConnectionAdded { was_empty });
271    }
272
273    fn remove(&mut self, entry: &S) {
274        let AcceptQueueInner { ready_queue, all_sockets, notifier } = self;
275        // Avoid a HashMap lookup if the queue is already closed.
276        let notifier = match notifier.as_mut() {
277            Some(notifier) => notifier,
278            None => {
279                // Maintain the invariant that a closed queue has cleared all
280                // its entries.
281                debug_assert!(ready_queue.is_empty());
282                debug_assert!(all_sockets.is_empty());
283                return;
284            }
285        };
286
287        match all_sockets.remove(entry) {
288            Some(EntryState::Pending) | None => (),
289            Some(EntryState::Ready) => {
290                let before_len = ready_queue.len();
291                ready_queue.retain(|(s, _ready_data)| s != entry);
292                let after_len = ready_queue.len();
293                // Item must only be in ready queue once.
294                assert_eq!(after_len, before_len - 1);
295                notifier.accept_queue_changed(AcceptQueueChange::ConnectionRemoved {
296                    is_empty: ready_queue.is_empty(),
297                });
298            }
299        }
300    }
301
302    fn close(&mut self) -> (impl Iterator<Item = S> + use<S, R, N>, N) {
303        let AcceptQueueInner { ready_queue, all_sockets, notifier } = self;
304        // Remove the notifier, this signals that the queue is closed.
305        let notifier = notifier.take().expect("queue is already closed");
306        // Remove all items from the ready queue by replacing it with an empty
307        // VecDeque. This ensures that it'll immediately free all of its memory
308        // as opposed to just calling clear.
309        *ready_queue = Default::default();
310        // Take all the entries so we can return from the caller, leaving a
311        // fresh empty HashMap in its place.
312        let entries = core::mem::take(all_sockets);
313        (entries.into_keys(), notifier)
314    }
315
316    fn is_closed(&self) -> bool {
317        let AcceptQueueInner { notifier, .. } = self;
318        notifier.is_none()
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use assert_matches::assert_matches;
325    use netstack3_hashmap::HashSet;
326
327    use super::AcceptQueueChange;
328
329    #[test]
330    fn push_ready_pop() {
331        let mut queue = AcceptQueueInner::new(Notifier::default());
332        assert_eq!(queue.pop_ready(), None);
333        assert_eq!(queue.len(), 0);
334        assert_eq!(queue.ready_len(), 0);
335        assert_eq!(queue.pending_len(), 0);
336        assert_eq!(queue.clear_notifier(), None);
337
338        queue.push_pending(Socket(0));
339        assert_eq!(queue.pop_ready(), None);
340        assert_eq!(queue.len(), 1);
341        assert_eq!(queue.ready_len(), 0);
342        assert_eq!(queue.pending_len(), 1);
343        assert_eq!(queue.clear_notifier(), None);
344
345        queue.notify_ready(&Socket(0), Ready(2));
346        assert_eq!(queue.len(), 1);
347        assert_eq!(queue.ready_len(), 1);
348        assert_eq!(queue.pending_len(), 0);
349        assert_eq!(
350            queue.clear_notifier(),
351            Some(AcceptQueueChange::ConnectionAdded { was_empty: true })
352        );
353
354        queue.push_pending(Socket(1));
355        queue.notify_ready(&Socket(1), Ready(3));
356        assert_eq!(queue.len(), 2);
357        assert_eq!(queue.ready_len(), 2);
358        assert_eq!(queue.pending_len(), 0);
359        assert_eq!(
360            queue.clear_notifier(),
361            Some(AcceptQueueChange::ConnectionAdded { was_empty: false })
362        );
363
364        assert_eq!(queue.pop_ready(), Some((Socket(0), Ready(2))));
365        assert_eq!(
366            queue.clear_notifier(),
367            Some(AcceptQueueChange::ConnectionRemoved { is_empty: false })
368        );
369        assert_eq!(queue.len(), 1);
370        assert_eq!(queue.ready_len(), 1);
371        assert_eq!(queue.pending_len(), 0);
372
373        assert_eq!(queue.pop_ready(), Some((Socket(1), Ready(3))));
374        assert_eq!(
375            queue.clear_notifier(),
376            Some(AcceptQueueChange::ConnectionRemoved { is_empty: true })
377        );
378        assert_eq!(queue.len(), 0);
379        assert_eq!(queue.ready_len(), 0);
380        assert_eq!(queue.pending_len(), 0);
381        assert_eq!(queue.pop_ready(), None);
382    }
383
384    #[test]
385    fn close() {
386        let mut queue = AcceptQueueInner::new(Notifier::default());
387        let mut expect = HashSet::new();
388        for i in 0..3 {
389            let s = Socket(i);
390            queue.push_pending(s.clone());
391            assert!(expect.insert(s));
392        }
393        let (socks, _notifier) = queue.close();
394        let got = socks.collect::<HashSet<_>>();
395        assert_eq!(got, expect);
396
397        assert!(queue.is_closed());
398        assert_eq!(queue.len(), 0);
399    }
400
401    #[test]
402    fn remove() {
403        let mut queue = AcceptQueueInner::new(Notifier::default());
404        let s1 = Socket(1);
405        let s2 = Socket(2);
406        queue.push_pending(s1.clone());
407        queue.push_pending(s2.clone());
408        queue.notify_ready(&s2, Ready(2));
409        assert_eq!(queue.len(), 2);
410        assert_eq!(queue.ready_len(), 1);
411        assert_eq!(queue.pending_len(), 1);
412        assert_eq!(
413            queue.clear_notifier(),
414            Some(AcceptQueueChange::ConnectionAdded { was_empty: true })
415        );
416
417        queue.remove(&s1);
418        assert_eq!(queue.len(), 1);
419        assert_eq!(queue.ready_len(), 1);
420        assert_eq!(queue.pending_len(), 0);
421        assert_eq!(queue.clear_notifier(), None);
422
423        queue.remove(&s2);
424        assert_eq!(queue.len(), 0);
425        assert_eq!(queue.ready_len(), 0);
426        assert_eq!(queue.pending_len(), 0);
427        assert_eq!(
428            queue.clear_notifier(),
429            Some(AcceptQueueChange::ConnectionRemoved { is_empty: true })
430        );
431
432        // Removing already removed entry is a no-op.
433        queue.remove(&s1);
434        queue.remove(&s2);
435        assert_eq!(queue.len(), 0);
436        assert_eq!(queue.ready_len(), 0);
437        assert_eq!(queue.pending_len(), 0);
438        assert_eq!(queue.clear_notifier(), None);
439    }
440
441    #[derive(Default, Eq, PartialEq, Debug, Hash, Clone)]
442    struct Socket(usize);
443    #[derive(Default, Eq, PartialEq, Debug)]
444    struct Ready(usize);
445
446    #[derive(Default, Eq, PartialEq, Debug)]
447    struct Notifier(Option<AcceptQueueChange>);
448
449    type AcceptQueueInner = super::AcceptQueueInner<Socket, Ready, Notifier>;
450
451    impl AcceptQueueInner {
452        fn clear_notifier(&mut self) -> Option<AcceptQueueChange> {
453            let Self { notifier, .. } = self;
454            let Notifier(v) = notifier.as_mut().unwrap();
455            v.take()
456        }
457    }
458
459    impl super::ListenerNotifier for Notifier {
460        fn accept_queue_changed(&mut self, change: AcceptQueueChange) {
461            let Self(n) = self;
462            assert_matches!(n.replace(change), None);
463        }
464    }
465}
netstack3_tcp/socket/accept_queue.rs

netstack3_tcp/socket/
accept_queue.rs
