Move record_call_enter/exit to Heap

Summary:
We need to push logic of tracking of allocations per function down to arena to be able to measure retained memory, not allocated memory.

The idea is: `Arena` should have a map from pointer to call stack.

Reviewed By: ndmitchell

Differential Revision: D37606255

fbshipit-source-id: 01bae7d9f5b43d707e9d55741afc8298f5d82f7e

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

@@ -21,22 +21,16 @@ use std::collections::HashMap;
 use std::fmt::Debug;
 use std::fs::File;
 use std::io::Write;
-use std::mem;
 use std::path::Path;
 use std::rc::Rc;
 use std::time::Instant;
 
 use anyhow::Context;
-use derive_more::Display;
-use gazebo::any::ProvidesStaticType;
 use gazebo::prelude::*;
 
-use crate as starlark;
 use crate::eval::runtime::profile::csv::CsvWriter;
 use crate::eval::runtime::small_duration::SmallDuration;
 use crate::values::Heap;
-use crate::values::StarlarkValue;
-use crate::values::Trace;
 use crate::values::Value;
 use crate::values::ValueLike;
 
@@ -50,50 +44,6 @@ pub(crate) struct HeapProfile {
     enabled: bool,
 }
 
-/// A type which is either drop or non-drop.
-trait MaybeDrop: Debug + Sync + Send + 'static {}
-
-/// Type which has `Drop`.
-#[derive(ProvidesStaticType, Debug, Trace)]
-struct NeedsDrop;
-impl Drop for NeedsDrop {
-    fn drop(&mut self) {
-        // Just make this type `Drop`.
-        // Note `mem::needs_drop()` would return `true` for this type,
-        // even if `drop` is optimized away: https://rust.godbolt.org/z/1cxKoMzdM
-    }
-}
-
-/// Type which doesn't have `Drop`.
-#[derive(ProvidesStaticType, Debug, Trace)]
-struct NoDrop;
-
-impl MaybeDrop for NeedsDrop {}
-impl MaybeDrop for NoDrop {}
-
-#[derive(Trace, Debug, Display, ProvidesStaticType, NoSerialize)]
-#[display(fmt = "CallEnter")]
-struct CallEnter<'v, D: MaybeDrop + 'static> {
-    function: Value<'v>,
-    time: Instant,
-    maybe_drop: D,
-}
-
-impl<'v, D: MaybeDrop + Trace<'v> + 'v> StarlarkValue<'v> for CallEnter<'v, D> {
-    starlark_type!("call_enter");
-}
-
-#[derive(Debug, Display, ProvidesStaticType, NoSerialize)]
-#[display(fmt = "CallExit")]
-struct CallExit<D: MaybeDrop + 'static> {
-    time: Instant,
-    maybe_drop: D,
-}
-
-impl<'v, D: MaybeDrop> StarlarkValue<'v> for CallExit<D> {
-    starlark_type!("call_exit");
-}
-
 #[derive(Copy, Clone, Dupe, Debug, Eq, PartialEq, Hash)]
 struct FunctionId(usize);
 
@@ -161,37 +111,15 @@ impl HeapProfile {
     #[inline(never)]
     pub(crate) fn record_call_enter<'v>(&self, function: Value<'v>, heap: &'v Heap) {
         if self.enabled {
-            let time = Instant::now();
-            assert!(mem::needs_drop::<CallEnter<NeedsDrop>>());
-            assert!(!mem::needs_drop::<CallEnter<NoDrop>>());
-            heap.alloc_complex_no_freeze(CallEnter {
-                function,
-                time,
-                maybe_drop: NeedsDrop,
-            });
-            heap.alloc_complex_no_freeze(CallEnter {
-                function,
-                time,
-                maybe_drop: NoDrop,
-            });
+            heap.record_call_enter(function);
         }
     }
 
     #[cold]
     #[inline(never)]
     pub(crate) fn record_call_exit<'v>(&self, heap: &'v Heap) {
         if self.enabled {
-            let time = Instant::now();
-            assert!(mem::needs_drop::<CallExit<NeedsDrop>>());
-            assert!(!mem::needs_drop::<CallExit<NoDrop>>());
-            heap.alloc_simple(CallExit {
-                time,
-                maybe_drop: NeedsDrop,
-            });
-            heap.alloc_simple(CallExit {
-                time,
-                maybe_drop: NoDrop,
-            });
+            heap.record_call_exit();
         }
     }
 
