karyon_p2p/
slots.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use karyon_core::async_util::CondWait;
4
5/// Manages available inbound and outbound slots.
6pub struct ConnectionSlots {
7    /// A condvar for notifying when a slot become available.
8    signal: CondWait,
9    /// The number of occupied slots
10    slots: AtomicUsize,
11    /// The maximum number of slots.
12    max_slots: usize,
13}
14
15impl ConnectionSlots {
16    /// Creates a new ConnectionSlots
17    pub fn new(max_slots: usize) -> Self {
18        Self {
19            signal: CondWait::new(),
20            slots: AtomicUsize::new(0),
21            max_slots,
22        }
23    }
24
25    /// Increases the occupied slots by one.
26    pub fn add(&self) {
27        self.slots.fetch_add(1, Ordering::SeqCst);
28    }
29
30    /// Decreases the occupied slots by one and notifies the waiting signal
31    /// to start accepting/connecting new connections.
32    pub async fn remove(&self) {
33        self.slots.fetch_sub(1, Ordering::SeqCst);
34        if self.slots.load(Ordering::SeqCst) < self.max_slots {
35            self.signal.signal().await;
36        }
37    }
38
39    /// Waits for a slot to become available.
40    pub async fn wait_for_slot(&self) {
41        if self.slots.load(Ordering::SeqCst) < self.max_slots {
42            return;
43        }
44
45        // Wait for a signal
46        self.signal.wait().await;
47        self.signal.reset().await;
48    }
49}