karyon_p2p/
conn_queue.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
use std::{collections::VecDeque, sync::Arc};

use karyon_core::{async_runtime::lock::Mutex, async_util::CondVar};

use crate::{connection::ConnDirection, connection::Connection, ConnRef, Result};

/// Connection queue
pub struct ConnQueue {
    queue: Mutex<VecDeque<Connection>>,
    conn_available: CondVar,
}

impl ConnQueue {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            queue: Mutex::new(VecDeque::new()),
            conn_available: CondVar::new(),
        })
    }

    /// Handle a connection by pushing it into the queue and wait for the disconnect signal
    pub async fn handle(&self, conn: ConnRef, direction: ConnDirection) -> Result<()> {
        let endpoint = conn.peer_endpoint()?;

        let (disconnect_tx, disconnect_rx) = async_channel::bounded(1);
        let new_conn = Connection::new(conn, disconnect_tx, direction, endpoint);

        // Push a new conn to the queue
        self.queue.lock().await.push_back(new_conn);
        self.conn_available.signal();

        // Wait for the disconnect signal from the connection handler
        if let Ok(result) = disconnect_rx.recv().await {
            return result;
        }

        Ok(())
    }

    /// Waits for the next connection in the queue
    pub async fn next(&self) -> Connection {
        let mut queue = self.queue.lock().await;
        while queue.is_empty() {
            queue = self.conn_available.wait(queue).await;
        }
        queue.pop_front().unwrap()
    }
}