mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-03-31 17:55:36 +00:00
Move all crates to new crates dir
This commit is contained in:
parent
3835da22ac
commit
9b032f7e24
128 changed files with 27 additions and 26 deletions
|
|
@ -1,254 +0,0 @@
|
|||
use std::{
|
||||
cell::RefCell,
|
||||
convert::TryInto,
|
||||
io::Cursor,
|
||||
rc::Rc,
|
||||
sync::{atomic::Ordering, Arc},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use aquatic_http_protocol::response::Response;
|
||||
use futures_lite::{AsyncReadExt, AsyncWriteExt};
|
||||
use futures_rustls::{client::TlsStream, TlsConnector};
|
||||
use glommio::net::TcpStream;
|
||||
use glommio::{prelude::*, timer::TimerActionRepeat};
|
||||
use rand::{prelude::SmallRng, SeedableRng};
|
||||
|
||||
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));
|
||||
let rng = Rc::new(RefCell::new(SmallRng::from_entropy()));
|
||||
|
||||
let interval = config.connection_creation_interval_ms;
|
||||
|
||||
if interval == 0 {
|
||||
loop {
|
||||
if *num_active_connections.borrow() < config.num_connections {
|
||||
if let Err(err) = Connection::run(
|
||||
config.clone(),
|
||||
tls_config.clone(),
|
||||
load_test_state.clone(),
|
||||
num_active_connections.clone(),
|
||||
rng.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
::log::error!("connection creation error: {:?}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let interval = Duration::from_millis(interval);
|
||||
|
||||
TimerActionRepeat::repeat(move || {
|
||||
periodically_open_connections(
|
||||
config.clone(),
|
||||
interval,
|
||||
tls_config.clone(),
|
||||
load_test_state.clone(),
|
||||
num_active_connections.clone(),
|
||||
rng.clone(),
|
||||
)
|
||||
});
|
||||
}
|
||||
|
||||
futures_lite::future::pending::<bool>().await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn periodically_open_connections(
|
||||
config: Rc<Config>,
|
||||
interval: Duration,
|
||||
tls_config: Arc<rustls::ClientConfig>,
|
||||
load_test_state: LoadTestState,
|
||||
num_active_connections: Rc<RefCell<usize>>,
|
||||
rng: Rc<RefCell<SmallRng>>,
|
||||
) -> Option<Duration> {
|
||||
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,
|
||||
rng.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
::log::error!("connection creation error: {:?}", err);
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
|
||||
Some(interval)
|
||||
}
|
||||
|
||||
struct Connection {
|
||||
config: Rc<Config>,
|
||||
load_test_state: LoadTestState,
|
||||
rng: Rc<RefCell<SmallRng>>,
|
||||
stream: TlsStream<TcpStream>,
|
||||
buffer: [u8; 2048],
|
||||
}
|
||||
|
||||
impl Connection {
|
||||
async fn run(
|
||||
config: Rc<Config>,
|
||||
tls_config: Arc<rustls::ClientConfig>,
|
||||
load_test_state: LoadTestState,
|
||||
num_active_connections: Rc<RefCell<usize>>,
|
||||
rng: Rc<RefCell<SmallRng>>,
|
||||
) -> anyhow::Result<()> {
|
||||
let stream = TcpStream::connect(config.server_address)
|
||||
.await
|
||||
.map_err(|err| anyhow::anyhow!("connect: {:?}", err))?;
|
||||
|
||||
let stream = TlsConnector::from(tls_config)
|
||||
.connect("example.com".try_into().unwrap(), stream)
|
||||
.await?;
|
||||
|
||||
let mut connection = Connection {
|
||||
config,
|
||||
load_test_state,
|
||||
rng,
|
||||
stream,
|
||||
buffer: [0; 2048],
|
||||
};
|
||||
|
||||
*num_active_connections.borrow_mut() += 1;
|
||||
|
||||
if let Err(err) = connection.run_connection_loop().await {
|
||||
::log::info!("connection error: {:?}", err);
|
||||
}
|
||||
|
||||
*num_active_connections.borrow_mut() -= 1;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_connection_loop(&mut self) -> anyhow::Result<()> {
|
||||
loop {
|
||||
self.send_request().await?;
|
||||
self.read_response().await?;
|
||||
|
||||
if !self.config.keep_alive {
|
||||
break Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_request(&mut self) -> anyhow::Result<()> {
|
||||
let request = create_random_request(
|
||||
&self.config,
|
||||
&self.load_test_state,
|
||||
&mut self.rng.borrow_mut(),
|
||||
);
|
||||
|
||||
let mut cursor = Cursor::new(&mut self.buffer[..]);
|
||||
|
||||
request.write(&mut cursor, self.config.url_suffix.as_bytes())?;
|
||||
|
||||
let cursor_position = cursor.position() as usize;
|
||||
|
||||
let bytes_sent = self
|
||||
.stream
|
||||
.write(&cursor.into_inner()[..cursor_position])
|
||||
.await?;
|
||||
|
||||
self.stream.flush().await?;
|
||||
|
||||
self.load_test_state
|
||||
.statistics
|
||||
.bytes_sent
|
||||
.fetch_add(bytes_sent, Ordering::Relaxed);
|
||||
|
||||
self.load_test_state
|
||||
.statistics
|
||||
.requests
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn read_response(&mut self) -> anyhow::Result<()> {
|
||||
let mut buffer_position = 0;
|
||||
|
||||
loop {
|
||||
let bytes_read = self
|
||||
.stream
|
||||
.read(&mut self.buffer[buffer_position..])
|
||||
.await?;
|
||||
|
||||
if bytes_read == 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
buffer_position += bytes_read;
|
||||
|
||||
let interesting_bytes = &self.buffer[..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 {
|
||||
match Response::from_bytes(&interesting_bytes[body_start_index..]) {
|
||||
Ok(response) => {
|
||||
match response {
|
||||
Response::Announce(_) => {
|
||||
self.load_test_state
|
||||
.statistics
|
||||
.responses_announce
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
Response::Scrape(_) => {
|
||||
self.load_test_state
|
||||
.statistics
|
||||
.responses_scrape
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
Response::Failure(response) => {
|
||||
self.load_test_state
|
||||
.statistics
|
||||
.responses_failure
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
println!("failure response: reason: {}", response.failure_reason);
|
||||
}
|
||||
}
|
||||
|
||||
self.load_test_state
|
||||
.statistics
|
||||
.bytes_received
|
||||
.fetch_add(interesting_bytes.len(), Ordering::Relaxed);
|
||||
|
||||
break;
|
||||
}
|
||||
Err(err) => {
|
||||
::log::warn!(
|
||||
"deserialize response error with {} bytes read: {:?}, text: {}",
|
||||
buffer_position,
|
||||
err,
|
||||
interesting_bytes.escape_ascii()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue