feat: fix memory leak & improve memory allocation (toeverything#478)

* chore: format codes

* feat: add memory leak test

* feat: add array ins del test

* feat: replace left & right to weak arc

* feat: use weak item in ItemBuilder

* feat: replace more strong arc

* feat: use weak arc in search marker

* chore: restore random crdt item test

* chore: cleanup codes

* perf(jwst-codec): reduce mem alloc times

* feat: move out `std::sync::*` to use loom

* fix: test

* fix: deps lock glitch

* fix: crate pointing in macro

---------

Co-authored-by: forehalo <[email protected]>

Author: darkskygit

Cargo.lock

@@ -37,7 +37,11 @@ impl JwstStorage {
 
         let storage = rt
             .block_on(
-                AutoStorage::new_with_migration(&format!("sqlite:{path}?mode=rwc"), BlobStorageType::DB).or_else(|e| {
+                AutoStorage::new_with_migration(
+                    &format!("sqlite:{path}?mode=rwc"),
+                    BlobStorageType::DB,
+                )
+                .or_else(|e| {
                     warn!(
                         "Failed to open storage, falling back to memory storage: {}",
                         e

libs/jwst-binding/jwst-jni/src/storage.rs

@@ -27,7 +27,11 @@ impl Storage {
 
         let storage = rt
             .block_on(
-                AutoStorage::new_with_migration(&format!("sqlite:{path}?mode=rwc"), BlobStorageType::DB).or_else(|e| {
+                AutoStorage::new_with_migration(
+                    &format!("sqlite:{path}?mode=rwc"),
+                    BlobStorageType::DB,
+                )
+                .or_else(|e| {
                     warn!(
                         "Failed to open storage, falling back to memory storage: {}",
                         e

libs/jwst-binding/jwst-swift/src/storage.rs

@@ -1,3 +1,3 @@
 pub mod doc_operation;
 
-pub use doc_operation::*;
+pub use doc_operation::*;

libs/jwst-codec-util/src/lib.rs

@@ -14,6 +14,7 @@ nanoid = "0.4.0"
 nom = "7.1.3"
 ordered-float = "3.6.0"
 rand = "0.8.5"
+rand_chacha = "0.3.1"
 serde_json = "1.0.94"
 thiserror = "1.0.40"
 
@@ -24,6 +25,9 @@ jwst-logger = { workspace = true }
 arbitrary = { version = "1.3.0", features = ["derive"] }
 ordered-float = { version = "3.6.0", features = ["arbitrary"] }
 
+[target.'cfg(loom)'.dependencies]
+loom = "0.6"
+
 [dev-dependencies]
 assert-json-diff = "2.0.2"
 criterion = { version = "0.5.1", features = ["html_reports"] }
@@ -32,12 +36,14 @@ ordered-float = { version = "3.6.0", features = ["proptest"] }
 path-ext = "0.1.0"
 proptest = "1.1.0"
 proptest-derive = "0.3.0"
-rand = "0.8.5"
-rand_chacha = "0.3.1"
 serde = { version = "1.0.155", features = ["derive"] }
 y-sync = { git = "https://github.com/toeverything/y-sync", rev = "aeb0010" }
 yrs = "0.16.5"
 
+[[bin]]
+name = "memory_leak_test"
+path = "bin/memory_leak_test.rs"
+
 [[bench]]
 name = "array_ops_benchmarks"
 harness = false

libs/jwst-codec/Cargo.toml

@@ -0,0 +1,64 @@
+use jwst_codec::*;
+use rand::{Rng, SeedableRng};
+use rand_chacha::ChaCha20Rng;
+
+fn run_text_test(seed: u64) {
+    let doc = Doc::with_client(1);
+    let mut rand = ChaCha20Rng::seed_from_u64(seed);
+    let mut text = doc.get_or_create_text("test").unwrap();
+    text.insert(0, "This is a string with length 32.").unwrap();
+
+    let iteration = 20;
+    let mut len = 32;
+
+    for i in 0..iteration {
+        let mut text = text.clone();
+        let ins = i % 2 == 0;
+        let pos = rand.gen_range(0..if ins { text.len() } else { len / 2 });
+        if ins {
+            let str = format!("hello {i}");
+            text.insert(pos, &str).unwrap();
+            len += str.len() as u64;
+        } else {
+            text.remove(pos, 6).unwrap();
+            len -= 6;
+        }
+    }
+
+    assert_eq!(text.to_string().len(), len as usize);
+    assert_eq!(text.len(), len);
+}
+
+fn run_array_test(seed: u64) {
+    let doc = Doc::with_client(1);
+    let mut rand = ChaCha20Rng::seed_from_u64(seed);
+    let mut array = doc.get_or_create_array("test").unwrap();
+    array.push(1).unwrap();
+
+    let iteration = 20;
+    let mut len = 1;
+
+    for i in 0..iteration {
+        let mut array = array.clone();
+        let ins = i % 2 == 0;
+        let pos = rand.gen_range(0..if ins { array.len() } else { len / 2 });
+        if ins {
+            array.insert(pos, 1).unwrap();
+            len += 1;
+        } else {
+            array.remove(pos, 1).unwrap();
+            len -= 1;
+        }
+    }
+
+    assert_eq!(array.len(), len);
+}
+
+fn main() {
+    let mut rand = ChaCha20Rng::seed_from_u64(rand::thread_rng().gen());
+    for _ in 0..10000 {
+        let seed = rand.gen();
+        run_array_test(seed);
+        run_text_test(seed);
+    }
+}

libs/jwst-codec/bin/memory_leak_test.rs

@@ -152,7 +152,7 @@ impl AwarenessEventBuilder {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::sync::Mutex;
+    use crate::sync::{Mutex, MutexGuard};
 
     #[test]
     fn test_awareness() {
@@ -219,7 +219,7 @@ mod tests {
             awareness.set_local_state("test".to_string());
             awareness.clear_local_state();
 
-            let values: std::sync::MutexGuard<Vec<AwarenessEvent>> = values.lock().unwrap();
+            let values: MutexGuard<Vec<AwarenessEvent>> = values.lock().unwrap();
             assert_eq!(values.len(), 4);
             let event = values.get(0).unwrap();
 

libs/jwst-codec/src/doc/awareness.rs

@@ -1,9 +1,5 @@
-use std::sync::{
-    atomic::{AtomicU8, Ordering},
-    Arc,
-};
-
 use super::*;
+use crate::sync::{Arc, AtomicU8, Ordering};
 
 #[derive(Debug, Clone)]
 #[cfg_attr(test, derive(proptest_derive::Arbitrary))]
@@ -116,22 +112,22 @@ impl ItemFlags {
 }
 
 #[derive(Clone)]
-#[cfg_attr(test, derive(proptest_derive::Arbitrary))]
+#[cfg_attr(all(test, not(loom)), derive(proptest_derive::Arbitrary))]
 pub struct Item {
     pub id: Id,
     pub origin_left_id: Option<Id>,
     pub origin_right_id: Option<Id>,
-    #[cfg_attr(test, proptest(value = "None"))]
+    #[cfg_attr(all(test, not(loom)), proptest(value = "None"))]
     pub left: Option<StructInfo>,
-    #[cfg_attr(test, proptest(value = "None"))]
+    #[cfg_attr(all(test, not(loom)), proptest(value = "None"))]
     pub right: Option<StructInfo>,
     pub parent: Option<Parent>,
     pub parent_sub: Option<String>,
     // make content Arc, so we can share the content between items
     // and item can be readonly and cloned fast.
     // TODO: considering using Cow
     pub content: Arc<Content>,
-    #[cfg_attr(test, proptest(value = "ItemFlags::default()"))]
+    #[cfg_attr(all(test, not(loom)), proptest(value = "ItemFlags::default()"))]
     pub flags: ItemFlags,
 }
 
@@ -190,6 +186,32 @@ impl Default for Item {
 }
 
 impl Item {
+    pub fn new(
+        id: Id,
+        content: Content,
+        left: Option<ItemRef>,
+        right: Option<ItemRef>,
+        parent: Option<Parent>,
+        parent_sub: Option<String>,
+    ) -> Self {
+        let flags = ItemFlags::from(if content.countable() {
+            item_flags::ITEM_COUNTABLE
+        } else {
+            0
+        });
+
+        Self {
+            id,
+            origin_left_id: left.as_ref().map(|l| l.last_id()),
+            left: left.map(|l| StructInfo::WeakItem(Arc::downgrade(&l))),
+            origin_right_id: right.as_ref().map(|r| r.id),
+            right: right.map(|r| StructInfo::WeakItem(Arc::downgrade(&r))),
+            parent,
+            parent_sub,
+            content: Arc::new(content),
+            flags,
+        }
+    }
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
@@ -235,15 +257,6 @@ impl Item {
         &mut *(Arc::as_ptr(item) as *mut Item)
     }
 
-    /// Safety:
-    /// see [inner_mut]
-    pub(crate) unsafe fn swap(one: &Arc<Item>, other: &Arc<Item>) {
-        let one = Self::inner_mut(one);
-        let other = Self::inner_mut(other);
-
-        std::mem::swap(one, other);
-    }
-
     #[allow(dead_code)]
     #[cfg(any(debug, test))]
     pub(crate) fn print_left(&self) {
@@ -259,6 +272,9 @@ impl Item {
                 StructInfo::Item(item) => {
                     ret.push(format!("{item}"));
                 }
+                StructInfo::WeakItem(item) => {
+                    ret.push(format!("{:?}", item.upgrade()));
+                }
                 StructInfo::GC { id, len } => {
                     ret.push(format!("GC{id}: {len}"));
                     break;
@@ -289,6 +305,9 @@ impl Item {
                 StructInfo::Item(item) => {
                     ret.push(format!("{item}"));
                 }
+                StructInfo::WeakItem(item) => {
+                    ret.push(format!("{:?}", item.upgrade()));
+                }
                 StructInfo::GC { id, len } => {
                     ret.push(format!("GC{id}: {len}"));
                     break;
@@ -354,7 +373,9 @@ impl Item {
                 debug_assert_ne!(first_5_bit, 10);
                 Arc::new(Content::read(decoder, first_5_bit)?)
             },
-            ..Default::default()
+            left: None,
+            right: None,
+            flags: ItemFlags::from(0),
         };
 
         if item.content.countable() {

libs/jwst-codec/src/doc/codec/item.rs

@@ -6,6 +6,7 @@ mod io;
 mod item;
 mod refs;
 mod update;
+#[cfg(any(fuzzing, test))]
 mod utils;
 
 pub use any::Any;
@@ -16,7 +17,6 @@ pub use io::{CrdtRead, CrdtReader, CrdtWrite, CrdtWriter, RawDecoder, RawEncoder
 pub use item::{item_flags, Item, ItemFlags, ItemRef, Parent};
 pub use refs::StructInfo;
 pub use update::Update;
-pub use utils::ItemBuilder;
 #[cfg(any(fuzzing, test))]
 pub use utils::*;