cmkl/spring-2025/sen-109/00010/rust-chat-server/server/examples/stress_test.rs

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 {}
}