feat(s2n-quic-qns): add ability to set congestion controller in inter…

…op/perf
aws · Aug 17, 2023 · 529cd9b · 529cd9b
1 parent c7e3f51
commit 529cd9b
Show file tree

Hide file tree

Showing 9 changed files with 203 additions and 72 deletions.
diff --git a/quic/s2n-quic-qns/src/client.rs b/quic/s2n-quic-qns/src/client.rs
@@ -8,3 +8,62 @@ pub mod perf;
 
 pub use interop::Interop;
 pub use perf::Perf;
+
+use crate::{
+    congestion_control::{CongestionControl, CongestionController},
+    tls,
+    tls::TlsProviders,
+    Result,
+};
+use s2n_quic::{
+    client,
+    client::ClientProviders,
+    provider::congestion_controller::{Bbr, Cubic},
+};
+
+/// Build and start a client with the given TLS configuration and Congestion Controller
+pub fn build(
+    builder: client::Builder<impl ClientProviders>,
+    alpns: &[String],
+    tls_client: &tls::Client,
+    congestion_control: &CongestionControl,
+) -> Result<s2n_quic::Client> {
+    macro_rules! build {
+        ($build_tls:ident, $cc:ident $(, $alpns:ident)?) => {
+            {
+                let tls = tls_client.$build_tls($($alpns)?)?;
+
+                builder
+                    .with_tls(tls)?
+                    .with_congestion_controller($cc::default())?
+                    .start()
+                    .unwrap()
+            }
+        }
+    }
+
+    Ok(
+        match (tls_client.tls, congestion_control.congestion_controller) {
+            #[cfg(unix)]
+            (TlsProviders::S2N, CongestionController::Cubic) => {
+                build!(build_s2n_tls, Cubic, alpns)
+            }
+            #[cfg(unix)]
+            (TlsProviders::S2N, CongestionController::Bbr) => {
+                build!(build_s2n_tls, Bbr, alpns)
+            }
+            (TlsProviders::Rustls, CongestionController::Cubic) => {
+                build!(build_rustls, Cubic, alpns)
+            }
+            (TlsProviders::Rustls, CongestionController::Bbr) => {
+                build!(build_rustls, Bbr, alpns)
+            }
+            (TlsProviders::Null, CongestionController::Cubic) => {
+                build!(build_null, Cubic)
+            }
+            (TlsProviders::Null, CongestionController::Bbr) => {
+                build!(build_null, Bbr)
+            }
+        },
+    )
+}
diff --git a/quic/s2n-quic-qns/src/client/interop.rs b/quic/s2n-quic-qns/src/client/interop.rs
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
+    client,
     client::{h09, h3},
     interop::Testcase,
     task, tls, Result,
@@ -54,6 +55,9 @@ pub struct Interop {
 
     #[structopt(flatten)]
     runtime: crate::runtime::Runtime,
+
+    #[structopt(flatten)]
+    congestion_controller: crate::congestion_control::CongestionControl,
 }
 
 impl Interop {
@@ -157,9 +161,12 @@ impl Interop {
             .with_limits(limits)?
             .with_event(event::tracing::Subscriber::default())?;
 
-        let client = self.tls.build(client, &self.application_protocols)?;
-
-        Ok(client)
+        client::build(
+            client,
+            &self.application_protocols,
+            &self.tls,
+            &self.congestion_controller,
+        )
     }
 
     async fn endpoints(&self) -> Result<HashMap<Host<&str>, Connect>> {

diff --git a/quic/s2n-quic-qns/src/client/perf.rs b/quic/s2n-quic-qns/src/client/perf.rs
@@ -1,8 +1,8 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{perf, task, tls, Result};
-use s2n_quic::{client, provider::event, Client, Connection};
+use crate::{client, perf, task, tls, Result};
+use s2n_quic::{client::Connect, provider::event, Client, Connection};
 use structopt::StructOpt;
 
 #[derive(Debug, StructOpt)]
@@ -53,6 +53,9 @@ pub struct Perf {
 
     #[structopt(flatten)]
     runtime: crate::runtime::Runtime,
+
+    #[structopt(flatten)]
+    congestion_controller: crate::congestion_control::CongestionControl,
 }
 
 impl Perf {
@@ -69,7 +72,7 @@ impl Perf {
         let send = self.send;
         let receive = self.receive;
 
-        let mut connect = client::Connect::new((self.ip, self.port));
+        let mut connect = Connect::new((self.ip, self.port));
         if let Some(server_name) = self.server_name.as_deref() {
             connect = connect.with_server_name(server_name);
         } else {
@@ -153,8 +156,11 @@ impl Perf {
             .with_io(io)?
             .with_event(subscriber)?;
 
-        let client = self.tls.build(client, &self.application_protocols)?;
-
-        Ok(client)
+        client::build(
+            client,
+            &self.application_protocols,
+            &self.tls,
+            &self.congestion_controller,
+        )
     }
 }
diff --git a/quic/s2n-quic-qns/src/congestion_control.rs b/quic/s2n-quic-qns/src/congestion_control.rs
@@ -0,0 +1,36 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use core::str::FromStr;
+use std::io;
+use structopt::StructOpt;
+
+#[derive(Debug, StructOpt)]
+pub struct CongestionControl {
+    /// The congestion controller to use
+    #[structopt(long = "cc", default_value = "bbr", possible_values = &["cubic","bbr"])]
+    pub congestion_controller: CongestionController,
+}
+
+#[derive(Copy, Clone, Debug)]
+pub enum CongestionController {
+    Cubic,
+    Bbr,
+}
+
+impl FromStr for CongestionController {
+    type Err = crate::Error;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "cubic" => Ok(Self::Cubic),
+            "bbr" => Ok(Self::Bbr),
+            _ => {
+                Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    format!("Unsupported congestion controller: {s}"),
+                )
+                .into())
+            }
+        }
+    }
+}
diff --git a/quic/s2n-quic-qns/src/main.rs b/quic/s2n-quic-qns/src/main.rs
@@ -7,6 +7,7 @@ pub type Error = Box<dyn 'static + std::error::Error + Send + Sync>;
 pub type Result<T, E = Error> = core::result::Result<T, E>;
 
 mod client;
