From 3469fa67a4e991e1f465115c2326197677f40baa Mon Sep 17 00:00:00 2001
From: Vecna <vvecna@uwaterloo.ca>
Date: Tue, 4 Jun 2024 06:26:58 -0400
Subject: [PATCH] Track memory usage in simulation

---
 Cargo.toml            |  3 ++-
 src/bin/simulation.rs | 45 ++++++++++++++++++++++++++++++++++++++++++-
 src/main.rs           | 23 ++++++++++++++++++++++
 3 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 22fd969..dbf3631 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -25,6 +25,7 @@ julianday = "1.2.0"
 lazy_static = "1"
 lox_cli = { path = "../lox_cli", version = "0.1", optional = true }
 lox-library = { git = "https://gitlab.torproject.org/vecna/lox.git", version = "0.1.0" }
+memory-stats = { version = "1.0.0", optional = true }
 nalgebra = "0.29"
 rand = { version = "0.8" }
 #select = "0.6.0"
@@ -45,7 +46,7 @@ base64 = "0.21.7"
 faketime = "0.2"
 
 [features]
-simulation = ["faketime", "lox_cli"]
+simulation = ["faketime", "lox_cli", "memory-stats"]
 
 [[bin]]
 name = "simulation"
diff --git a/src/bin/simulation.rs b/src/bin/simulation.rs
index 5c6e783..6965170 100644
--- a/src/bin/simulation.rs
+++ b/src/bin/simulation.rs
@@ -18,6 +18,7 @@ use troll_patrol::{
 use clap::Parser;
 use lox_cli::{networking::*, *};
 use lox_library::proto::{level_up::LEVEL_INTERVAL, trust_promotion::UNTRUSTED_INTERVAL};
+use memory_stats::memory_stats;
 use rand::{prelude::SliceRandom, Rng};
 use serde::Deserialize;
 use std::{
@@ -223,9 +224,33 @@ pub async fn main() {
     let mut true_neg = 0;
     let mut true_pos = 0;
 
+    // Track memory use during simulation
+    let mut max_physical_mem = 0;
+    let mut max_virtual_mem = 0;
+
     // Main loop
     for day in 1..=config.num_days {
         println!("Starting day {} of the simulation", day);
+        println!(
+            "We have {} users and {} bridges",
+            users.len(),
+            bridges.len()
+        );
+        println!(
+            "The censor has learned {} bridges",
+            censor.known_bridges.len()
+        );
+
+        if let Some(usage) = memory_stats() {
+            if usage.physical_mem > max_physical_mem {
+                max_physical_mem = usage.physical_mem;
+            }
+            if usage.virtual_mem > max_virtual_mem {
+                max_virtual_mem = usage.virtual_mem;
+            }
+        } else {
+            println!("Failed to get the current memory usage");
+        }
 
         // USER TASKS
 
@@ -305,7 +330,6 @@ pub async fn main() {
         let new_blockages: HashMap<String, HashSet<String>> =
             serde_json::from_slice(&new_blockages_resp).unwrap();
 
-        // TODO: Track more stats about new blockages
         // Since we have only one censor, just convert to a set of bridges
         let mut blocked_bridges = HashSet::<[u8; 20]>::new();
         for (bridge, ccs) in new_blockages {
@@ -356,11 +380,30 @@ pub async fn main() {
         increment_simulated_date();
     }
 
+    // Print various information about the simulation run
+    println!(
+        "\nSimulation ended with {} users and {} bridges",
+        users.len(),
+        bridges.len()
+    );
+    println!("The censor learned {} bridges", censor.known_bridges.len());
+
+    println!(
+        "\nMaximum physical memory usage during simulation: {}",
+        max_physical_mem
+    );
+    println!(
+        "Maximum virtual memory usage during simulation: {}\n",
+        max_virtual_mem
+    );
+
     println!("True Positives: {}", true_pos);
     println!("True Negatives: {}", true_neg);
     println!("False Positives: {}", false_pos);
     println!("False Negatives: {}", false_neg);
 
+    println!("\nFull stats per bridge:");
+
     println!(
         "Fingerprint,first_distributed,first_blocked,first_detected_blocked,first_positive_report"
     );
diff --git a/src/main.rs b/src/main.rs
index 5868a57..21ab4c5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -27,6 +27,9 @@ use tokio::{
 #[cfg(not(feature = "simulation"))]
 use tokio_cron::{Job, Scheduler};
 
+#[cfg(feature = "simulation")]
+use memory_stats::memory_stats;
+
 async fn shutdown_signal() {
     tokio::signal::ctrl_c()
         .await
@@ -151,12 +154,27 @@ async fn context_manager(
     max_historical_days: u32,
     mut context_rx: mpsc::Receiver<Command>,
 ) {
+    #[cfg(feature = "simulation")]
+    let (mut max_physical_mem, mut max_virtual_mem) = (0, 0);
+
     let db: Db = sled::open(&db_config.db_path).unwrap();
 
     // Create negative report key for today if we don't have one
     new_negative_report_key(&db, get_date());
 
     while let Some(cmd) = context_rx.recv().await {
+        #[cfg(feature = "simulation")]
+        if let Some(usage) = memory_stats() {
+            if usage.physical_mem > max_physical_mem {
+                max_physical_mem = usage.physical_mem;
+            }
+            if usage.virtual_mem > max_virtual_mem {
+                max_virtual_mem = usage.virtual_mem;
+            }
+        } else {
+            println!("Failed to get the current memory usage");
+        }
+
         use Command::*;
         match cmd {
             Request { req, sender } => {
@@ -170,6 +188,11 @@ async fn context_manager(
                 println!("Sending Shutdown Signal, all threads should shutdown.");
                 drop(shutdown_sig);
                 println!("Shutdown Sent.");
+                #[cfg(feature = "simulation")]
+                println!(
+                    "\nMaximum physical memory usage: {}\nMaximum virtual memory usage: {}\n",
+                    max_physical_mem, max_virtual_mem
+                );
             }
             Update { _req, sender } => {
                 let blockages = update_daily_info(