Use TimedCache to avoid OOM (#39)

* use TimedCache to avoid OOM

* add comment

simlib/mixnet-sims/Cargo.toml

@@ -26,3 +26,4 @@ multiaddr = "0.18"
 sha2 = "0.10"
 uuid = { version = "1", features = ["fast-rng", "v4"] }
 tracing-appender = "0.2"
+cached = "0.54.0"

simlib/mixnet-sims/src/log.rs

@@ -76,7 +76,7 @@ pub fn config_tracing(
     }
 
     tracing_subscriber::registry()
-        .with(LevelFilter::from(Level::DEBUG))
+        .with(LevelFilter::from(Level::INFO))
         .with(layers)
         .init();
 

simlib/mixnet-sims/src/node/mix/mod.rs

@@ -6,6 +6,7 @@ pub mod state;
 pub mod stream_wrapper;
 
 use crate::node::mix::consensus_streams::{Epoch, Slot};
+use cached::{Cached, TimedCache};
 use crossbeam::channel;
 use futures::Stream;
 use lottery::StakeLottery;
@@ -35,7 +36,7 @@ use scheduler::{Interval, TemporalRelease};
 use serde::Deserialize;
 use sha2::{Digest, Sha256};
 use state::MixnodeState;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::{pin::pin, task::Poll, time::Duration};
 use stream_wrapper::CrossbeamReceiverStream;
 
@@ -70,7 +71,7 @@ pub struct MixNode {
     state: MixnodeState,
     settings: MixnodeSettings,
     network_interface: InMemoryNetworkInterface<MixMessage>,
-    message_cache: HashSet<Sha256Hash>,
+    message_cache: TimedCache<Sha256Hash, ()>,
 
     data_msg_lottery_update_time_sender: channel::Sender<Duration>,
     data_msg_lottery_interval: Interval,
@@ -185,7 +186,9 @@ impl MixNode {
         Self {
             id,
             network_interface,
-            message_cache: HashSet::new(),
+            // We're not coupling this lifespan with the steps now, but it's okay
+            // We expected that a message will be delivered to most of nodes within 60s.
+            message_cache: TimedCache::with_lifespan(60),
             settings,
             state: MixnodeState {
                 node_id: id,
@@ -210,7 +213,11 @@ impl MixNode {
     }
 
     fn forward(&mut self, message: MixMessage, exclude_node: Option<NodeId>) {
-        if !self.message_cache.insert(Self::sha256(&message.0)) {
+        if self
+            .message_cache
+            .cache_set(Self::sha256(&message.0), ())
+            .is_some()
+        {
             return;
         }
         for node_id in self
@@ -229,7 +236,11 @@ impl MixNode {
             .receive_messages()
             .into_iter()
             // Retain only messages that have not been seen before
-            .filter(|msg| self.message_cache.insert(Self::sha256(&msg.payload().0)))
+            .filter(|msg| {
+                self.message_cache
+                    .cache_set(Self::sha256(&msg.payload().0), ())
+                    .is_none()
+            })
             .collect()
     }