Cleanup profiling code

Reviewed By: milend

Differential Revision: D38680952

fbshipit-source-id: e7bbbb45ca3158780ce0dde48579d26c6b02d10c

Author: stepancheg

starlark/src/eval/runtime/evaluator.rs

@@ -80,18 +80,10 @@ pub(crate) enum EvaluatorError {
     InstrumentationEnabled,
     #[error("Profile data already collected")]
     ProfileDataAlreadyCollected,
-    #[error("Can't call `write_heap_profile` unless you first call `enable_heap_profile`.")]
-    HeapProfilingNotEnabled,
-    #[error("Can't call `write_stmt_profile` unless you first call `enable_stmt_profile`.")]
-    StmtProfilingNotEnabled,
-    #[error("Can't call `write_flame_profile` unless you first call `enable_flame_profile`.")]
-    FlameProfilingNotEnabled,
     #[error("Retained memory profiling can be only obtained from `FrozenModule`")]
     RetainedMemoryProfilingCannotBeObtainedFromEvaluator,
     #[error("Can't call `write_bc_profile` unless you first call `enable_bc_profile`.")]
     BcProfilingNotEnabled,
-    #[error("Typecheck profiling not enabled")]
-    TypecheckProfilingNotEnabled,
     #[error("Profile or instrumentation already enabled")]
     ProfileOrInstrumentationAlreadyEnabled,
     #[error("Top frame is not def (internal error)")]
@@ -324,31 +316,20 @@ impl<'v, 'a> Evaluator<'v, 'a> {
         match mode {
             ProfileMode::HeapSummaryAllocated => self
                 .heap_profile
-                .gen(self.heap(), HeapProfileFormat::Summary)
-                .ok_or_else(|| EvaluatorError::HeapProfilingNotEnabled.into()),
+                .gen(self.heap(), HeapProfileFormat::Summary),
             ProfileMode::HeapFlameAllocated => self
                 .heap_profile
-                .gen(self.heap(), HeapProfileFormat::FlameGraph)
-                .ok_or_else(|| EvaluatorError::HeapProfilingNotEnabled.into()),
+                .gen(self.heap(), HeapProfileFormat::FlameGraph),
             ProfileMode::HeapSummaryRetained | ProfileMode::HeapFlameRetained => {
                 Err(EvaluatorError::RetainedMemoryProfilingCannotBeObtainedFromEvaluator.into())
             }
-            ProfileMode::Statement => self
-                .stmt_profile
-                .gen()
-                .ok_or_else(|| EvaluatorError::StmtProfilingNotEnabled.into()),
+            ProfileMode::Statement => self.stmt_profile.gen(),
             ProfileMode::Bytecode | ProfileMode::BytecodePairs => self
                 .bc_profile
                 .gen_csv()
                 .map(|csv| ProfileData::new(mode, csv)),
-            ProfileMode::TimeFlame => self
-                .flame_profile
-                .gen()
-                .ok_or_else(|| EvaluatorError::FlameProfilingNotEnabled.into()),
-            ProfileMode::Typecheck => self
-                .typecheck_profile
-                .gen()
-                .ok_or_else(|| EvaluatorError::TypecheckProfilingNotEnabled.into()),
+            ProfileMode::TimeFlame => self.flame_profile.gen(),
+            ProfileMode::Typecheck => self.typecheck_profile.gen(),
         }
     }
 

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

@@ -32,6 +32,12 @@ use crate::values::Trace;
 use crate::values::Tracer;
 use crate::values::Value;
 
+#[derive(Debug, thiserror::Error)]
+enum FlameProfileError {
+    #[error("Flame profile not enabled")]
+    NotEnabled,
+}
+
 /// Index into FlameData.values
 #[derive(Hash, PartialEq, Eq, Clone, Copy, Dupe)]
 struct ValueIndex(usize);
@@ -169,8 +175,11 @@ impl<'v> FlameProfile<'v> {
     }
 
     // We could expose profile on the Heap, but it's an implementation detail that it works here.
-    pub(crate) fn gen(&self) -> Option<ProfileData> {
-        self.0.as_ref().map(|box x| Self::gen_profile(x))
+    pub(crate) fn gen(&self) -> anyhow::Result<ProfileData> {
+        match &self.0 {
+            None => Err(FlameProfileError::NotEnabled.into()),
+            Some(x) => Ok(Self::gen_profile(x)),
+        }
     }
 
     fn gen_profile(x: &FlameData) -> ProfileData {

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

@@ -25,6 +25,12 @@ use crate::values::layout::heap::profile::aggregated::AggregateHeapProfileInfo;
 use crate::values::Heap;
 use crate::values::Value;
 
+#[derive(Debug, thiserror::Error)]
+enum HeapProfileError {
+    #[error("heap profile not enabled")]
+    NotEnabled,
+}
+
 #[derive(Copy, Clone, Dupe, Debug)]
 pub(crate) enum HeapProfileFormat {
     Summary,
@@ -61,12 +67,15 @@ impl HeapProfile {
     }
 
     // We could expose profile on the Heap, but it's an implementation detail that it works here.
-    pub(crate) fn gen(&self, heap: &Heap, format: HeapProfileFormat) -> Option<ProfileData> {
+    pub(crate) fn gen(
+        &self,
+        heap: &Heap,
+        format: HeapProfileFormat,
+    ) -> anyhow::Result<ProfileData> {
         if !self.enabled {
-            None
-        } else {
-            Some(Self::gen_enabled(heap, format))
+            return Err(HeapProfileError::NotEnabled.into());
         }
+        Ok(Self::gen_enabled(heap, format))
     }
 
     pub(crate) fn gen_enabled(heap: &Heap, format: HeapProfileFormat) -> ProfileData {

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

@@ -183,11 +183,15 @@ impl StmtProfile {
     }
 
     // None = not applicable because not enabled
-    pub(crate) fn gen(&self) -> Option<ProfileData> {
+    pub(crate) fn gen(&self) -> anyhow::Result<ProfileData> {
         let now = Instant::now();
-        self.0
-            .as_ref()
-            .map(|data| ProfileData::new(ProfileMode::Statement, data.write_to_string(now)))
+        match &self.0 {
+            Some(data) => Ok(ProfileData::new(
+                ProfileMode::Statement,
+                data.write_to_string(now),
+            )),
+            None => Err(StmtProfileError::NotEnabled.into()),
+        }
     }
 
     pub(crate) fn coverage(&self) -> anyhow::Result<HashSet<ResolvedFileSpan>> {

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

@@ -26,6 +26,12 @@ use crate::eval::runtime::small_duration::SmallDuration;
 use crate::eval::ProfileMode;
 use crate::values::FrozenStringValue;
 
+#[derive(Debug, thiserror::Error)]
+enum TypecheckProfileErorr {
+    #[error("Typecheck profile not enabled")]
+    NotEnabled,
+}
+
 #[derive(Default, Debug)]
 pub(crate) struct TypecheckProfile {
     pub(crate) enabled: bool,
@@ -62,11 +68,11 @@ impl TypecheckProfile {
         w.finish()
     }
 
-    pub(crate) fn gen(&self) -> Option<ProfileData> {
+    pub(crate) fn gen(&self) -> anyhow::Result<ProfileData> {
         if !self.enabled {
-            return None;
+            return Err(TypecheckProfileErorr::NotEnabled.into());
         }
-        Some(ProfileData::new(ProfileMode::Typecheck, self.gen_csv()))
+        Ok(ProfileData::new(ProfileMode::Typecheck, self.gen_csv()))
     }
 }