Support inverted temporal inverted flame graphs

Allow `memray flamegraph` to support `--temporal` and `--inverted` at the same time. Signed-off-by: Ivona Stojanovic <[email protected]>
bloomberg · Sep 8, 2023 · 514af3e · 514af3e
1 parent bb9c0db
commit 514af3e
Show file tree

Hide file tree

Showing 6 changed files with 151 additions and 49 deletions.
diff --git a/src/memray/commands/common.py b/src/memray/commands/common.py
@@ -271,9 +271,6 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
 
         if hasattr(args, "inverted"):
             kwargs["inverted"] = args.inverted
-            if temporal and args.inverted:
-                parser.error("Can't create an inverted temporal flame graph.")
-
         self.write_report(
             result_path,
             output_file,

diff --git a/src/memray/reporters/assets/flamegraph_common.js b/src/memray/reporters/assets/flamegraph_common.js
@@ -146,12 +146,17 @@ export function onFilterImportSystem() {
       });
     } else {
       data = invertedNoImportsData;
+      if (temporal) {
+        hideImports = true;
+        intervals = invertedNoImportsIntervals;
+      }
     }
   } else {
     filteredChart.unRegisterFilter(FILTER_IMPORT_SYSTEM);
-    if (inverted) {
-      //We can remove this check when we add ``flamegraphData`` to the temporal flamegraph
-      data = flamegraphData;
+    data = flamegraphData;
+    if (temporal) {
+      hideImports = false;
+      intervals = flamegraphIntervals;
     }
   }
   this.hideImportSystemFrames = !this.hideImportSystemFrames;

diff --git a/src/memray/reporters/assets/temporal_flamegraph.js b/src/memray/reporters/assets/temporal_flamegraph.js
@@ -17,19 +17,22 @@ import {
 var active_plot = null;
 var current_dimensions = null;
 
-var parent_index_by_child_index = (function () {
-  let ret = new Array(packed_data.nodes.children.length);
+var parent_index_by_child_index = generateParentIndexes(packed_data.nodes);
+var inverted_no_imports_parent_index_by_child_index = inverted
+  ? generateParentIndexes(packed_data.inverted_no_imports_nodes)
+  : null;
+
+function generateParentIndexes(nodes) {
+  let ret = new Array(nodes.children.length);
   console.log("finding parent index for each node");
-  for (const [parentIndex, children] of packed_data.nodes.children.entries()) {
+  for (const [parentIndex, children] of nodes.children.entries()) {
     children.forEach((idx) => (ret[idx] = parentIndex));
   }
   console.assert(ret[0] === undefined, "root node has a parent");
   return ret;
-})();
-
-function packedDataToTree(packedData, rangeStart, rangeEnd) {
-  const { strings, nodes, unique_threads } = packedData;
+}
 
+function generateNodeObjects(strings, nodes) {
   console.log("constructing nodes");
   const node_objects = nodes.name.map((_, i) => ({
     name: strings[nodes["name"][i]],
@@ -51,6 +54,83 @@ function packedDataToTree(packedData, rangeStart, rangeEnd) {
     node["children"] = node["children"].map((idx) => node_objects[idx]);
   }
 
+  return node_objects;
+}
+
+function initTrees(packedData) {
+  const {
+    strings,
+    nodes,
+    inverted_no_imports_nodes,
+    unique_threads,
+    intervals,
+    no_imports_interval_list,
+  } = packedData;
+
+  const flamegraphNodeObjects = generateNodeObjects(strings, nodes);
+  const invertedNoImportsNodeObjects = inverted
+    ? generateNodeObjects(strings, inverted_no_imports_nodes)
+    : null;
+
+  flamegraphIntervals = intervals;
+  invertedNoImportsIntervals = no_imports_interval_list;
+
+  return {
+    flamegraphNodeObjects: flamegraphNodeObjects,
+    invertedNoImportsNodeObjects: invertedNoImportsNodeObjects,
+  };
+}
+
+function findHWMAllocations(
+  intervals,
+  node_objects,
+  hwmSnapshot,
+  parent_index_by_child_index
+) {
+  intervals.forEach((interval) => {
+    let [allocBefore, deallocBefore, nodeIndex, count, bytes] = interval;
+
+    if (
+      allocBefore <= hwmSnapshot &&
+      (deallocBefore === null || deallocBefore > hwmSnapshot)
+    ) {
+      while (nodeIndex !== undefined) {
+        node_objects[nodeIndex].n_allocations += count;
+        node_objects[nodeIndex].value += bytes;
+        nodeIndex = parent_index_by_child_index[nodeIndex];
+      }
+    }
+  });
+}
+
+function findLeakedAllocations(
+  intervals,
+  node_objects,
+  rangeStart,
+  rangeEnd,
+  parent_index_by_child_index
+) {
+  intervals.forEach((interval) => {
+    let [allocBefore, deallocBefore, nodeIndex, count, bytes] = interval;
+
+    if (
+      allocBefore >= rangeStart &&
+      allocBefore <= rangeEnd &&
+      (deallocBefore === null || deallocBefore > rangeEnd)
+    ) {
+      while (nodeIndex !== undefined) {
+        node_objects[nodeIndex].n_allocations += count;
+        node_objects[nodeIndex].value += bytes;
+        nodeIndex = parent_index_by_child_index[nodeIndex];
+      }
+    }
+  });
+}
+
+function packedDataToTree(packedData, rangeStart, rangeEnd) {
+  const { flamegraphNodeObjects, invertedNoImportsNodeObjects } =
+    initTrees(packedData);
+
   const hwms = packedData.high_water_mark_by_snapshot;
   if (hwms) {
     console.log("finding highest high water mark in range");
@@ -113,48 +193,53 @@ function packedDataToTree(packedData, rangeStart, rangeEnd) {
 
     // We could binary search rather than using a linear scan...
     console.log("finding hwm allocations");
-    packedData.intervals.forEach((interval) => {
-      let [allocBefore, deallocBefore, nodeIndex, count, bytes] = interval;
-
-      if (
-        allocBefore <= hwmSnapshot &&
-        (deallocBefore === null || deallocBefore > hwmSnapshot)
-      ) {
-        while (nodeIndex !== undefined) {
-          node_objects[nodeIndex].n_allocations += count;
-          node_objects[nodeIndex].value += bytes;
-          nodeIndex = parent_index_by_child_index[nodeIndex];
-        }
-      }
-    });
+    findHWMAllocations(
+      flamegraphIntervals,
+      flamegraphNodeObjects,
+      hwmSnapshot,
+      parent_index_by_child_index
+    );
+    if (inverted) {
+      findHWMAllocations(
+        invertedNoImportsIntervals,
+        invertedNoImportsNodeObjects,
+        hwmSnapshot,
+        inverted_no_imports_parent_index_by_child_index
+      );
+    }
   } else {
     // We could binary search rather than using a linear scan...
     console.log("finding leaked allocations");
-    packedData.intervals.forEach((interval) => {
-      let [allocBefore, deallocBefore, nodeIndex, count, bytes] = interval;
-
-      if (
-        allocBefore >= rangeStart &&
-        allocBefore <= rangeEnd &&
-        (deallocBefore === null || deallocBefore > rangeEnd)
-      ) {
-        while (nodeIndex !== undefined) {
-          node_objects[nodeIndex].n_allocations += count;
-          node_objects[nodeIndex].value += bytes;
-          nodeIndex = parent_index_by_child_index[nodeIndex];
-        }
-      }
-    });
+    findLeakedAllocations(
+      flamegraphIntervals,
+      flamegraphNodeObjects,
+      rangeStart,
+      rangeEnd,
+      parent_index_by_child_index
+    );
+    if (inverted) {
+      findLeakedAllocations(
+        invertedNoImportsIntervals,
+        invertedNoImportsNodeObjects,
+        rangeStart,
+        rangeEnd,
+        inverted_no_imports_parent_index_by_child_index
+      );
+    }
   }
 
-  console.log("total allocations in range: " + node_objects[0].n_allocations);
-  console.log("total bytes in range: " + node_objects[0].value);
-
-  node_objects.forEach((node) => {
+  flamegraphNodeObjects.forEach((node) => {
     node.children = node.children.filter((node) => node.n_allocations > 0);
   });
 
-  return node_objects[0];
+  if (inverted) {
+    invertedNoImportsNodeObjects.forEach((node) => {
+      node.children = node.children.filter((node) => node.n_allocations > 0);
+    });
+  }
+
+  flamegraphData = flamegraphNodeObjects[0];
+  invertedNoImportsData = inverted ? invertedNoImportsNodeObjects[0] : null;
 }
 
 function initMemoryGraph(memory_records) {
@@ -259,7 +344,14 @@ function refreshFlamegraph(event) {
   console.log("last possible index is " + memory_records.length);
 
   console.log("constructing tree");
-  data = packedDataToTree(packed_data, idx0, idx1);
+  packedDataToTree(packed_data, idx0, idx1);
+
+  data = inverted && hideImports ? invertedNoImportsData : flamegraphData;
+  intervals =
+    inverted && hideImports ? invertedNoImportsIntervals : flamegraphIntervals;
+
+  console.log("total allocations in range: " + data.n_allocations);
+  console.log("total bytes in range: " + data.value);
 
   console.log("drawing chart");
   getFilteredChart().drawChart(data);

diff --git a/src/memray/reporters/flamegraph.py b/src/memray/reporters/flamegraph.py
@@ -255,6 +255,7 @@ def _from_any_snapshot(
             inverted_no_imports_frames = [cls._create_root_node()]
 
         interval_list: List[Tuple[int, Optional[int], int, int, int]] = []
+        no_imports_interval_list: List[Tuple[int, Optional[int], int, int, int]] = []
 
         NodeKey = Tuple[int, StackFrame, str]
         node_index_by_key: Dict[NodeKey, int] = {}
@@ -316,7 +317,7 @@ def _from_any_snapshot(
                     node_index_by_key=inverted_no_imports_node_index_by_key,
                     record=record_data,
                     inverted=inverted,
-                    interval_list=interval_list,
+                    interval_list=no_imports_interval_list,
                 )
 
         all_strings = StringRegistry()
@@ -340,6 +341,8 @@ def _from_any_snapshot(
 
         if interval_list:
             data["intervals"] = interval_list
+        if no_imports_interval_list:
+            data["no_imports_interval_list"] = no_imports_interval_list
 
         return cls(data, memory_records=memory_records)
 

diff --git a/src/memray/reporters/templates/base.html b/src/memray/reporters/templates/base.html
@@ -144,6 +144,7 @@ <h5 class="modal-title" id="helpModalLabel">How to interpret {{ kind }} reports<
     const merge_threads = {{ merge_threads|tojson }};
     const memory_records = {{ memory_records|tojson }};
     const inverted = {{ inverted|tojson }};
+    const temporal = packed_data.intervals != null;
   </script>
   {% endblock scripts %}
 </body>

diff --git a/src/memray/reporters/templates/temporal_flamegraph.html b/src/memray/reporters/templates/temporal_flamegraph.html
@@ -39,5 +39,9 @@
 {% block flamegraph_script %}
 <script type="text/javascript">
   {{ include_file("assets/temporal_flamegraph.js") }}
+  var hideImports = false;
+  var intervals = null;
+  var flamegraphIntervals = null;
+  var invertedNoImportsIntervals = null;
 </script>
 {% endblock %}