cf_mach/nq_core/connection/
map.rs

1// Copyright (c) 2023-2024 Cloudflare, Inc.
2// Licensed under the BSD-3-Clause license found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
3
4use std::collections::VecDeque;
5use std::net::SocketAddr;
6use std::sync::Arc;
7
8use anyhow::Result;
9use tokio::sync::RwLock;
10use tokio_util::sync::CancellationToken;
11
12use crate::nq_core::connection::http::{
13    EstablishedConnection, start_h1_conn, start_h2_conn, tls_connection,
14};
15use crate::nq_core::util::ByteStream;
16use crate::nq_core::{ConnectionTiming, ConnectionType, Time};
17
18/// Creates and holds [`EstablishedConnection`]s in a VecDeque.
19#[derive(Default, Debug)]
20pub struct ConnectionManager {
21    connections: RwLock<VecDeque<Arc<RwLock<EstablishedConnection>>>>,
22}
23
24impl ConnectionManager {
25    /// Creates a new connection on the given io.
26    #[allow(clippy::too_many_arguments)]
27    pub async fn new_connection(
28        &self,
29        mut timing: ConnectionTiming,
30        remote_addr: SocketAddr,
31        domain: String,
32        conn_type: ConnectionType,
33        io: Box<dyn ByteStream>,
34        time: &dyn Time,
35        shutdown: CancellationToken,
36    ) -> Result<Arc<RwLock<EstablishedConnection>>> {
37        let connection = match conn_type {
38            ConnectionType::H1 { use_tls } => {
39                if use_tls {
40                    let stream = tls_connection(conn_type, &domain, &mut timing, io, time).await?;
41                    start_h1_conn(domain, timing, stream, time, shutdown).await?
42                } else {
43                    start_h1_conn(domain, timing, io, time, shutdown).await?
44                }
45            }
46            ConnectionType::H2 => {
47                let stream = tls_connection(conn_type, &domain, &mut timing, io, time).await?;
48                start_h2_conn(remote_addr, domain, timing, stream, time, shutdown).await?
49            }
50            ConnectionType::H3 => todo!(),
51        };
52
53        let connection = Arc::new(RwLock::new(connection));
54        self.connections.write().await.push_back(connection.clone());
55        Ok(connection)
56    }
57
58    /// Drop all `SendRequest` structs, effectively cancelling all connections.
59    pub async fn shutdown(&self) {
60        for connection in self.connections.write().await.iter_mut() {
61            let mut conn = connection.write().await;
62            conn.drop_send_request();
63        }
64    }
65}
cf_mach/nq_core/connection/map.rs

cf_mach/nq_core/connection/
map.rs
