From 303f63afd6a3773912c3ad1d36b5570abf2ea683 Mon Sep 17 00:00:00 2001 From: Gabriel Kaszewski Date: Thu, 16 May 2024 00:39:03 +0200 Subject: [PATCH] Add rate limiting, improve memory usage --- Cargo.lock | 215 ++++++++++++++++++++++++++++++++++++ Cargo.toml | 4 + painter-js/src/canvas.js | 78 ++++++++----- painter-js/src/constants.js | 6 +- painter-js/src/main.js | 60 ++++++---- painter-js/src/socket.js | 18 +-- painter-js/src/style.css | 4 - src/main.rs | 213 +++++++++++++++++++++++++---------- 8 files changed, 481 insertions(+), 117 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 00de12a..9f43e9f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -125,6 +125,15 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + [[package]] name = "bitflags" version = "2.5.0" @@ -202,6 +211,12 @@ dependencies = [ "libc", ] +[[package]] +name = "crossbeam-utils" +version = "0.8.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345" + [[package]] name = "crypto-common" version = "0.1.6" @@ -212,6 +227,19 @@ dependencies = [ "typenum", ] +[[package]] +name = "dashmap" +version = "5.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +dependencies = [ + "cfg-if", + "hashbrown", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "data-encoding" version = "2.6.0" @@ -275,6 +303,31 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "forwarded-header-value" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8835f84f38484cc86f110a805655697908257fb9a7af005234060891557198e9" +dependencies = [ + "nonempty", + "thiserror", +] + +[[package]] +name = "futures" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.30" @@ -282,6 +335,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -290,6 +344,34 @@ version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" +[[package]] +name = "futures-executor" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" + +[[package]] +name = "futures-macro" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.30" @@ -302,15 +384,25 @@ version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + [[package]] name = "futures-util" version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" dependencies = [ + "futures-channel", "futures-core", + "futures-io", + "futures-macro", "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", "slab", @@ -343,6 +435,32 @@ version = "0.28.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4271d37baee1b8c7e4b708028c57d816cf9d2434acb33a549475f78c181f6253" +[[package]] +name = "governor" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68a7f542ee6b35af73b06abc0dad1c1bae89964e4e253bc4b587b91c9637867b" +dependencies = [ + "cfg-if", + "dashmap", + "futures", + "futures-timer", + "no-std-compat", + "nonzero_ext", + "parking_lot", + "portable-atomic", + "quanta", + "rand", + "smallvec", + "spinning_top", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + [[package]] name = "hermit-abi" version = "0.3.9" @@ -524,6 +642,16 @@ version = "2.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c8640c5d730cb13ebd907d8d04b52f55ac9a2eec55b440c8892f40d56c76c1d" +[[package]] +name = "memory-stats" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34f79cf9964c5c9545493acda1263f1912f8d2c56c8a2ffee2606cb960acaacc" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "mime" version = "0.3.17" @@ -560,6 +688,24 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "no-std-compat" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" + +[[package]] +name = "nonempty" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9e591e719385e6ebaeb5ce5d3887f7d5676fceca6411d1925ccc95745f3d6f7" + +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + [[package]] name = "nu-ansi-term" version = "0.46.0" @@ -615,13 +761,17 @@ name = "painter" version = "0.1.0" dependencies = [ "axum", + "bincode", + "bytes", "chrono", "dotenv", + "memory-stats", "serde", "serde_json", "socketioxide", "tokio", "tower-http", + "tower_governor", "tracing", "tracing-subscriber", ] @@ -687,6 +837,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "portable-atomic" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0" + [[package]] name = "ppv-lite86" version = "0.2.17" @@ -702,6 +858,21 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "quanta" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5167a477619228a0b284fac2674e3c388cba90631d7b7de620e6f1fcd08da5" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + [[package]] name = "quote" version = "1.0.36" @@ -741,6 +912,15 @@ dependencies = [ "getrandom", ] +[[package]] +name = "raw-cpuid" +version = "11.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e29830cbb1290e404f24c73af91c5d8d631ce7e128691e9477556b540cd01ecd" +dependencies = [ + "bitflags", +] + [[package]] name = "redox_syscall" version = "0.5.1" @@ -903,6 +1083,15 @@ dependencies = [ "tower", ] +[[package]] +name = "spinning_top" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96d2d1d716fb500937168cc09353ffdc7a012be8475ac7308e1bdf0e3923300" +dependencies = [ + "lock_api", +] + [[package]] name = "syn" version = "2.0.63" @@ -1079,6 +1268,22 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6bc1c9ce2b5135ac7f93c72918fc37feb872bdc6a5533a8b85eb4b86bfdae52" +[[package]] +name = "tower_governor" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "313fa625fea5790ed56360a30ea980e41229cf482b4835801a67ef1922bf63b9" +dependencies = [ + "axum", + "forwarded-header-value", + "governor", + "http", + "pin-project", + "thiserror", + "tower", + "tracing", +] + [[package]] name = "tracing" version = "0.1.40" @@ -1281,6 +1486,16 @@ version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af190c94f2773fdb3729c55b007a722abb5384da03bc0986df4c289bf5567e96" +[[package]] +name = "web-sys" +version = "0.3.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77afa9a11836342370f4817622a2f0f418b134426d91a82dfb48f532d2ec13ef" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/Cargo.toml b/Cargo.toml index 1c64523..d0f737a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,12 +7,16 @@ edition = "2021" [dependencies] axum = "0.7.5" +bincode = "1.3.3" +bytes = "1.6.0" chrono = "0.4.38" dotenv = "0.15.0" +memory-stats = "1.1.0" serde = { version = "1.0.201", features = ["derive"] } serde_json = "1.0.117" socketioxide = "0.13.1" tokio = { version = "1.37.0", features = ["full"] } tower-http = { version = "0.5.2", features = ["cors", "fs"] } +tower_governor = "0.4.2" tracing = "0.1.40" tracing-subscriber = "0.3.18" diff --git a/painter-js/src/canvas.js b/painter-js/src/canvas.js index 8c65b77..31366d3 100644 --- a/painter-js/src/canvas.js +++ b/painter-js/src/canvas.js @@ -1,11 +1,18 @@ -import socket from "./socket.js"; import { hexToU32, u32ToHex, getColorFromElementCSS, rgbToHex, } from "./utils.js"; -import { pixelSize, pixelCooldown } from "./constants.js"; +import { + pixelSize, + pixelCooldown, + canvasEndpoint, + WIDTH, + HEIGHT, +} from "./constants.js"; + +let socket = null; const canvas = document.getElementById("canvas"); const ctx = canvas.getContext("2d"); @@ -97,37 +104,29 @@ const handleColorPicker = () => { }); }; -handleColorPicker(); +const fetchCanvasState = async () => { + fetch(canvasEndpoint) + .then((response) => response.json()) + .then((data) => { + canvasState = data; + drawCanvasState(data); + }) + .catch((error) => { + alert("Error fetching canvas state from server. Please try again later."); + }); +}; const drawCanvasState = (canvasState) => { - for (let y = 0; y < canvasState.length; y++) { - for (let x = 0; x < canvasState[y].length; x++) { - const color = u32ToHex(canvasState[y][x]); + for (let y = 0; y < HEIGHT; y++) { + for (let x = 0; x < WIDTH; x++) { + const index = y * WIDTH + x; + const color = u32ToHex(canvasState[index]); ctx.fillStyle = color; ctx.fillRect(x * pixelSize, y * pixelSize, pixelSize, pixelSize); } } }; -socket.on("init-canvas", (data) => { - const canvasData = JSON.parse(data); - drawCanvasState(canvasData); - canvasState = canvasData; -}); - -socket.on("pixel-updated", (update) => { - const color = u32ToHex(update.color); - ctx.fillStyle = color; - ctx.fillRect( - update.x * pixelSize, - update.y * pixelSize, - pixelSize, - pixelSize - ); - - canvasState[update.y][update.x] = update.color; -}); - const checkIfCanPlacePixel = () => { const now = Date.now(); return now - lastPixelTime >= pixelCooldown; @@ -145,7 +144,8 @@ const handlePlacePixel = (pixelData) => { } socket.emit("place-pixel", pixelData); - canvasState[pixelData.y][pixelData.x] = pixelData.color; + const index = pixelData.y * WIDTH + pixelData.x; + canvasState[index] = pixelData.color; setLastPixelTime(); pixelData = null; }; @@ -218,8 +218,6 @@ const handleToggleGrid = () => { }); }; -handleToggleGrid(); - window.onkeydown = (event) => { // on enter (keycode 13 is enter) if (event.keyCode === 13) { @@ -243,3 +241,27 @@ saveCanvasButton.addEventListener("click", () => { a.download = "canvas.png"; a.click(); }); + +handleColorPicker(); +handleToggleGrid(); + +export const handleSocketEvents = (_socket) => { + socket = _socket; + socket.on("connect", () => { + fetchCanvasState(); + }); + + socket.on("pixel-updated", (update) => { + const color = u32ToHex(update.color); + ctx.fillStyle = color; + ctx.fillRect( + update.x * pixelSize, + update.y * pixelSize, + pixelSize, + pixelSize + ); + + const index = update.y * WIDTH + update.x; + canvasState[index] = update.color; + }); +}; diff --git a/painter-js/src/constants.js b/painter-js/src/constants.js index 2a8443a..3a2d004 100644 --- a/painter-js/src/constants.js +++ b/painter-js/src/constants.js @@ -1,2 +1,6 @@ export const pixelSize = 10; -export const pixelCooldown = 60 * 1000; // 1 minute +export const pixelCooldown = 10 * 1000; // 10 seconds +export const canvasEndpoint = "http://localhost:3000/canvas/"; +export const checkEndpoint = "http://localhost:3000/check/"; +export const WIDTH = 500; +export const HEIGHT = 500; diff --git a/painter-js/src/main.js b/painter-js/src/main.js index 73f36ab..2c55247 100644 --- a/painter-js/src/main.js +++ b/painter-js/src/main.js @@ -1,8 +1,9 @@ -import socket from "./socket.js"; +import { connectToWS } from "./socket.js"; import "./canvas.js"; import "./counter.js"; import { updateCountdown } from "./counter.js"; -import { pixelSize } from "./constants.js"; +import { checkEndpoint, pixelSize } from "./constants.js"; +import { handleSocketEvents } from "./canvas.js"; const isDebug = import.meta.env.VITE_IS_DEBUG === "true"; @@ -11,27 +12,46 @@ const currentSoldiersSpan = document.getElementById("current-soldiers"); let coords = []; const canvas = document.getElementById("canvas"); const coordsText = document.getElementById("coords"); +const ogCanvasStyle = canvas.style.display; +canvas.style.display = "none"; -socket.on("connect", () => { - console.log("connect"); -}); +fetch(checkEndpoint) + .then((response) => { + if (response.ok) { + const socket = connectToWS(); -socket.on("error", (message) => { - alert(message); -}); + socket.on("connect", () => { + canvas.style.display = ogCanvasStyle; + console.log("connect"); + }); -socket.on("current_soldiers", (currentSoldiers) => { - currentSoldiersSpan.textContent = currentSoldiers; -}); + socket.on("error", (message) => { + alert(message); + }); -requestAnimationFrame(updateCountdown); + socket.on("current_soldiers", (currentSoldiers) => { + currentSoldiersSpan.textContent = currentSoldiers; + }); -window.addEventListener("mousemove", (event) => { - // get coordinates of the mouse inside the canvas - const rect = canvas.getBoundingClientRect(); - const x = Math.floor((event.clientX - rect.left) / pixelSize); - const y = Math.floor((event.clientY - rect.top) / pixelSize); - coords = [x, y]; + handleSocketEvents(socket); - coordsText.textContent = `${x}, ${y}`; -}); + requestAnimationFrame(updateCountdown); + + window.addEventListener("mousemove", (event) => { + // get coordinates of the mouse inside the canvas + const rect = canvas.getBoundingClientRect(); + const x = Math.floor((event.clientX - rect.left) / pixelSize); + const y = Math.floor((event.clientY - rect.top) / pixelSize); + coords = [x, y]; + + coordsText.textContent = `${x}, ${y}`; + }); + } else { + throw new Error("Can't connect to the server"); + } + }) + .catch((error) => { + alert( + "You have already connected to the server from another tab or window. Please close the other tab or window and refresh this page." + ); + }); diff --git a/painter-js/src/socket.js b/painter-js/src/socket.js index 83c292f..5774edb 100644 --- a/painter-js/src/socket.js +++ b/painter-js/src/socket.js @@ -6,12 +6,16 @@ const wsHost = window.location.host; let socket; -if (isDebug) { - socket = io("ws://localhost:3000"); -} else { - socket = io(`${wsProtocol}//${wsHost}`, { - transports: ["websocket"], - }); -} +export const connectToWS = () => { + if (isDebug) { + socket = io("ws://localhost:3000"); + } else { + socket = io(`${wsProtocol}//${wsHost}`, { + transports: ["websocket"], + }); + } + + return socket; +}; export default socket; diff --git a/painter-js/src/style.css b/painter-js/src/style.css index 961e1ee..15550e5 100644 --- a/painter-js/src/style.css +++ b/painter-js/src/style.css @@ -6,8 +6,6 @@ .rainbow-border { position: relative; display: inline-block; - border-radius: 0.5rem; /* Rounded corners */ - padding: 5px; /* Space for the border */ background: linear-gradient(to right, red, orange, yellow, green, blue, indigo, violet); } .rainbow-border::before { @@ -17,8 +15,6 @@ right: 0; bottom: 0; left: 0; - margin: -5px; /* Negative margin to overlap the gradient */ - border-radius: 0.5rem; /* Match the rounded corners */ background: white; z-index: -1; /* Place behind the content */ } diff --git a/src/main.rs b/src/main.rs index e7640c9..572ac4a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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>>; +type LastUpdate = Arc>>>; +type Soldiers = Arc>>; + #[derive(Serialize, Deserialize, Debug)] struct PixelUpdate { x: u32, @@ -25,99 +35,178 @@ struct PixelUpdate { color: u32, } -struct AppState { - canvas: Mutex>>, - last_update: Mutex>>, - soldiers: Mutex>, -} - -async fn on_connect(socket: SocketRef, state: Arc) { +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::(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) -> Json> { + let canvas = canvas.lock().unwrap(); + Json(canvas.clone()) +} + #[tokio::main] async fn main() -> Result<(), Box> { 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> { }; 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> { 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::(), + ) + .await?; Ok(()) }