Add extra_memory function

Summary:
For profiling, we can observe how much memory a Value stores on the Starlark heap, but it's very difficult to see what memory the Starlark heap points at. Add an extra function to give more accurate memory profiling. This can never be perfect, think of things like Arc which are shared, but it's an improvement over the status quo.

Next steps will be to integrate this into heap profiling modes.

Reviewed By: krallin

Differential Revision: D30866799

fbshipit-source-id: 99aa514e29de441422fa495ed27ba0a88a021fea

Author: ndmitchell

starlark/src/collections/small_map.rs

@@ -569,6 +569,21 @@ impl<K, V> SmallMap<K, V> {
         }
     }
 
+    /// Give a best guess as to how much heap memory is being used.
+    /// Used internally, but not exported as this isn't a usual API.
+    pub(crate) fn extra_memory(&self) -> usize {
+        match &self.state {
+            MapHolder::Empty => 0,
+            MapHolder::Vec(x) => x.capacity() * mem::size_of::<(K, V)>(),
+            MapHolder::Map(x) => {
+                // For each entry, IndexMap stores the (K,V), a usize in the hash table, and about 1 byte overhead.
+                // Not canonical details, just based on the current implementation.
+                let factor = mem::size_of::<(K, V)>() + mem::size_of::<usize>() + 1;
+                factor * x.capacity()
+            }
+        }
+    }
+
     fn upgrade_empty_to_vec(&mut self) -> &mut VecMap<K, V> {
         self.state = MapHolder::Vec(VecMap::default());
         if let MapHolder::Vec(ref mut v) = self.state {

starlark/src/values/layout/avalue.rs

@@ -331,6 +331,9 @@ impl<'v, Mode: 'static, T: StarlarkValue<'v>> StarlarkValue<'v> for Wrapper<Mode
     fn compare(&self, other: Value<'v>) -> anyhow::Result<Ordering> {
         self.1.compare(other)
     }
+    fn extra_memory(&self) -> usize {
+        self.1.extra_memory()
+    }
     fn invoke(
         &self,
         me: Value<'v>,

starlark/src/values/traits.rs

@@ -444,6 +444,13 @@ pub trait StarlarkValue<'v>: 'v + AnyLifetime<'v> + AsStarlarkValue<'v> + Debug
         }
     }
 
+    /// Return how much extra memory is consumed by this data type, in bytes, in addition to the
+    /// direct `size_of` measurements. Used for profiling, so best effort rather than precise. Defaults to 0.
+    /// Should not reported any memory held on to by a [`Value`].
+    fn extra_memory(&self) -> usize {
+        0
+    }
+
     /// Compare `self` with `other` for equality.
     /// Should only return an error on excessive recursion.
     fn equals(&self, _other: Value<'v>) -> anyhow::Result<bool> {

starlark/src/values/types/dict.rs

@@ -320,6 +320,10 @@ where
         }
     }
 
+    fn extra_memory(&self) -> usize {
+        self.0.content().extra_memory()
+    }
+
     fn length(&self) -> anyhow::Result<i32> {
         Ok(self.0.content().len() as i32)
     }

starlark/src/values/types/enumeration.rs

@@ -221,6 +221,11 @@ where
         self.constructor.invoke(location, args, eval)
     }
 
+    fn extra_memory(&self) -> usize {
+        let typ = self.typ.as_aref();
+        typ.as_ref().map_or(0, |s| s.capacity()) + self.elements.extra_memory()
+    }
+
     fn length(&self) -> anyhow::Result<i32> {
         Ok(self.elements.len() as i32)
     }

starlark/src/values/types/function.rs

@@ -165,6 +165,10 @@ impl<'v> StarlarkValue<'v> for NativeFunction {
         })
     }
 
+    fn extra_memory(&self) -> usize {
+        self.name.capacity()
+    }
+
     fn get_attr(&self, attribute: &str, _heap: &'v Heap) -> Option<Value<'v>> {
         if let Some(s) = &self.typ {
             if attribute == "type" {

starlark/src/values/types/list.rs

@@ -43,6 +43,7 @@ use std::{
     fmt::Debug,
     intrinsics::unlikely,
     marker::PhantomData,
+    mem,
     ops::Deref,
 };
 
@@ -221,13 +222,13 @@ impl FrozenList {
 }
 
 trait ListLike<'v>: Debug {
-    fn content(&self) -> ARef<[Value<'v>]>;
+    fn content(&self) -> ARef<Vec<Value<'v>>>;
     fn set_at(&self, i: usize, v: Value<'v>) -> anyhow::Result<()>;
 }
 
 impl<'v> ListLike<'v> for RefCell<List<'v>> {
-    fn content(&self) -> ARef<[Value<'v>]> {
-        ARef::new_ref(Ref::map(self.borrow(), |x| x.content.as_slice()))
+    fn content(&self) -> ARef<Vec<Value<'v>>> {
+        ARef::new_ref(Ref::map(self.borrow(), |x| &x.content))
     }
 
     fn set_at(&self, i: usize, v: Value<'v>) -> anyhow::Result<()> {
@@ -242,8 +243,8 @@ impl<'v> ListLike<'v> for RefCell<List<'v>> {
 }
 
 impl<'v> ListLike<'v> for FrozenList {
-    fn content(&self) -> ARef<[Value<'v>]> {
-        ARef::new_ptr(coerce_ref(&self.content).as_slice())
+    fn content(&self) -> ARef<Vec<Value<'v>>> {
+        ARef::new_ptr(coerce_ref(&self.content))
     }
 
     fn set_at(&self, _i: usize, _v: Value<'v>) -> anyhow::Result<()> {
@@ -313,6 +314,10 @@ where
         Ok(self.0.content()[i])
     }
 
+    fn extra_memory(&self) -> usize {
+        self.0.content().capacity() * mem::size_of::<Value>()
+    }
+
     fn length(&self) -> anyhow::Result<i32> {
         Ok(self.0.content().len() as i32)
     }

starlark/src/values/types/record.rs

@@ -294,6 +294,13 @@ where
         self.constructor.invoke(location, args, eval)
     }
 
+    fn extra_memory(&self) -> usize {
+        // We don't capture the memory beneath the TypeCompiled, since we don't know how big
+        // those closures are.
+        let typ = self.typ.as_aref();
+        typ.as_ref().map_or(0, |s| s.capacity()) + self.fields.extra_memory()
+    }
+
     fn dir_attr(&self) -> Vec<String> {
         vec!["type".to_owned()]
     }

starlark/src/values/types/string.rs

@@ -293,6 +293,10 @@ impl<'v> StarlarkValue<'v> for StarlarkStr {
         }
     }
 
+    fn extra_memory(&self) -> usize {
+        self.len
+    }
+
     fn length(&self) -> anyhow::Result<i32> {
         Ok(fast_string::len(self.unpack()) as i32)
     }

starlark/src/values/types/structs.rs

@@ -124,6 +124,10 @@ where
         RES.methods(crate::stdlib::structs::struct_methods)
     }
 
+    fn extra_memory(&self) -> usize {
+        self.fields.extra_memory()
+    }
+
     fn to_json(&self) -> anyhow::Result<String> {
         let mut s = "{".to_owned();
         s += &self