@@ -325,6 +253,11 @@ impl HeapProfile {
 mod summary {
     use super::*;
     use crate::eval::runtime::small_duration::SmallDuration;
+    use crate::values::layout::heap::call_enter_exit::CallEnter;
+    use crate::values::layout::heap::call_enter_exit::CallExit;
+    use crate::values::layout::heap::call_enter_exit::MaybeDrop;
+    use crate::values::layout::heap::call_enter_exit::NeedsDrop;
+    use crate::values::layout::heap::call_enter_exit::NoDrop;
 
     /// Information relating to a function.
     #[derive(Default, Debug, Clone)]
@@ -439,6 +372,10 @@ mod summary {
 
 mod flame {
     use super::*;
+    use crate::values::layout::heap::call_enter_exit::CallEnter;
+    use crate::values::layout::heap::call_enter_exit::CallExit;
+    use crate::values::layout::heap::call_enter_exit::NeedsDrop;
+    use crate::values::layout::heap::call_enter_exit::NoDrop;
 
     /// Allocations made in a given stack frame for a given type.
     #[derive(Default)]

starlark-rust/starlark/src/values/layout/heap/call_enter_exit.rs

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2019 The Starlark in Rust Authors.
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! Marker objects to track allocations.
+
+use std::fmt::Debug;
+use std::time::Instant;
+
+use gazebo::any::ProvidesStaticType;
+
+use crate as starlark;
+use crate::values::StarlarkValue;
+use crate::values::Trace;
+use crate::values::Value;
+
+/// A type which is either drop or non-drop.
+pub(crate) trait MaybeDrop: Debug + Sync + Send + 'static {}
+
+/// Type which has `Drop`.
+#[derive(ProvidesStaticType, Debug, Trace)]
+pub(crate) struct NeedsDrop;
+impl Drop for NeedsDrop {
+    fn drop(&mut self) {
+        // Just make this type `Drop`.
+        // Note `mem::needs_drop()` would return `true` for this type,
+        // even if `drop` is optimized away: https://rust.godbolt.org/z/1cxKoMzdM
+    }
+}
+
+/// Type which doesn't have `Drop`.
+#[derive(ProvidesStaticType, Debug, Trace)]
+pub(crate) struct NoDrop;
+
+impl MaybeDrop for NeedsDrop {}
+impl MaybeDrop for NoDrop {}
+
+#[derive(Trace, Debug, derive_more::Display, ProvidesStaticType, NoSerialize)]
+#[display(fmt = "CallEnter")]
+pub(crate) struct CallEnter<'v, D: MaybeDrop + 'static> {
+    pub(crate) function: Value<'v>,
+    pub(crate) time: Instant,
+    pub(crate) maybe_drop: D,
+}
+
+impl<'v, D: MaybeDrop + Trace<'v> + 'v> StarlarkValue<'v> for CallEnter<'v, D> {
+    starlark_type!("call_enter");
+}
+
+#[derive(Debug, derive_more::Display, ProvidesStaticType, NoSerialize)]
+#[display(fmt = "CallExit")]
+pub(crate) struct CallExit<D: MaybeDrop + 'static> {
+    pub(crate) time: Instant,
+    pub(crate) maybe_drop: D,
+}
+
+impl<'v, D: MaybeDrop> StarlarkValue<'v> for CallExit<D> {
+    starlark_type!("call_exit");
+}

starlark-rust/starlark/src/values/layout/heap/heap_type.rs

@@ -27,11 +27,13 @@ use std::hash::Hash;
 use std::hash::Hasher;
 use std::intrinsics::copy_nonoverlapping;
 use std::marker::PhantomData;
+use std::mem;
 use std::mem::MaybeUninit;
 use std::ops::Deref;
 use std::ptr;
 use std::slice;
 use std::sync::Arc;
+use std::time::Instant;
 use std::usize;
 
 use either::Either;
@@ -61,6 +63,10 @@ use crate::values::layout::avalue::VALUE_EMPTY_TUPLE;
 use crate::values::layout::heap::arena::Arena;
 use crate::values::layout::heap::arena::HeapSummary;
 use crate::values::layout::heap::arena::Reservation;
+use crate::values::layout::heap::call_enter_exit::CallEnter;
+use crate::values::layout::heap::call_enter_exit::CallExit;
+use crate::values::layout::heap::call_enter_exit::NeedsDrop;
+use crate::values::layout::heap::call_enter_exit::NoDrop;
 use crate::values::layout::heap::fast_cell::FastCell;
 use crate::values::layout::heap::repr::AValueRepr;
 use crate::values::layout::static_string::constant_string;
@@ -696,6 +702,36 @@ impl Heap {
     pub fn allocated_summary(&self) -> HeapSummary {
         self.arena.borrow().allocated_summary()
     }
+
+    pub(crate) fn record_call_enter<'v>(&'v self, function: Value<'v>) {
+        let time = Instant::now();
+        assert!(mem::needs_drop::<CallEnter<NeedsDrop>>());
+        assert!(!mem::needs_drop::<CallEnter<NoDrop>>());
+        self.alloc_complex_no_freeze(CallEnter {
+            function,
+            time,
+            maybe_drop: NeedsDrop,
+        });
+        self.alloc_complex_no_freeze(CallEnter {
+            function,
+            time,
+            maybe_drop: NoDrop,
+        });
+    }
+
+    pub(crate) fn record_call_exit<'v>(&'v self) {
+        let time = Instant::now();
+        assert!(mem::needs_drop::<CallExit<NeedsDrop>>());
+        assert!(!mem::needs_drop::<CallExit<NoDrop>>());
+        self.alloc_simple(CallExit {
+            time,
+            maybe_drop: NeedsDrop,
+        });
+        self.alloc_simple(CallExit {
+            time,
+            maybe_drop: NoDrop,
+        });
+    }
 }
 
 /// Used to perform garbage collection by [`Trace::trace`](crate::values::Trace::trace).

starlark-rust/starlark/src/values/layout/heap/mod.rs

@@ -18,6 +18,7 @@
 //! Starlark heap implementation.
 
 pub(crate) mod arena;
+pub(crate) mod call_enter_exit;
 mod fast_cell;
 pub(crate) mod heap_type;
 pub(crate) mod repr;