169 lines
5.5 KiB
Rust
169 lines
5.5 KiB
Rust
|
use std::time::Duration;
|
||
|
|
||
|
use comms::{
|
||
|
command::{JoinRoomCommand, UserCommand},
|
||
|
event::Event,
|
||
|
transport,
|
||
|
};
|
||
|
use nanoid::nanoid;
|
||
|
use rand::{rngs::StdRng, Rng, SeedableRng};
|
||
|
use serde::{Deserialize, Serialize};
|
||
|
use tokio::{net::TcpStream, task::JoinSet};
|
||
|
use tokio_stream::StreamExt;
|
||
|
|
||
|
/// Stress Test for the Chat Server
|
||
|
///
|
||
|
/// Generates synthetic load with users who joins and sends messages to random roms.
|
||
|
/// The number of users, number of rooms joined per user and chatting of users can be configured.
|
||
|
///
|
||
|
/// !IMPORTANT! Be sure to check and configure your socket limits, before you run the tests
|
||
|
|
||
|
const SERVER_ADDR: &str = "localhost:8080";
|
||
|
const CHAT_ROOMS_METADATA: &str = include_str!("../resources/chat_rooms_metadata.json");
|
||
|
|
||
|
/// Stress Test Configuration
|
||
|
// The number of users to spawn, distributed across the load increments
|
||
|
const LOAD_INCREMENTS: &str = r#"[
|
||
|
{ "user_count": 1200, "after": { "secs": 60, "nanos": 0 }, "steps": 60 },
|
||
|
{ "user_count": 2400, "after": { "secs": 120, "nanos": 0 }, "steps": 60 }
|
||
|
]"#;
|
||
|
// How many rooms a user should join, this affects the total tokio task count
|
||
|
const NUMBER_OF_ROOMS_TO_JOIN: usize = 5;
|
||
|
// How many milliseconds to wait between each user message
|
||
|
const USER_CHAT_DELAY_MILLIS: u64 = 10_000;
|
||
|
|
||
|
/// [RotatingIterator] is a simple iterator that rotates through a list of items
|
||
|
/// and starts from the beginning when the end is reached.
|
||
|
struct RotatingIterator<T> {
|
||
|
items: Vec<T>,
|
||
|
current: usize,
|
||
|
}
|
||
|
|
||
|
impl<T> RotatingIterator<T> {
|
||
|
fn new(items: Vec<T>) -> Self {
|
||
|
Self { items, current: 0 }
|
||
|
}
|
||
|
}
|
||
|
|
||
|
impl<T: Clone> Iterator for RotatingIterator<T> {
|
||
|
type Item = T;
|
||
|
|
||
|
fn next(&mut self) -> Option<Self::Item> {
|
||
|
let item = self.items.get(self.current).cloned();
|
||
|
self.current = (self.current + 1) % self.items.len();
|
||
|
item
|
||
|
}
|
||
|
}
|
||
|
|
||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
|
struct ChatRoomMetadata {
|
||
|
name: String,
|
||
|
description: String,
|
||
|
}
|
||
|
|
||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
|
struct LoadIncrements {
|
||
|
user_count: usize,
|
||
|
after: Duration,
|
||
|
steps: usize,
|
||
|
}
|
||
|
|
||
|
async fn spawn_single_user(rooms_to_join: Vec<String>) -> anyhow::Result<()> {
|
||
|
let result = spawn_single_user_raw(rooms_to_join).await;
|
||
|
|
||
|
match result.as_ref() {
|
||
|
Ok(_) => println!("exited without problems"),
|
||
|
Err(err) => println!("some error occurred = {}", err.to_string()),
|
||
|
}
|
||
|
|
||
|
result
|
||
|
}
|
||
|
|
||
|
async fn spawn_single_user_raw(rooms_to_join: Vec<String>) -> anyhow::Result<()> {
|
||
|
let tcp_stream = TcpStream::connect(SERVER_ADDR).await?;
|
||
|
let (mut event_stream, mut command_writer) = transport::client::split_tcp_stream(tcp_stream);
|
||
|
|
||
|
let _login_event = match event_stream.next().await {
|
||
|
Some(Ok(Event::LoginSuccessful(login_event))) => login_event,
|
||
|
_ => return Err(anyhow::anyhow!("server did not send login successful")),
|
||
|
};
|
||
|
|
||
|
for room_name in rooms_to_join.iter() {
|
||
|
command_writer
|
||
|
.write(&UserCommand::JoinRoom(JoinRoomCommand {
|
||
|
room: String::from(room_name),
|
||
|
}))
|
||
|
.await?;
|
||
|
}
|
||
|
|
||
|
let join_handle = tokio::spawn({
|
||
|
let mut rng = StdRng::from_entropy();
|
||
|
let mut rooms_iterator = RotatingIterator::new(rooms_to_join);
|
||
|
let to_sleep = Duration::from_millis(USER_CHAT_DELAY_MILLIS);
|
||
|
|
||
|
async move {
|
||
|
// sleep initially for a time to distribute the messaging times
|
||
|
tokio::time::sleep(Duration::from_millis(
|
||
|
rng.gen_range(1..USER_CHAT_DELAY_MILLIS),
|
||
|
))
|
||
|
.await;
|
||
|
|
||
|
loop {
|
||
|
let room_name = rooms_iterator.next().unwrap();
|
||
|
let _ = command_writer
|
||
|
.write(&UserCommand::SendMessage(
|
||
|
comms::command::SendMessageCommand {
|
||
|
room: String::from(room_name),
|
||
|
content: nanoid!(),
|
||
|
},
|
||
|
))
|
||
|
.await;
|
||
|
|
||
|
tokio::time::sleep(to_sleep).await;
|
||
|
}
|
||
|
}
|
||
|
});
|
||
|
|
||
|
while let Some(_) = event_stream.next().await {}
|
||
|
|
||
|
join_handle.abort();
|
||
|
Ok(())
|
||
|
}
|
||
|
|
||
|
#[tokio::main]
|
||
|
async fn main() {
|
||
|
let load_increments: Vec<LoadIncrements> =
|
||
|
serde_json::from_str(LOAD_INCREMENTS).expect("could not parse the load increments");
|
||
|
let chat_room_metadata: Vec<ChatRoomMetadata> = serde_json::from_str(CHAT_ROOMS_METADATA)
|
||
|
.expect("could not parse the chat rooms metadata");
|
||
|
|
||
|
let mut room_iterator = RotatingIterator::new(chat_room_metadata);
|
||
|
let mut join_set: JoinSet<anyhow::Result<()>> = JoinSet::new();
|
||
|
|
||
|
let mut current: usize = 0;
|
||
|
for li in load_increments {
|
||
|
let diff = li.user_count - current;
|
||
|
let sleep_duration =
|
||
|
Duration::from_millis((li.after.as_millis() / li.steps as u128) as u64);
|
||
|
let to_increment = diff / li.steps;
|
||
|
|
||
|
for _ in 0..li.steps {
|
||
|
for _ in 0..to_increment {
|
||
|
let rooms_to_join = room_iterator
|
||
|
.by_ref()
|
||
|
.take(NUMBER_OF_ROOMS_TO_JOIN)
|
||
|
.map(|metadata| metadata.name.clone())
|
||
|
.collect();
|
||
|
|
||
|
join_set.spawn(spawn_single_user(rooms_to_join));
|
||
|
}
|
||
|
|
||
|
current += to_increment;
|
||
|
println!("total users: {}", current);
|
||
|
tokio::time::sleep(sleep_duration).await;
|
||
|
}
|
||
|
}
|
||
|
|
||
|
while let Some(_) = join_set.join_next().await {}
|
||
|
}
|