aquatic_http_load_test: add glommio implementation

aquatic_http_load_test/src/glommio.rs

@@ -0,0 +1,260 @@
+use std::{cell::RefCell, convert::TryInto, io::{Cursor, ErrorKind, Read}, rc::Rc, sync::{Arc, atomic::Ordering}, time::Duration};
+
+use aquatic_http_protocol::response::Response;
+use futures_lite::{AsyncReadExt, AsyncWriteExt};
+use glommio::{enclose, prelude::*, timer::TimerActionRepeat};
+use glommio::net::TcpStream;
+use rand::{SeedableRng, prelude::SmallRng};
+use rustls::ClientConnection;
+
+use crate::{common::LoadTestState, config::Config, utils::create_random_request};
+
+pub async fn run_socket_thread(
+    config: Config,
+    tls_config: Arc<rustls::ClientConfig>,
+    load_test_state: LoadTestState,
+) -> anyhow::Result<()> {
+    let config = Rc::new(config);
+    let num_active_connections = Rc::new(RefCell::new(0usize));
+
+    TimerActionRepeat::repeat(enclose!((config, tls_config, load_test_state, num_active_connections) move || {
+        enclose!((config, tls_config, load_test_state, num_active_connections) move || async move {
+            if *num_active_connections.borrow() < config.num_connections {
+                spawn_local(async move {
+                    if let Err(err) = Connection::run(config, tls_config, load_test_state, num_active_connections).await {
+                        eprintln!("connection creation error: {:?}", err);
+                    }
+                }).detach();
+            }
+
+            Some(Duration::from_secs(1))
+        })()
+    }));
+
+    futures_lite::future::pending::<bool>().await;
+
+    Ok(())
+}
+
+struct Connection {
+    config: Rc<Config>,
+    load_test_state: LoadTestState,
+    rng: SmallRng,
+    stream: TcpStream,
+    tls: ClientConnection,
+    response_buffer: [u8; 2048],
+    response_buffer_position: usize,
+    send_new_request: bool,
+    queued_responses: usize,
+}
+
+impl Connection {
+    async fn run(
+        config: Rc<Config>,
+        tls_config: Arc<rustls::ClientConfig>,
+        load_test_state: LoadTestState,
+        num_active_connections: Rc<RefCell<usize>>,
+    ) -> anyhow::Result<()> {
+        let stream = TcpStream::connect(config.server_address).await
+            .map_err(|err| anyhow::anyhow!("connect: {:?}", err))?;
+        let tls = ClientConnection::new(tls_config, "example.com".try_into().unwrap()).unwrap();
+        let rng = SmallRng::from_entropy();
+
+        let mut connection = Connection {
+            config,
+            load_test_state,
+            rng,
+            stream,
+            tls,
+            response_buffer: [0; 2048],
+            response_buffer_position: 0,
+            send_new_request: true,
+            queued_responses: 0,
+        };
+
+        *num_active_connections.borrow_mut() += 1;
+
+        println!("run connection");
+
+        if let Err(err) = connection.run_connection_loop().await {
+            eprintln!("connection error: {:?}", err);
+        }
+
+        *num_active_connections.borrow_mut() -= 1;
+
+        Ok(())
+    }
+
+    async fn run_connection_loop(&mut self) -> anyhow::Result<()> {
+        loop {
+            if self.send_new_request {
+                let request = create_random_request(&self.config, &self.load_test_state, &mut self.rng);
+
+                request.write(&mut self.tls.writer())?;
+                self.queued_responses += 1;
+
+                self.send_new_request = false;
+            }
+
+            self.write_tls().await?;
+            self.read_tls().await?;
+        }
+    }
+
+    async fn read_tls(&mut self) -> anyhow::Result<()> {
+        loop {
+            let mut buf = [0u8; 1024];
+
+            let bytes_read = self.stream.read(&mut buf).await?;
+
+            if bytes_read == 0 {
+                return Err(anyhow::anyhow!("Peer has closed connection"));
+            }
+
+            self
+                .load_test_state
+                .statistics
+                .bytes_received
+                .fetch_add(bytes_read, Ordering::SeqCst);
+
+            let _ = self.tls.read_tls(&mut &buf[..bytes_read]).unwrap();
+
+            let io_state = self.tls.process_new_packets()?;
+
+            let mut added_plaintext = false;
+
+            if io_state.plaintext_bytes_to_read() != 0 {
+                loop {
+                    match self.tls.reader().read(&mut buf) {
+                        Ok(0) => {
+                            break;
+                        }
+                        Ok(amt) => {
+                            let end = self.response_buffer_position + amt;
+
+                            if end > self.response_buffer.len() {
+                                return Err(anyhow::anyhow!("response too large"));
+                            } else {
+                                let response_buffer_slice = &mut self.response_buffer[self.response_buffer_position..end];
+
+                                response_buffer_slice.copy_from_slice(&buf[..amt]);
+
+                                self.response_buffer_position = end;
+
+                                added_plaintext = true;
+                            }
+                        }
+                        Err(err) if err.kind() == ErrorKind::WouldBlock => {
+                            break;
+                        }
+                        Err(err) => {
+                            break;
+                        }
+                    }
+                }
+            }
+
+            if added_plaintext {
+                let interesting_bytes = &self.response_buffer[..self.response_buffer_position];
+
+                let mut opt_body_start_index = None;
+
+                for (i, chunk) in interesting_bytes.windows(4).enumerate() {
+                    if chunk == b"\r\n\r\n" {
+                        opt_body_start_index = Some(i + 4);
+
+                        break;
+                    }
+                }
+
+                if let Some(body_start_index) = opt_body_start_index {
+                    let interesting_bytes = &interesting_bytes[body_start_index..];
+
+                    match Response::from_bytes(interesting_bytes) {
+                        Ok(response) => {
+
+                            match response {
+                                Response::Announce(_) => {
+                                    self
+                                        .load_test_state
+                                        .statistics
+                                        .responses_announce
+                                        .fetch_add(1, Ordering::SeqCst);
+                                }
+                                Response::Scrape(_) => {
+                                    self
+                                        .load_test_state
+                                        .statistics
+                                        .responses_scrape
+                                        .fetch_add(1, Ordering::SeqCst);
+                                }
+                                Response::Failure(response) => {
+                                    self
+                                        .load_test_state
+                                        .statistics
+                                        .responses_failure
+                                        .fetch_add(1, Ordering::SeqCst);
+                                    println!(
+                                        "failure response: reason: {}",
+                                        response.failure_reason
+                                    );
+                                }
+                            }
+
+                            self.response_buffer_position = 0;
+                            self.send_new_request = true;
+
+                            break;
+                        }
+                        Err(err) => {
+                            eprintln!(
+                                "deserialize response error with {} bytes read: {:?}, text: {}",
+                                self.response_buffer_position,
+                                err,
+                                String::from_utf8_lossy(interesting_bytes)
+                            );
+                        }
+                    }
+                }
+            }
+
+            if self.tls.wants_write() {
+                break;
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn write_tls(&mut self) -> anyhow::Result<()> {
+        if !self.tls.wants_write() {
+            return Ok(());
+        }
+
+        let mut buf = Vec::new();
+        let mut buf = Cursor::new(&mut buf);
+
+        while self.tls.wants_write() {
+            self.tls.write_tls(&mut buf).unwrap();
+        }
+
+        let len = buf.get_ref().len();
+
+        self.stream.write_all(&buf.into_inner()).await?;
+        self.stream.flush().await?;
+
+        self
+            .load_test_state
+            .statistics
+            .bytes_sent
+            .fetch_add(len, Ordering::SeqCst);
+
+        if self.queued_responses != 0 {
+            self.load_test_state.statistics.requests.fetch_add(self.queued_responses, Ordering::SeqCst);
+
+            self.queued_responses = 0;
+        }
+
+        Ok(())
+    }
+}