From 7ecffc1de5f7c42f8a51dce02de0971c483f319b Mon Sep 17 00:00:00 2001
From: Neil Mitchell <ndmitchell@fb.com>
Date: Tue, 14 Sep 2021 15:47:04 -0700
Subject: [PATCH] Add available_bytes to the heaps

Summary: Useful for figuring out memory that is allocated but not used.

Reviewed By: swgillespie, krallin

Differential Revision: D30939477

fbshipit-source-id: 17c4ab1f40c94197a844aec86e9cadc0da25b6a8
---
 .../starlark/src/values/layout/arena.rs         |  6 +++++-
 .../starlark/src/values/layout/heap.rs          | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/third-party/rust/starlark-rust/starlark/src/values/layout/arena.rs b/third-party/rust/starlark-rust/starlark/src/values/layout/arena.rs
index 8b58a2e62a3..d319e5be22c 100644
--- a/third-party/rust/starlark-rust/starlark/src/values/layout/arena.rs
+++ b/third-party/rust/starlark-rust/starlark/src/values/layout/arena.rs
@@ -117,6 +117,10 @@ impl Arena {
         self.drop.allocated_bytes() + self.non_drop.allocated_bytes()
     }
 
+    pub fn available_bytes(&self) -> usize {
+        self.drop.chunk_capacity() + self.non_drop.chunk_capacity()
+    }
+
     /// Bytes allocated which can't be iterated over
     pub fn allocated_bytes_inline(&self) -> usize {
         self.non_drop.allocated_bytes()
@@ -478,6 +482,6 @@ mod test {
         assert_eq!(res.len(), 1);
         let entry = res.values().next().unwrap();
         assert_eq!(entry.0, 2);
-        assert_eq!(entry.1, arena.allocated_bytes())
+        assert_eq!(entry.1, arena.allocated_bytes());
     }
 }
diff --git a/third-party/rust/starlark-rust/starlark/src/values/layout/heap.rs b/third-party/rust/starlark-rust/starlark/src/values/layout/heap.rs
index 0fe3a284b10..d2a63789b32 100644
--- a/third-party/rust/starlark-rust/starlark/src/values/layout/heap.rs
+++ b/third-party/rust/starlark-rust/starlark/src/values/layout/heap.rs
@@ -133,6 +133,13 @@ impl FrozenHeapRef {
         self.0.arena.allocated_bytes()
     }
 
+    /// Number of bytes allocated by the heap but not filled.
+    /// Note that these bytes will _never_ be filled as no further allocations can
+    /// be made on this heap (it has been sealed).
+    pub fn available_bytes(&self) -> usize {
+        self.0.arena.available_bytes()
+    }
+
     /// Obtain a summary of how much memory is currently allocated by this heap.
     /// Doesn't include the heaps it keeps alive by reference.
     pub fn allocated_summary(&self) -> HeapSummary {
@@ -204,6 +211,11 @@ impl FrozenHeap {
         self.arena.allocated_bytes()
     }
 
+    /// Number of bytes allocated by the heap but not yet filled.
+    pub fn available_bytes(&self) -> usize {
+        self.arena.available_bytes()
+    }
+
     /// Obtain a summary of how much memory is currently allocated by this heap.
     pub fn allocated_summary(&self) -> HeapSummary {
         self.arena.allocated_summary()
@@ -289,6 +301,11 @@ impl Heap {
         self.arena.borrow().allocated_bytes()
     }
 
+    /// Number of bytes allocated by the heap but not yet filled.
+    pub fn available_bytes(&self) -> usize {
+        self.arena.borrow().available_bytes()
+    }
+
     /// Only those allocated on the inline heap (mostly strings)
     pub(crate) fn allocated_bytes_inline(&self) -> usize {
         self.arena.borrow().allocated_bytes_inline()