AggregateHeapProfileInfo::merge

Summary: Utility to merge multiple `AggregateHeapProfileInfo`.

Reviewed By: bobyangyf

Differential Revision: D38518064

fbshipit-source-id: dca83ccedf08029f461da6b1a812cb5a7cd04ba7

Author: stepancheg

starlark/src/eval/runtime/small_duration.rs

@@ -92,3 +92,12 @@ impl<'a> Sum<&'a SmallDuration> for SmallDuration {
         iter.fold(SmallDuration::default(), |acc, x| acc + *x)
     }
 }
+
+impl Sum<SmallDuration> for SmallDuration {
+    fn sum<I>(iter: I) -> SmallDuration
+    where
+        I: Iterator<Item = SmallDuration>,
+    {
+        iter.fold(SmallDuration::default(), |acc, x| acc + x)
+    }
+}

starlark/src/values/layout/heap/profile/aggregated.rs

@@ -209,6 +209,47 @@ pub(crate) struct StackFrame {
     pub(crate) calls_x2: u32,
 }
 
+impl StackFrame {
+    fn merge_callees<'a>(
+        frames: &'a [StackFrameWithContext<'a>],
+        strings: &mut StringIndex,
+    ) -> SmallMap<StringId, StackFrame> {
+        let mut group_by_callee: SmallMap<&str, Vec<StackFrameWithContext>> = SmallMap::new();
+        for frame in frames {
+            for (name, callee) in frame.callees() {
+                group_by_callee
+                    .entry(name)
+                    .or_insert_with(Vec::new)
+                    .push(callee);
+            }
+        }
+        group_by_callee
+            .into_iter()
+            .map(|(name, frames)| {
+                let name = strings.index(name);
+                (name, StackFrame::merge(frames, strings))
+            })
+            .collect()
+    }
+
+    fn merge<'a>(
+        frames: impl IntoIterator<Item = StackFrameWithContext<'a>>,
+        strings: &mut StringIndex,
+    ) -> StackFrame {
+        let frames = Vec::from_iter(frames);
+        let callees = StackFrame::merge_callees(&frames, strings);
+        let allocs = HeapSummary::merge(frames.iter().map(|f| &f.frame.allocs));
+        let time_x2 = frames.iter().map(|f| f.frame.time_x2).sum();
+        let calls_x2 = frames.iter().map(|f| f.frame.calls_x2).sum();
+        StackFrame {
+            callees,
+            allocs,
+            time_x2,
+            calls_x2,
+        }
+    }
+}
+
 struct StackFrameWithContext<'c> {
     frame: &'c StackFrame,
     strings: &'c StringIndex,
@@ -292,6 +333,25 @@ impl AggregateHeapProfileInfo {
         }
     }
 
+    #[allow(dead_code)] // TODO: used later.
+    fn merge<'a>(
+        profiles: impl IntoIterator<Item = &'a AggregateHeapProfileInfo>,
+    ) -> AggregateHeapProfileInfo {
+        let mut strings = StringIndex::default();
+        let totals_id = strings.index(Self::TOTALS_STR);
+        let root_id = strings.index(Self::ROOT_STR);
+        let blank_id = strings.index(Self::BLANK_STR);
+        let roots = profiles.into_iter().map(|p| p.root());
+        let root = StackFrame::merge(roots, &mut strings);
+        AggregateHeapProfileInfo {
+            strings,
+            root,
+            totals_id,
+            root_id,
+            blank_id,
+        }
+    }
+
     /// Write this out recursively to a file.
     pub(crate) fn gen_flame_graph(&self) -> String {
         let mut writer = FlameGraphWriter::new();
@@ -312,6 +372,7 @@ mod tests {
     use crate::values::layout::heap::heap_type::HeapKind;
     use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
     use crate::values::layout::heap::profile::aggregated::StackFrame;
+    use crate::values::layout::heap::profile::summary::HeapSummaryByFunction;
     use crate::values::Freezer;
     use crate::values::FrozenHeap;
     use crate::values::Heap;
@@ -372,4 +433,25 @@ mod tests {
         );
         assert_eq!(2, total_alloc_count(&stacks.root));
     }
+
+    #[test]
+    fn test_merge() {
+        fn make() -> AggregateHeapProfileInfo {
+            let heap = Heap::new();
+            heap.record_call_enter(const_frozen_string!("xx").to_value());
+            let s = heap.alloc_str("abc");
+            heap.record_call_exit();
+            let freezer = Freezer::new(FrozenHeap::new());
+            freezer.freeze(s.to_value()).unwrap();
+
+            AggregateHeapProfileInfo::collect(&heap, Some(HeapKind::Frozen))
+        }
+
+        let merge = AggregateHeapProfileInfo::merge([&make(), &make(), &make()]);
+        let summary = HeapSummaryByFunction::init(&merge);
+        assert_eq!(1, summary.info().len());
+        let (xx_id, xx_info) = summary.info()[0];
+        assert_eq!("xx", merge.strings.get(xx_id));
+        assert_eq!(3, xx_info.alloc.get("string").unwrap().count);
+    }
 }

starlark/src/values/layout/heap/profile/by_type.rs

@@ -53,4 +53,14 @@ impl HeapSummary {
     pub(crate) fn add(&mut self, t: &'static str, s: AllocCounts) {
         *self.summary.entry(t).or_default() += s;
     }
+
+    pub(crate) fn merge<'a>(heaps: impl IntoIterator<Item = &'a HeapSummary>) -> HeapSummary {
+        let mut summary = SmallMap::new();
+        for heap in heaps {
+            for (k, v) in heap.summary.iter() {
+                *summary.entry(*k).or_default() += *v;
+            }
+        }
+        HeapSummary { summary }
+    }
 }

starlark/src/values/layout/heap/profile/summary.rs

@@ -29,7 +29,7 @@ use crate::values::layout::heap::profile::string_index::StringIndexMap;
 
 /// Information relating to a function.
 #[derive(Default, Debug, Clone)]
-struct FuncInfo {
+pub(crate) struct FuncInfo {
     /// Number of times this function was called
     pub calls: usize,
     /// Who called this function (and how many times each)
@@ -108,7 +108,7 @@ impl HeapSummaryByFunction {
         FuncInfo::merge(self.info.values())
     }
 
-    fn info(&self) -> Vec<(StringId, &FuncInfo)> {
+    pub(crate) fn info(&self) -> Vec<(StringId, &FuncInfo)> {
         self.info.iter().collect::<Vec<_>>()
     }