flame+summary profiling

Summary: Running `buck2 profile` with both `heap-summary-*` and `heap-flame-*` gathers all the same data... this provides new `heap-*` modes that produce both outputs.

Reviewed By: cjhopman

Differential Revision: D69315461

fbshipit-source-id: b875c6f624cbec36aeda3485d90e97b9c47593e7

starlark/src/eval/runtime/evaluator.rs

@@ -278,7 +278,9 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
         self.profile_or_instrumentation_mode = ProfileOrInstrumentationMode::Profile(mode.dupe());
 
         match mode {
-            ProfileMode::HeapSummaryAllocated
+            ProfileMode::HeapAllocated
+            | ProfileMode::HeapRetained
+            | ProfileMode::HeapSummaryAllocated
             | ProfileMode::HeapFlameAllocated
             | ProfileMode::HeapSummaryRetained
             | ProfileMode::HeapFlameRetained => {
@@ -291,6 +293,10 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
                     ProfileMode::HeapSummaryRetained => self
                         .module_env
                         .enable_retained_heap_profile(RetainedHeapProfileMode::Summary),
+                    ProfileMode::HeapRetained => {
+                        self.module_env
+                            .enable_retained_heap_profile(RetainedHeapProfileMode::FlameAndSummary);
+                    }
                     _ => {}
                 }
 
@@ -342,17 +348,20 @@ impl<'v, 'a, 'e: 'a> Evaluator<'v, 'a, 'e> {
         };
         self.profile_or_instrumentation_mode = ProfileOrInstrumentationMode::Collected;
         match mode {
+            ProfileMode::HeapAllocated => self
+                .heap_profile
+                .gen(self.heap(), HeapProfileFormat::FlameGraphAndSummary),
             ProfileMode::HeapSummaryAllocated => self
                 .heap_profile
                 .gen(self.heap(), HeapProfileFormat::Summary),
             ProfileMode::HeapFlameAllocated => self
                 .heap_profile
                 .gen(self.heap(), HeapProfileFormat::FlameGraph),
-            ProfileMode::HeapSummaryRetained | ProfileMode::HeapFlameRetained => {
-                Err(crate::Error::new_other(
-                    EvaluatorError::RetainedMemoryProfilingCannotBeObtainedFromEvaluator,
-                ))
-            }
+            ProfileMode::HeapSummaryRetained
+            | ProfileMode::HeapFlameRetained
+            | ProfileMode::HeapRetained => Err(crate::Error::new_other(
+                EvaluatorError::RetainedMemoryProfilingCannotBeObtainedFromEvaluator,
+            )),
             ProfileMode::Statement => self.stmt_profile.gen(),
             ProfileMode::Coverage => self.stmt_profile.gen_coverage(),
             ProfileMode::Bytecode => self.gen_bc_profile(),

starlark/src/eval/runtime/profile/data.rs

@@ -20,8 +20,10 @@ use crate::eval::runtime::profile::bc::BcPairsProfilerType;
 use crate::eval::runtime::profile::bc::BcProfileData;
 use crate::eval::runtime::profile::bc::BcProfilerType;
 use crate::eval::runtime::profile::flamegraph::FlameGraphData;
+use crate::eval::runtime::profile::heap::HeapAllocatedProfilerType;
 use crate::eval::runtime::profile::heap::HeapFlameAllocatedProfilerType;
 use crate::eval::runtime::profile::heap::HeapFlameRetainedProfilerType;
+use crate::eval::runtime::profile::heap::HeapRetainedProfilerType;
 use crate::eval::runtime::profile::heap::HeapSummaryAllocatedProfilerType;
 use crate::eval::runtime::profile::heap::HeapSummaryRetainedProfilerType;
 use crate::eval::runtime::profile::mode::ProfileMode;
@@ -46,6 +48,8 @@ enum ProfileDataError {
 pub(crate) enum ProfileDataImpl {
     Bc(Box<BcProfileData>),
     BcPairs(BcPairsProfileData),
+    HeapRetained(Box<AggregateHeapProfileInfo>),
+    HeapAllocated(Box<AggregateHeapProfileInfo>),
     HeapFlameRetained(Box<AggregateHeapProfileInfo>),
     HeapFlameAllocated(Box<AggregateHeapProfileInfo>),
     HeapSummaryRetained(Box<AggregateHeapProfileInfo>),
@@ -63,6 +67,8 @@ impl ProfileDataImpl {
         match self {
             ProfileDataImpl::Bc(_) => ProfileMode::Bytecode,
             ProfileDataImpl::BcPairs(_) => ProfileMode::BytecodePairs,
+            ProfileDataImpl::HeapRetained(_) => ProfileMode::HeapRetained,
+            ProfileDataImpl::HeapAllocated(_) => ProfileMode::HeapAllocated,
             ProfileDataImpl::HeapFlameRetained(_) => ProfileMode::HeapFlameRetained,
             ProfileDataImpl::HeapFlameAllocated(_) => ProfileMode::HeapFlameAllocated,
             ProfileDataImpl::HeapSummaryRetained(_) => ProfileMode::HeapSummaryRetained,
@@ -97,7 +103,9 @@ impl ProfileData {
     pub fn gen_flame_data(&self) -> crate::Result<String> {
         match &self.profile {
             ProfileDataImpl::TimeFlameProfile(profile) => Ok(profile.write()),
-            ProfileDataImpl::HeapFlameRetained(profile)
+            ProfileDataImpl::HeapRetained(profile)
+            | ProfileDataImpl::HeapAllocated(profile)
+            | ProfileDataImpl::HeapFlameRetained(profile)
             | ProfileDataImpl::HeapFlameAllocated(profile) => Ok(profile.gen_flame_graph_data()),
             _ => Ok("".to_owned()),
         }
@@ -108,6 +116,9 @@ impl ProfileData {
         match &self.profile {
             ProfileDataImpl::Bc(bc) => Ok(bc.gen_csv()),
             ProfileDataImpl::BcPairs(bc_pairs) => Ok(bc_pairs.gen_csv()),
+            ProfileDataImpl::HeapRetained(profile) | ProfileDataImpl::HeapAllocated(profile) => {
+                Ok(profile.gen_summary_csv())
+            }
             ProfileDataImpl::HeapSummaryRetained(profile)
             | ProfileDataImpl::HeapSummaryAllocated(profile) => Ok(profile.gen_summary_csv()),
             ProfileDataImpl::TimeFlameProfile(data) => Ok(data.write()),
@@ -145,6 +156,12 @@ impl ProfileData {
         let profile = match &profile_mode {
             ProfileMode::Bytecode => BcProfilerType::merge_profiles(&profiles)?.profile,
             ProfileMode::BytecodePairs => BcPairsProfilerType::merge_profiles(&profiles)?.profile,
+            ProfileMode::HeapAllocated => {
+                HeapAllocatedProfilerType::merge_profiles(&profiles)?.profile
+            }
+            ProfileMode::HeapRetained => {
+                HeapRetainedProfilerType::merge_profiles(&profiles)?.profile
+            }
             ProfileMode::HeapSummaryAllocated => {
                 HeapSummaryAllocatedProfilerType::merge_profiles(&profiles)?.profile
             }

starlark/src/eval/runtime/profile/golden/heap_flame_allocated.flame.golden

@@ -0,0 +1,13 @@
+# @generated
+# To regenerate, run:
+# ```
+# STARLARK_RUST_REGENERATE_GOLDEN_TESTS=1 cargo test -p starlark --lib
+# ```
+
+range 8
+module;function 16
+module;test.star.test;dict 192
+module;test.star.test;list 384
+module;test.star.test;array 480
+module;test.star.test;tuple 192
+module;test.star.test;test.star.inner;dict 1760

starlark/src/eval/runtime/profile/golden/heap_retained.flame.golden

@@ -0,0 +1,8 @@
+# @generated
+# To regenerate, run:
+# ```
+# STARLARK_RUST_REGENERATE_GOLDEN_TESTS=1 cargo test -p starlark --lib
+# ```
+
+module;function 16
+module;test.star.test;list 8

starlark/src/eval/runtime/profile/golden/heap_summary_allocated.csv.golden

@@ -0,0 +1,11 @@
+# @generated
+# To regenerate, run:
+# ```
+# STARLARK_RUST_REGENERATE_GOLDEN_TESTS=1 cargo test -p starlark --lib
+# ```
+
+Function,Time(s),TimeRec(s),Calls,Callers,TopCaller,TopCallerCount,Allocs,AllocBytes,dict,array,list,tuple,function
+"TOTALS",0.378,0.378,24,0,"",0,378,3024,244,60,48,24,2
+"test.star.test",0.168,0.308,20,1,"module",1,156,1248,24,60,48,24,0
+"test.star.inner",0.140,0.140,0,1,"test.star.test",1,220,1760,220,0,0,0,0
+"module",0.070,0.378,4,1,"(root)",1,2,16,0,0,0,0,2

starlark/src/eval/runtime/profile/golden/heap_summary_retained.csv.golden

@@ -0,0 +1,11 @@
+# @generated
+# To regenerate, run:
+# ```
+# STARLARK_RUST_REGENERATE_GOLDEN_TESTS=1 cargo test -p starlark --lib
+# ```
+
+Function,Time(s),TimeRec(s),Calls,Callers,TopCaller,TopCallerCount,Allocs,AllocBytes,function,list
+"TOTALS",0.378,0.378,24,0,"",0,3,24,2,1
+"test.star.test",0.168,0.308,20,1,"module",1,1,8,0,1
+"test.star.inner",0.140,0.140,0,1,"test.star.test",1,0,0,0,0
+"module",0.070,0.378,4,1,"(root)",1,2,16,2,0

starlark/src/eval/runtime/profile/heap.rs

@@ -28,11 +28,35 @@ use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
 use crate::values::Heap;
 use crate::values::Value;
 
+pub(crate) struct HeapAllocatedProfilerType;
+pub(crate) struct HeapRetainedProfilerType;
 pub(crate) struct HeapSummaryAllocatedProfilerType;
 pub(crate) struct HeapFlameAllocatedProfilerType;
 pub(crate) struct HeapSummaryRetainedProfilerType;
 pub(crate) struct HeapFlameRetainedProfilerType;
 
+impl ProfilerType for HeapAllocatedProfilerType {
+    type Data = Box<AggregateHeapProfileInfo>;
+    const PROFILE_MODE: ProfileMode = ProfileMode::HeapAllocated;
+
+    fn data_from_generic(profile_data: &ProfileDataImpl) -> Option<&Self::Data> {
+        match profile_data {
+            ProfileDataImpl::HeapAllocated(data) => Some(data),
+            _ => None,
+        }
+    }
+
+    fn data_to_generic(data: Self::Data) -> ProfileDataImpl {
+        ProfileDataImpl::HeapAllocated(data)
+    }
+
+    fn merge_profiles_impl(profiles: &[&Self::Data]) -> starlark_syntax::Result<Self::Data> {
+        Ok(Box::new(AggregateHeapProfileInfo::merge(
+            profiles.iter().map(|x| &***x),
+        )))
+    }
+}
+
 impl ProfilerType for HeapSummaryAllocatedProfilerType {
     type Data = Box<AggregateHeapProfileInfo>;
     const PROFILE_MODE: ProfileMode = ProfileMode::HeapSummaryAllocated;
@@ -77,6 +101,28 @@ impl ProfilerType for HeapFlameAllocatedProfilerType {
     }
 }
 
+impl ProfilerType for HeapRetainedProfilerType {
+    type Data = Box<AggregateHeapProfileInfo>;
+    const PROFILE_MODE: ProfileMode = ProfileMode::HeapRetained;
+
+    fn data_from_generic(profile_data: &ProfileDataImpl) -> Option<&Self::Data> {
+        match profile_data {
+            ProfileDataImpl::HeapRetained(data) => Some(data),
+            _ => None,
+        }
+    }
+
+    fn data_to_generic(data: Self::Data) -> ProfileDataImpl {
+        ProfileDataImpl::HeapRetained(data)
+    }
+
+    fn merge_profiles_impl(profiles: &[&Self::Data]) -> starlark_syntax::Result<Self::Data> {
+        Ok(Box::new(AggregateHeapProfileInfo::merge(
+            profiles.iter().map(|x| &***x),
+        )))
+    }
+}
+
 impl ProfilerType for HeapSummaryRetainedProfilerType {
     type Data = Box<AggregateHeapProfileInfo>;
     const PROFILE_MODE: ProfileMode = ProfileMode::HeapSummaryRetained;
@@ -125,6 +171,7 @@ impl ProfilerType for HeapFlameRetainedProfilerType {
 pub(crate) enum RetainedHeapProfileMode {
     Flame,
     Summary,
+    FlameAndSummary,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -135,8 +182,9 @@ enum HeapProfileError {
 
 #[derive(Copy, Clone, Dupe, Debug)]
 pub(crate) enum HeapProfileFormat {
-    Summary,
     FlameGraph,
+    Summary,
+    FlameGraphAndSummary,
 }
 
 pub(crate) struct HeapProfile {
@@ -178,6 +226,9 @@ impl HeapProfile {
 
     pub(crate) fn gen_enabled(heap: &Heap, format: HeapProfileFormat) -> ProfileData {
         match format {
+            HeapProfileFormat::FlameGraphAndSummary => {
+                Self::write_flame_and_summarized_heap_profile(heap)
+            }
             HeapProfileFormat::Summary => Self::write_summarized_heap_profile(heap),
             HeapProfileFormat::FlameGraph => Self::write_flame_heap_profile(heap),
         }
@@ -196,6 +247,13 @@ impl HeapProfile {
             profile: ProfileDataImpl::HeapSummaryAllocated(Box::new(stacks)),
         }
     }
+
+    fn write_flame_and_summarized_heap_profile(heap: &Heap) -> ProfileData {
+        let stacks = AggregateHeapProfileInfo::collect(heap, None);
+        ProfileData {
+            profile: ProfileDataImpl::HeapAllocated(Box::new(stacks)),
+        }
+    }
 }
 
 #[cfg(test)]

starlark/src/eval/runtime/profile/mode.rs

@@ -36,6 +36,10 @@ pub enum ProfileMode {
     HeapFlameAllocated,
     /// Like heap flame, but information about retained memory after module is frozen.
     HeapFlameRetained,
+    /// HeapSummaryAllocated+HeapFlameAllocated
+    HeapAllocated,
+    /// HeapSummaryRetained+HeapFlameRetained
+    HeapRetained,
     /// The statement profile mode provides information about time spent in each statement.
     Statement,
     /// Code coverage.
@@ -60,11 +64,13 @@ impl Display for ProfileMode {
 }
 
 impl ProfileMode {
-    pub(crate) const ALL: [ProfileMode; 11] = [
+    pub(crate) const ALL: [ProfileMode; 13] = [
         ProfileMode::HeapSummaryAllocated,
         ProfileMode::HeapSummaryRetained,
         ProfileMode::HeapFlameAllocated,
         ProfileMode::HeapFlameRetained,
+        ProfileMode::HeapAllocated,
+        ProfileMode::HeapRetained,
         ProfileMode::Statement,
         ProfileMode::Coverage,
         ProfileMode::Bytecode,
@@ -80,6 +86,8 @@ impl ProfileMode {
             ProfileMode::HeapSummaryRetained => "heap-summary-retained",
             ProfileMode::HeapFlameAllocated => "heap-flame-allocated",
             ProfileMode::HeapFlameRetained => "heap-flame-retained",
+            ProfileMode::HeapAllocated => "heap-allocated",
+            ProfileMode::HeapRetained => "heap-retained",
             ProfileMode::Statement => "statement",
             ProfileMode::Coverage => "coverage",
             ProfileMode::Bytecode => "bytecode",
@@ -94,7 +102,9 @@ impl ProfileMode {
     /// [`FrozenModule::heap_profile`](crate::environment::FrozenModule::heap_profile).
     pub fn requires_frozen_module(&self) -> bool {
         match self {
-            ProfileMode::HeapSummaryRetained | ProfileMode::HeapFlameRetained => true,
+            ProfileMode::HeapSummaryRetained
+            | ProfileMode::HeapFlameRetained
+            | ProfileMode::HeapRetained => true,
             _ => false,
         }
     }

starlark/src/eval/runtime/profile/tests.rs

@@ -74,31 +74,61 @@ R = test()
         _ => eval.gen_profile().unwrap(),
     };
 
-    if let ProfileDataImpl::HeapFlameRetained(profile)
+    if let ProfileDataImpl::HeapRetained(profile)
+    | ProfileDataImpl::HeapAllocated(profile)
+    | ProfileDataImpl::HeapFlameRetained(profile)
     | ProfileDataImpl::HeapFlameAllocated(profile)
     | ProfileDataImpl::HeapSummaryRetained(profile)
     | ProfileDataImpl::HeapSummaryAllocated(profile) = &mut profile_data.profile
     {
         profile.normalize_for_golden_tests();
     }
 
-    golden_test_template(
-        &format!(
-            "src/eval/runtime/profile/golden/{}.golden",
-            mode.name().replace('-', "_")
-        ),
-        &match mode {
-            ProfileMode::HeapFlameAllocated | ProfileMode::HeapFlameRetained => {
-                profile_data.gen_flame_data().unwrap()
-            }
-            _ => profile_data.gen_csv().unwrap(),
-        },
-    );
+    match mode {
+        ProfileMode::HeapRetained
+        | ProfileMode::HeapFlameRetained
+        | ProfileMode::HeapAllocated
+        | ProfileMode::HeapFlameAllocated
+        | ProfileMode::TimeFlame => {
+            golden_test_template(
+                &format!(
+                    "src/eval/runtime/profile/golden/{}.flame.golden",
+                    mode.name().replace('-', "_")
+                ),
+                &profile_data.gen_flame_data().unwrap(),
+            );
+        }
+        _ => {}
+    }
 
+    match mode {
+        ProfileMode::HeapFlameRetained
+        | ProfileMode::HeapFlameAllocated
+        | ProfileMode::TimeFlame => {}
+        _ => {
+            golden_test_template(
+                &format!(
+                    "src/eval/runtime/profile/golden/{}.csv.golden",
+                    mode.name().replace('-', "_")
+                ),
+                &profile_data.gen_csv().unwrap(),
+            );
+        }
+    }
     // Smoke test for profile merging.
     ProfileData::merge([&profile_data, &profile_data]).unwrap();
 }
 
+#[test]
+fn test_profile_golden_heap_allocated() {
+    test_profile_golden_for_mode(ProfileMode::HeapAllocated);
+}
+
+#[test]
+fn test_profile_golden_heap_retained() {
+    test_profile_golden_for_mode(ProfileMode::HeapAllocated);
+}
+
 #[test]
 fn test_profile_golden_heap_summary_allocated() {
     test_profile_golden_for_mode(ProfileMode::HeapSummaryAllocated);

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

@@ -387,6 +387,9 @@ impl RetainedHeapProfile {
     pub(crate) fn to_profile(&self) -> ProfileData {
         ProfileData {
             profile: match self.mode {
+                RetainedHeapProfileMode::FlameAndSummary => {
+                    ProfileDataImpl::HeapRetained(Box::new(self.info.clone()))
+                }
                 RetainedHeapProfileMode::Flame => {
                     ProfileDataImpl::HeapFlameRetained(Box::new(self.info.clone()))
                 }