GKaszewski/painter
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add rate limiting, improve memory usage

Browse Source
This commit is contained in:
Gabriel Kaszewski
2024-05-16 00:39:03 +02:00
parent f70cc988e9
commit 303f63afd6
8 changed files with 481 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

213
src/main.rs
View File

@@ -1,16 +1,22 @@
use std::{
collections::{HashMap, HashSet},
net::SocketAddr,
sync::{Arc, Mutex},
};
use axum::Router;
use axum::{routing::get, Extension, Json, Router, ServiceExt};
use chrono::{DateTime, Duration, Utc};
use memory_stats::memory_stats;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use socketioxide::{
extract::{Bin, Data, SocketRef},
SocketIo,
};
use tower_governor::{
governor::{GovernorConfig, GovernorConfigBuilder},
GovernorLayer,
};
use tower_http::{
cors::{Any, CorsLayer},
services::ServeDir,
@@ -18,6 +24,10 @@ use tower_http::{
use tracing::info;
use tracing_subscriber::FmtSubscriber;
type Canvas = Arc<Mutex<Vec<u32>>>;
type LastUpdate = Arc<Mutex<HashMap<String, DateTime<Utc>>>>;
type Soldiers = Arc<Mutex<HashSet<String>>>;
#[derive(Serialize, Deserialize, Debug)]
struct PixelUpdate {
x: u32,
@@ -25,99 +35,178 @@ struct PixelUpdate {
color: u32,
}
struct AppState {
canvas: Mutex<Vec<Vec<u32>>>,
last_update: Mutex<HashMap<String, DateTime<Utc>>>,
soldiers: Mutex<HashSet<String>>,
}
async fn on_connect(socket: SocketRef, state: Arc<AppState>) {
async fn on_connect(
socket: SocketRef,
canvas: Canvas,
last_update: LastUpdate,
soldiers: Soldiers,
) {
info!("Socket connected: {:?} {:?}", socket.ns(), socket.id,);
let canvas_state = state.canvas.lock().unwrap();
let serialized_canvas = serde_json::to_string(&*canvas_state).unwrap();
socket.emit("init-canvas", &serialized_canvas).ok();
let socket_id = socket.id.to_string();
let mut soldiers = state.soldiers.lock().unwrap();
soldiers.insert(socket_id.clone());
let soldiers_num = soldiers.len();
socket.emit("current_soldiers", &soldiers_num).ok();
socket
.broadcast()
.emit("current_soldiers", &soldiers_num)
.ok();
{
let mut soldiers = soldiers.lock().unwrap();
soldiers.insert(socket_id.clone());
let soldiers_num = soldiers.len();
let app_state = state.clone();
let app_state_for_disconnect = state.clone();
socket.emit("current_soldiers", &soldiers_num).ok();
socket
.broadcast()
.emit("current_soldiers", &soldiers_num)
.ok();
}
info!("Memory usage after connection and before cloning state: ");
print_memory_usage();
let canvas_clone = canvas.clone();
let last_update_clone = last_update.clone();
info!("Memory usage after connection and after cloning state: ");
print_memory_usage();
socket.on(
"place-pixel",
move |socket: SocketRef, Data::<PixelUpdate>(update), Bin(_bin)| {
let now = Utc::now();
let mut last_update = app_state.last_update.lock().unwrap();
let socket_id = socket.id.to_string();
let last_update = last_update_clone.clone();
let canvas = canvas_clone.clone();
if let Some(&last_time) = last_update.get(&socket_id) {
if now < last_time + Duration::minutes(1) {
let _ = socket.emit(
"error",
Value::String("You can only place one pixel per minute".to_string()),
);
return;
{
let now = Utc::now();
let mut last_update = last_update.lock().unwrap();
if let Some(&last_time) = last_update.get(&socket_id) {
if now < last_time + Duration::seconds(10) {
let _ = socket.emit(
"error",
Value::String("You can only place one pixel per minute".to_string()),
);
return;
}
}
last_update.insert(socket_id.clone(), now);
}
last_update.insert(socket_id.clone(), now);
info!("Received pixel update: {:?}", update);
let mut canvas = app_state.canvas.lock().unwrap();
canvas[update.y as usize][update.x as usize] = update.color;
let mut canvas = canvas.lock().unwrap();
canvas[update.y as usize * 500 + update.x as usize] = update.color;
info!("Emitting pixel update");
socket.emit("pixel-updated", &update).ok(); // Send to the user who placed the pixel
socket.broadcast().emit("pixel-updated", &update).ok(); // Send to all other users
info!("Memory usage after pixel update: ");
print_memory_usage();
},
);
socket.on_disconnect(move |socket: SocketRef| {
info!("Socket disconnected: {:?}", socket.id);
let mut soldiers = app_state_for_disconnect.soldiers.lock().unwrap();
soldiers.remove(&socket.id.to_string());
info!("Soldiers: {:?}", soldiers.len());
let mut last_update = app_state_for_disconnect.last_update.lock().unwrap();
info!("Memory usage after disconnection: ");
print_memory_usage();
let socket_id = socket.id.to_string();
last_update.remove(&socket_id);
info!("Last update: {:?}", last_update.len());
{
let mut soldiers = soldiers.lock().unwrap();
soldiers.remove(&socket_id.to_string());
info!("Soldiers: {:?}", soldiers.len());
let soldiers_num = soldiers.len();
socket
.broadcast()
.emit("current_soldiers", &soldiers_num)
.ok();
let soldiers_num = soldiers.len();
socket
.broadcast()
.emit("current_soldiers", &soldiers_num)
.ok();
}
{
let mut last_update = last_update.lock().unwrap();
last_update.remove(&socket_id);
info!("Last update: {:?}", last_update.len());
}
info!("Memory usage after disconnection and cleanup: ");
print_memory_usage();
});
}
fn print_memory_usage() {
if let Some(usage) = memory_stats() {
info!(
"Current physical memory usage: {} MB",
usage.physical_mem as f64 / 1024.0 / 1024.0
);
info!(
"Current virtual memory usage: {} MB",
usage.virtual_mem as f64 / 1024.0 / 1024.0
);
}
}
async fn get_canvas_state(Extension(canvas): Extension<Canvas>) -> Json<Vec<u32>> {
let canvas = canvas.lock().unwrap();
Json(canvas.clone())
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
dotenv::dotenv().ok();
tracing::subscriber::set_global_default(FmtSubscriber::new())?;
let rate_governor = Arc::new(
GovernorConfigBuilder::default()
.burst_size(6)
.per_second(10)
.finish()
.unwrap(),
);
let governor = rate_governor.limiter().clone();
let interval = std::time::Duration::from_secs(1);
tokio::spawn(async move {
loop {
tokio::time::sleep(interval).await;
governor.retain_recent();
}
});
let (layer, io) = SocketIo::new_layer();
let app_state = Arc::new(AppState {
canvas: Mutex::new(vec![vec![0xFFFFFFFF; 500]; 500]),
last_update: Mutex::new(HashMap::new()),
soldiers: Mutex::new(HashSet::new()),
});
let canvas = Arc::new(Mutex::new(vec![0xFFFFFFFF; 500 * 500]));
let last_update = Arc::new(Mutex::new(HashMap::new()));
let soldiers = Arc::new(Mutex::new(HashSet::new()));
let app_state_for_ns = app_state.clone();
let used_memory = std::mem::size_of_val(&*canvas)
+ std::mem::size_of_val(&*last_update)
+ std::mem::size_of_val(&*soldiers);
info!(
"Used memory of state: {} bytes, {} KB, {} MB",
used_memory,
used_memory / 1024,
used_memory / 1024 / 1024
);
info!("Memory usage after state setup and before socket.io setup: ");
print_memory_usage();
let canvas_for_socket = canvas.clone();
io.ns("/", move |socket: SocketRef| {
tokio::spawn(on_connect(socket, app_state_for_ns.clone()));
tokio::spawn(on_connect(
socket,
canvas_for_socket.clone(),
last_update.clone(),
soldiers.clone(),
));
});
info!("Memory usage after socket.io setup: ");
print_memory_usage();
let address = std::env::var("ADDRESS").unwrap_or_else(|_| "0.0.0.0".to_string());
let port: u16 = std::env::var("PORT")
.unwrap_or_else(|_| "3000".to_string())
@@ -137,8 +226,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
};
let app = Router::new()
.fallback_service(ServeDir::new("dist"))
.layer(layer);
.route("/canvas/", get(get_canvas_state))
.route("/check/", get(|| async { "OK" }))
.fallback_service(ServeDir::new("dist/"))
.layer(GovernorLayer {
config: rate_governor,
})
.layer(layer)
.layer(Extension(canvas.clone()));
let app = if let Some(cors) = cors {
app.layer(cors)
@@ -149,7 +244,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let server_address = format!("{}:{}", address, port);
info!("Starting server on {}", server_address);
let listener = tokio::net::TcpListener::bind(server_address).await?;
axum::serve(listener, app.into_make_service()).await?;
axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.await?;
Ok(())
}