karyon_p2p/
conn_queue.rs

1use std::{collections::VecDeque, sync::Arc};
2
3use karyon_core::{async_runtime::lock::Mutex, async_util::CondVar};
4
5use crate::{connection::ConnDirection, connection::Connection, ConnRef, Result};
6
7/// Connection queue
8pub struct ConnQueue {
9    queue: Mutex<VecDeque<Connection>>,
10    conn_available: CondVar,
11}
12
13impl ConnQueue {
14    pub fn new() -> Arc<Self> {
15        Arc::new(Self {
16            queue: Mutex::new(VecDeque::new()),
17            conn_available: CondVar::new(),
18        })
19    }
20
21    /// Handle a connection by pushing it into the queue and wait for the disconnect signal
22    pub async fn handle(&self, conn: ConnRef, direction: ConnDirection) -> Result<()> {
23        let endpoint = conn.peer_endpoint()?;
24
25        let (disconnect_tx, disconnect_rx) = async_channel::bounded(1);
26        let new_conn = Connection::new(conn, disconnect_tx, direction, endpoint);
27
28        // Push a new conn to the queue
29        self.queue.lock().await.push_back(new_conn);
30        self.conn_available.signal();
31
32        // Wait for the disconnect signal from the connection handler
33        if let Ok(result) = disconnect_rx.recv().await {
34            return result;
35        }
36
37        Ok(())
38    }
39
40    /// Waits for the next connection in the queue
41    pub async fn next(&self) -> Connection {
42        let mut queue = self.queue.lock().await;
43        while queue.is_empty() {
44            queue = self.conn_available.wait(queue).await;
45        }
46        queue.pop_front().unwrap()
47    }
48}