Add rate limiting, improve memory usage

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"

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"

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 y = 0; y < HEIGHT; y++) {
-    for (let x = 0; x < canvasState[y].length; x++) {
+    for (let x = 0; x < WIDTH; x++) {
-      const color = u32ToHex(canvasState[y][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;
+  });
+};

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;

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,22 +12,32 @@ 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", () => {
+fetch(checkEndpoint)
+  .then((response) => {
+    if (response.ok) {
+      const socket = connectToWS();
+
+      socket.on("connect", () => {
+        canvas.style.display = ogCanvasStyle;
         console.log("connect");
-});
+      });
 
-socket.on("error", (message) => {
+      socket.on("error", (message) => {
         alert(message);
-});
+      });
 
-socket.on("current_soldiers", (currentSoldiers) => {
+      socket.on("current_soldiers", (currentSoldiers) => {
         currentSoldiersSpan.textContent = currentSoldiers;
-});
+      });
 
-requestAnimationFrame(updateCountdown);
+      handleSocketEvents(socket);
 
-window.addEventListener("mousemove", (event) => {
+      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);
@@ -34,4 +45,13 @@ window.addEventListener("mousemove", (event) => {
         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."
+    );
+  });

painter-js/src/socket.js

@@ -6,12 +6,16 @@ const wsHost = window.location.host;
 
 let socket;
 
-if (isDebug) {
+export const connectToWS = () => {
+  if (isDebug) {
     socket = io("ws://localhost:3000");
-} else {
+  } else {
     socket = io(`${wsProtocol}//${wsHost}`, {
       transports: ["websocket"],
     });
-}
+  }
+
+  return socket;
+};
 
 export default socket;

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 */
 }

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<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,40 +35,50 @@ struct PixelUpdate {
     color: u32,
 }
 
-struct AppState {
+async fn on_connect(
-    canvas: Mutex<Vec<Vec<u32>>>,
+    socket: SocketRef,
-    last_update: Mutex<HashMap<String, DateTime<Utc>>>,
+    canvas: Canvas,
-    soldiers: Mutex<HashSet<String>>,
+    last_update: LastUpdate,
-}
+    soldiers: Soldiers,
-
+) {
-async fn on_connect(socket: SocketRef, state: Arc<AppState>) {
     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();
+    {
+        let mut soldiers = 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 app_state = state.clone();
+    info!("Memory usage after connection and before cloning state: ");
-    let app_state_for_disconnect = state.clone();
+    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 socket_id = socket.id.to_string();
+            let last_update = last_update_clone.clone();
+            let canvas = canvas_clone.clone();
+
+            {
                 let now = Utc::now();
-            let mut last_update = app_state.last_update.lock().unwrap();
+                let mut last_update = last_update.lock().unwrap();
 
                 if let Some(&last_time) = last_update.get(&socket_id) {
-                if now < last_time + Duration::minutes(1) {
+                    if now < last_time + Duration::seconds(10) {
                         let _ = socket.emit(
                             "error",
                             Value::String("You can only place one pixel per minute".to_string()),
@@ -68,56 +88,125 @@ async fn on_connect(socket: SocketRef, state: Arc<AppState>) {
                 }
 
                 last_update.insert(socket_id.clone(), now);
+            }
 
             info!("Received pixel update: {:?}", update);
-            let mut canvas = app_state.canvas.lock().unwrap();
+            let mut canvas = canvas.lock().unwrap();
-            canvas[update.y as usize][update.x as usize] = update.color;
+            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();
+        info!("Memory usage after disconnection: ");
-        soldiers.remove(&socket.id.to_string());
+        print_memory_usage();
-
-        info!("Soldiers: {:?}", soldiers.len());
-
-        let mut last_update = app_state_for_disconnect.last_update.lock().unwrap();
         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 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 {
+    let canvas = Arc::new(Mutex::new(vec![0xFFFFFFFF; 500 * 500]));
-        canvas: Mutex::new(vec![vec![0xFFFFFFFF; 500]; 500]),
+    let last_update = Arc::new(Mutex::new(HashMap::new()));
-        last_update: Mutex::new(HashMap::new()),
+    let soldiers = Arc::new(Mutex::new(HashSet::new()));
-        soldiers: 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"))
+        .route("/canvas/", get(get_canvas_state))
-        .layer(layer);
+        .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(())
 }