diff --git a/Cargo.lock b/Cargo.lock index 62adc83d..ec7a6dea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -618,6 +618,7 @@ dependencies = [ "color-backtrace", "criterion 0.4.0", "ctor 0.2.6", + "dev-utils", "flate2", "loro 0.16.2", "serde_json", diff --git a/crates/dev-utils/src/lib.rs b/crates/dev-utils/src/lib.rs index 16bed574..0927d72e 100644 --- a/crates/dev-utils/src/lib.rs +++ b/crates/dev-utils/src/lib.rs @@ -1,3 +1,4 @@ +use std::fmt::{Debug, Display}; use std::path::Path; use tracing_subscriber::fmt::format::FmtSpan; @@ -42,3 +43,62 @@ pub fn setup_test_log() { .unwrap(); } } + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::sync::atomic::{AtomicUsize, Ordering::Relaxed}; + +struct Counter; + +static ALLOCATED: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for Counter { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ret = System.alloc(layout); + if !ret.is_null() { + ALLOCATED.fetch_add(layout.size(), Relaxed); + } + ret + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + System.dealloc(ptr, layout); + ALLOCATED.fetch_sub(layout.size(), Relaxed); + } +} + +#[global_allocator] +static A: Counter = Counter; + +pub struct MemorySize { + pub bytes: usize, +} + +impl Debug for MemorySize { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let (size, unit) = match self.bytes { + bytes if bytes < 1024 => (bytes as f64, "B"), + bytes if bytes < 1024 * 1024 => (bytes as f64 / 1024.0, "KB"), + bytes if bytes < 1024 * 1024 * 1024 => (bytes as f64 / (1024.0 * 1024.0), "MB"), + bytes => (bytes as f64 / (1024.0 * 1024.0 * 1024.0), "GB"), + }; + write!(f, "{:.2} {}", size, unit) + } +} + +impl Display for MemorySize { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let (size, unit) = match self.bytes { + bytes if bytes < 1024 => (bytes as f64, "B"), + bytes if bytes < 1024 * 1024 => (bytes as f64 / 1024.0, "KB"), + bytes if bytes < 1024 * 1024 * 1024 => (bytes as f64 / (1024.0 * 1024.0), "MB"), + bytes => (bytes as f64 / (1024.0 * 1024.0 * 1024.0), "GB"), + }; + write!(f, "{:.2} {}", size, unit) + } +} + +pub fn get_allocated_bytes() -> MemorySize { + MemorySize { + bytes: ALLOCATED.load(Relaxed), + } +} diff --git a/crates/examples/Cargo.toml b/crates/examples/Cargo.toml index 58378c07..ac901899 100644 --- a/crates/examples/Cargo.toml +++ b/crates/examples/Cargo.toml @@ -8,6 +8,7 @@ publish = false [dependencies] bench-utils = { path = "../bench-utils" } +dev-utils = { path = "../dev-utils" } loro = { path = "../loro" } tabled = "0.15.0" arbitrary = { version = "1.3.0", features = ["derive"] } diff --git a/crates/examples/examples/sheet.rs b/crates/examples/examples/sheet.rs index e1266e9a..90eac82b 100644 --- a/crates/examples/examples/sheet.rs +++ b/crates/examples/examples/sheet.rs @@ -1,6 +1,16 @@ +use dev_utils::get_allocated_bytes; use examples::sheet::init_large_sheet; pub fn main() { + let allocated = get_allocated_bytes(); + println!("allocated bytes: {}", allocated); let doc = init_large_sheet(); + let allocated = get_allocated_bytes(); + println!("allocated bytes: {}", allocated); doc.export_snapshot(); + let allocated = get_allocated_bytes(); + println!("allocated bytes: {}", allocated); + drop(doc); + let allocated = get_allocated_bytes(); + println!("allocated bytes: {}", allocated); }