Эх сурвалжийг харах

bin/ircd: use BtreeMap for unread_msgs

ghassmo 3 жил өмнө
parent
commit
f680467ebc

+ 24 - 5
bin/ircd/src/buffers.rs

@@ -1,8 +1,10 @@
 use async_std::sync::{Arc, Mutex};
-use std::{cmp::Ordering, collections::VecDeque};
+use std::{
+    cmp::Ordering,
+    collections::{BTreeMap, VecDeque},
+};
 
 use chrono::Utc;
-use fxhash::FxHashMap;
 use ripemd::{Digest, Ripemd160};
 
 use crate::Privmsg;
@@ -31,20 +33,33 @@ pub fn create_buffers() -> Buffers {
 }
 
 #[derive(Clone)]
-pub struct UMsgs(pub FxHashMap<String, Privmsg>);
+pub struct UMsgs {
+    pub msgs: BTreeMap<String, Privmsg>,
+    capacity: usize,
+}
 
 impl UMsgs {
     pub fn new() -> Self {
-        Self(FxHashMap::default())
+        Self { msgs: BTreeMap::new(), capacity: SIZE_OF_MSGS_BUFFER }
     }
 
     pub fn insert(&mut self, msg: &Privmsg) -> String {
         let mut hasher = Ripemd160::new();
         hasher.update(msg.to_string());
         let key = hex::encode(hasher.finalize());
-        self.0.insert(key.clone(), msg.clone());
+
+        if self.msgs.len() == self.capacity {
+            self.pop_front();
+        }
+
+        self.msgs.insert(key.clone(), msg.clone());
         key
     }
+
+    fn pop_front(&mut self) {
+        let first_key = self.msgs.iter().next_back().unwrap().0.clone();
+        self.msgs.remove(&first_key);
+    }
 }
 
 impl Default for UMsgs {
@@ -126,6 +141,10 @@ impl PrivmsgsBuffer {
         self.buffer.iter()
     }
 
+    pub fn len(&self) -> usize {
+        self.buffer.len()
+    }
+
     pub fn last_term(&self) -> u64 {
         match self.buffer.len() {
             0 => 0,

+ 4 - 4
bin/ircd/src/protocol_privmsg.rs

@@ -113,7 +113,7 @@ impl ProtocolPrivmsg {
 
             let mut inv_requested = vec![];
             for inv_object in inv.invs.iter() {
-                let msgs = &mut self.buffers.unread_msgs.lock().await.0;
+                let msgs = &mut self.buffers.unread_msgs.lock().await.msgs;
                 if let Some(msg) = msgs.get_mut(&inv_object.0) {
                     msg.read_confirms += 1;
                 } else {
@@ -164,7 +164,7 @@ impl ProtocolPrivmsg {
             let getdata = self.getdata_sub.receive().await?;
             let getdata = (*getdata).to_owned();
 
-            let msgs = &self.buffers.unread_msgs.lock().await.0;
+            let msgs = &self.buffers.unread_msgs.lock().await.msgs;
             for inv in getdata.invs {
                 if let Some(msg) = msgs.get(&inv.0) {
                     self.channel.send(msg.clone()).await?;
@@ -178,7 +178,7 @@ impl ProtocolPrivmsg {
     }
 
     async fn update_unread_msgs(&self) -> Result<()> {
-        let msgs = &mut self.buffers.unread_msgs.lock().await.0;
+        let msgs = &mut self.buffers.unread_msgs.lock().await.msgs;
         for (hash, msg) in msgs.clone() {
             if msg.timestamp + UNREAD_MSG_EXPIRE_TIME < Utc::now().timestamp() {
                 msgs.remove(&hash);
@@ -203,7 +203,7 @@ impl ProtocolPrivmsg {
 
         self.update_unread_msgs().await?;
 
-        for msg in self.buffers.unread_msgs.lock().await.0.values() {
+        for msg in self.buffers.unread_msgs.lock().await.msgs.values() {
             self.channel.send(msg.clone()).await?;
         }
         Ok(())