feat: solve issues with memory explosion

Signed-off-by: Anthony Griffon <[email protected]>

Author: Miaxos

Cargo.lock

@@ -40,6 +40,8 @@ scc = "2"
 sharded-thread = "1"
 serde.workspace = true
 thiserror = "1"
+rand = "0.8"
+zstd = "0.13"
 
 # Logging
 tracing = { workspace = true, features = ["attributes"] }
@@ -63,6 +65,9 @@ tokio = { version = "1.35", features = [
   "net",
 ] }
 
+[profile.release]
+debug = true
+
 [[bench]]
 name = "parsing"
 harness = false

app/roster/Cargo.toml

@@ -127,11 +127,11 @@ impl CommandExecution for Set {
         let expired = self.expire.map(|dur| ctx.now() + dur.into());
 
         // let now = Instant::now();
-        let response = match ctx
-            .storage
-            .set_async(self.key, self.value, SetOptions { expired })
-            .await
-        {
+        let response = match ctx.storage.set_async(
+            self.key,
+            self.value,
+            SetOptions { expired },
+        ) {
             Ok(_) => Frame::Simple(OK_STR.clone()),
             Err(_) => Frame::Null,
         };

app/roster/src/application/server/cmd/set.rs

@@ -123,7 +123,9 @@ impl ReadConnection {
                 // left to `BytesMut`. This is often done by moving an internal
                 // cursor, but it may be done by reallocating and copying data.
                 // self.buffer.advance(len);
-                self.buffer.reserve(4 * 1024);
+                if self.buffer.len() < 1024 {
+                    self.buffer.reserve(4 * 1024);
+                }
 
                 // Return the parsed frame to the caller.
                 Ok(Some(frame))

app/roster/src/application/server/connection.rs

@@ -123,6 +123,7 @@ impl Handler {
 
                 let is_in_slot =
                     hash.map(|hash| ctx.is_in_slot(hash)).unwrap_or(true);
+                let is_in_slot = true;
 
                 if !is_in_slot {
                     let hash = hash.unwrap();

app/roster/src/application/server/handle.rs

@@ -34,6 +34,7 @@ impl ServerConfig {
     pub fn initialize(self) -> ServerHandle {
         let mut threads = Vec::new();
         let cpus: usize = std::thread::available_parallelism().unwrap().into();
+        // let cpus = 1;
 
         // The mesh used to pass a whole connection if needed.
         let mesh =
@@ -45,6 +46,7 @@ impl ServerConfig {
         let main_dialer = RootDialer::new(mesh, &storage);
 
         for cpu in 0..cpus {
+            dbg!(cpu);
             // TODO(@miaxos): There are some links between those two, mb we
             // should modelise it again.
             let config = self.clone();

app/roster/src/application/server/mod.rs

@@ -1,12 +1,14 @@
 //! Storage primitive which is used to interact with Keys
 
 use std::hash::BuildHasherDefault;
+use std::mem::size_of_val;
 use std::sync::atomic::AtomicU32;
 use std::sync::Arc;
 
 use bytes::Bytes;
 use bytestring::ByteString;
 use coarsetime::Instant;
+use rand::random;
 use rustc_hash::FxHasher;
 use scc::HashMap;
 
@@ -16,14 +18,14 @@ use crate::infrastructure::hash::HASH_SLOT_MAX;
 #[derive(Debug)]
 pub struct StorageValue {
     pub expired: Option<Instant>,
-    pub val: Bytes,
+    pub val: Vec<u8>,
 }
 
 /// A [StorageSegment] is shared across multiple threads and owns a part of the
 /// hashing keys.
 #[derive(Debug, Clone)]
 pub struct StorageSegment {
-    db: Arc<HashMap<ByteString, StorageValue, BuildHasherDefault<FxHasher>>>,
+    db: Arc<HashMap<Vec<u8>, StorageValue, BuildHasherDefault<FxHasher>>>,
     slot: Slot,
     count: Arc<AtomicU32>,
 }
@@ -36,15 +38,17 @@ pub struct SetOptions {
 impl StorageSegment {
     /// Create a new [StorageSegment] by specifying the hash slot it handles.
     pub fn new(slot: Slot) -> Self {
-        for _ in 0..4096 {
+        let h = HashMap::with_capacity_and_hasher(
+            2usize.pow(20),
+            Default::default(),
+        );
+
+        for _ in 0..(2usize.pow(20)) {
             drop(scc::ebr::Guard::new());
         }
 
         Self {
-            db: Arc::new(HashMap::with_capacity_and_hasher(
-                4096,
-                Default::default(),
-            )),
+            db: Arc::new(h),
             slot,
             count: Arc::new(AtomicU32::new(0)),
         }
@@ -54,36 +58,31 @@ impl StorageSegment {
         self.slot.contains(&i)
     }
 
-    /// Set a key
-    pub async fn set_async(
+    /// Set a key into the storage
+    pub fn set_async(
         &self,
         key: ByteString,
         val: Bytes,
         opt: SetOptions,
     ) -> Result<Option<StorageValue>, (String, StorageValue)> {
+        let mut val = val.to_vec();
+        val.shrink_to_fit();
+
         let val = StorageValue {
             expired: opt.expired,
             val,
         };
 
-        let _old = self
+        let mut key = key.into_bytes().to_vec();
+        key.shrink_to_fit();
+
+        let old = self
             .count
             .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-        /*
-        // Simulate some eviction mechanisme when we have too many keys
-        if old > 200_000 {
-            // dbg!("remove");
-            // TODO: If the RC is for the DB instead, we could have a spawn from
-            // monoio for this task instead, it would save us some
-            // time for the p99.9
-            let count = self.count.clone();
-            let db = self.db.clone();
-            monoio::spawn(async move {
-                db.retain_async(|_, _| false).await;
-                count.swap(0, std::sync::atomic::Ordering::Relaxed);
-            });
+
+        if old % 1_000_000 == 0 {
+            dbg!(old);
         }
-        */
 
         if let Err((key, val)) = self.db.insert(key, val) {
             let old = self.db.update(&key, |_, _| val);
@@ -101,6 +100,9 @@ impl StorageSegment {
         key: ByteString,
         now: Instant,
     ) -> Option<Bytes> {
+        let mut key = key.into_bytes().to_vec();
+        key.shrink_to_fit();
+
         match self.db.entry_async(key).await {
             scc::hash_map::Entry::Occupied(oqp) => {
                 let val = oqp.get();
@@ -109,10 +111,10 @@ impl StorageSegment {
 
                 // TODO: Better handle expiration
                 if is_expired {
-                    // oqp.remove()?;
+                    let _ = oqp.remove();
                     None
                 } else {
-                    Some(val.val.clone())
+                    Some(Bytes::from(val.val.clone()))
                 }
             }
             scc::hash_map::Entry::Vacant(_) => None,

app/roster/src/domain/storage/mod.rs

@@ -24,11 +24,13 @@ impl Cfg {
 
         let settings = Config::builder()
             .add_source(config::File::with_name(&file_location))
+            /*
             .add_source(
                 config::Environment::with_prefix("ROSTER")
                     .try_parsing(true)
                     .separator("_"),
             )
+            */
             .build()?;
 
         let config = settings.try_deserialize::<Cfg>()?;