+mod congestion_control;
 mod file;
 mod interop;
 mod io;

diff --git a/quic/s2n-quic-qns/src/server.rs b/quic/s2n-quic-qns/src/server.rs
@@ -10,3 +10,62 @@ mod unstable;
 
 pub use interop::Interop;
 pub use perf::Perf;
+
+use crate::{
+    congestion_control::{CongestionControl, CongestionController},
+    tls,
+    tls::TlsProviders,
+    Result,
+};
+use s2n_quic::{
+    provider::congestion_controller::{Bbr, Cubic},
+    server,
+    server::ServerProviders,
+};
+
+/// Build and start a server with the given TLS configuration and Congestion Controller
+pub fn build(
+    builder: server::Builder<impl ServerProviders>,
+    alpns: &[String],
+    tls_server: &tls::Server,
+    congestion_control: &CongestionControl,
+) -> Result<s2n_quic::Server> {
+    macro_rules! build {
+        ($build_tls:ident, $cc:ident $(, $alpns:ident)?) => {
+            {
+                let tls = tls_server.$build_tls($($alpns)?)?;
+
+                builder
+                    .with_tls(tls)?
+                    .with_congestion_controller($cc::default())?
+                    .start()
+                    .unwrap()
+            }
+        }
+    }
+
+    Ok(
+        match (tls_server.tls, congestion_control.congestion_controller) {
+            #[cfg(unix)]
+            (TlsProviders::S2N, CongestionController::Cubic) => {
+                build!(build_s2n_tls, Cubic, alpns)
+            }
+            #[cfg(unix)]
+            (TlsProviders::S2N, CongestionController::Bbr) => {
+                build!(build_s2n_tls, Bbr, alpns)
+            }
+            (TlsProviders::Rustls, CongestionController::Cubic) => {
+                build!(build_rustls, Cubic, alpns)
+            }
+            (TlsProviders::Rustls, CongestionController::Bbr) => {
+                build!(build_rustls, Bbr, alpns)
+            }
+            (TlsProviders::Null, CongestionController::Cubic) => {
+                build!(build_null, Cubic)
+            }
+            (TlsProviders::Null, CongestionController::Bbr) => {
+                build!(build_null, Bbr)
+            }
+        },
+    )
+}
diff --git a/quic/s2n-quic-qns/src/server/interop.rs b/quic/s2n-quic-qns/src/server/interop.rs
@@ -3,6 +3,7 @@
 
 use crate::{
     interop::Testcase,
+    server,
     server::{h09, h3},
     tls, Result,
 };
@@ -42,6 +43,9 @@ pub struct Interop {
 
     #[structopt(flatten)]
     runtime: crate::runtime::Runtime,
+
+    #[structopt(flatten)]
+    congestion_controller: crate::congestion_control::CongestionControl,
 }
 
 impl Interop {
@@ -100,7 +104,12 @@ impl Interop {
                 EventSubscriber(1),
                 s2n_quic::provider::event::tracing::Subscriber::default(),
             ))?;
-        let server = self.tls.build(server, &self.application_protocols)?;
+        let server = server::build(
+            server,
+            &self.application_protocols,
+            &self.tls,
+            &self.congestion_controller,
+        )?;
 
         eprintln!("Server listening on port {}", self.io.port);
 

diff --git a/quic/s2n-quic-qns/src/server/perf.rs b/quic/s2n-quic-qns/src/server/perf.rs
@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{perf, tls, Result};
+use crate::{perf, server, tls, Result};
 use futures::future::try_join_all;
 use s2n_quic::{
     provider::event,
@@ -36,6 +36,9 @@ pub struct Perf {
 
     #[structopt(flatten)]
     runtime: crate::runtime::Runtime,
+
+    #[structopt(flatten)]
+    congestion_controller: crate::congestion_control::CongestionControl,
 }
 
 impl Perf {
@@ -189,7 +192,12 @@ impl Perf {
             .with_io(io)?
             .with_event(subscriber)?;
 
-        let server = self.tls.build(server, &self.application_protocols)?;
+        let server = server::build(
+            server,
+            &self.application_protocols,
+            &self.tls,
+            &self.congestion_controller,
+        )?;
 
         eprintln!("Server listening on port {}", self.io.port);