Search code examples
memoryrusthashmap

How to get memory occupied by an object in rust from the code itself


I have a requirement to read the memory occupied by an object based on some events. Is there's a simple method that I can use from the standard library ? I need to access the memory usage of the following object which contains nested structs.

HashMap<String, HashMap<AppIdentifier, HashMap<String, u64>>>

Solution

  • I recommend you create a trait with a method that returns the size of self's transitively-owned allocations:

    trait AllocatedSize {
        fn allocated_size(&self) -> usize;
    }
    

    And then implement that for everything involved:

    impl AllocatedSize for u64 {
        fn allocated_size(&self) -> usize {
            0
        }
    }
    
    impl AllocatedSize for String {
        fn allocated_size(&self) -> usize {
            self.capacity()
        }
    }
    
    impl AllocatedSize for AppIdentifier {
        fn allocated_size(&self) -> usize {
            todo!()
        }
    }
    
    impl<K: AllocatedSize, V: AllocatedSize> AllocatedSize for HashMap<K, V> {
        fn allocated_size(&self) -> usize {
            // every element in the map directly owns its key and its value
            const ELEMENT_SIZE: usize = std::mem::size_of::<K>() + std::mem::size_of::<V>();
    
            // directly owned allocation
            // NB: self.capacity() may be an underestimate, see its docs
            // NB: also ignores control bytes, see hashbrown implementation
            let directly_owned = self.capacity() * ELEMENT_SIZE;
    
            // transitively owned allocations
            let transitively_owned = self
                .iter()
                .map(|(key, val)| key.allocated_size() + val.allocated_size())
                .sum();
    
            directly_owned + transitively_owned
        }
    }