How do I log WASM heap memory usage from Rust?

Question

I want to detect memory leaks by periodically logging the size of the used portion of the WASM heap. What's the easiest way of doing that?

I thought the "Rust and WebAssembly" book had some advice about this, but I can't find it.

Answer 1

As explained in the comments, you can get via JavaScript the total memory usage in your WASM (WebAssembly.Memory.prototype.buffer.byteLength). This never shrinks, but if it continuously grows then you probably have a leak. You can get the WebAssembly.Memory instance via wasm_bindgen::memory(), and the rest can be done with wasm_bindgen:

fn get_current_allocated_bytes() -> u64 {
    #[wasm_bindgen]
    extern "C" {
        type Memory;
        #[wasm_bindgen(method, getter)]
        fn buffer(this: &Memory) -> MaybeSharedArrayBuffer;

        type MaybeSharedArrayBuffer;
        #[wasm_bindgen(method, getter = byteLength)]
        fn byte_length(this: &MaybeSharedArrayBuffer) -> f64;
    }

    wasm_bindgen::memory()
        .unchecked_into::<Memory>()
        .buffer()
        .byte_length() as u64
}

If you want a more performant implementation (this will be quite slow), or if you want a more precise metric (as said, this won't count deallocations), you can implement a global allocator:

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicIsize, Ordering};

struct CountingAllocator<A> {
    inner: A,
    allocated_now: AtomicIsize,
}

impl<A> CountingAllocator<A> {
    const fn new(inner: A) -> Self {
        Self {
            inner,
            allocated_now: AtomicIsize::new(0),
        }
    }

    fn allocated_now(&self) -> usize {
        self.allocated_now
            .load(Ordering::Relaxed)
            .try_into()
            .unwrap_or(0)
    }
}

unsafe impl<A: GlobalAlloc> GlobalAlloc for CountingAllocator<A> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.allocated_now
            .fetch_add(layout.size() as isize, Ordering::Relaxed);
        self.inner.alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.allocated_now
            .fetch_sub(layout.size() as isize, Ordering::Relaxed);
        self.inner.dealloc(ptr, layout);
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        self.allocated_now
            .fetch_add(layout.size() as isize, Ordering::Relaxed);
        self.inner.alloc_zeroed(layout)
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        self.allocated_now.fetch_add(
            new_size as isize - layout.size() as isize,
            Ordering::Relaxed,
        );
        self.inner.realloc(ptr, layout, new_size)
    }
}

#[global_allocator]
static ALLOCATOR: CountingAllocator<System> = CountingAllocator::new(System);

Then call ALLOCATOR.allocated_now() to retrieve the exact number of currently allocated bytes.