Temporal inverted flame graph

Added --inverted option for the temporal flame graph. Signed-off-by: Ivona Stojanovic <[email protected]>
bloomberg · Sep 8, 2023 · 21bdff5 · 21bdff5
1 parent 68a453a
commit 21bdff5
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 45 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
@@ -27,9 +27,8 @@ var parent_index_by_child_index = (function () {
   return ret;
 })();
 
-function packedDataToTree(packedData, rangeStart, rangeEnd) {
-  const { strings, nodes, unique_threads } = packedData;
-
+function generateNodeObjects(packedData) {
+  const { strings, nodes } = packedData;
   console.log("constructing nodes");
   const node_objects = nodes.name.map((_, i) => ({
     name: strings[nodes["name"][i]],
@@ -51,6 +50,88 @@ 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: strings,
+    nodes: nodes,
+    unique_threads: unique_threads,
+  });
+
+  const invertedNoImportsNodeObjects = inverted
+    ? generateNodeObjects({
+        strings: strings,
+        nodes: inverted_no_imports_nodes,
+        unique_threads: unique_threads,
+      })
+    : null;
+
+  flamegraphIntervals = intervals;
+  invertedNoImportsIntervals = no_imports_interval_list;
+
+  return {
+    flamegraphNodeObjects: flamegraphNodeObjects,
+    invertedNoImportsNodeObjects: invertedNoImportsNodeObjects,
+  };
+}
+
+function findHWMAllocations(intervals, node_objects, hwmSnapshot) {
+  if (!node_objects) {
+    return;
+  }
+
+  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) {
+  if (!node_objects) {
+    return;
+  }
+
+  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 +194,42 @@ 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);
+    findHWMAllocations(
+      invertedNoImportsIntervals,
+      invertedNoImportsNodeObjects,
+      hwmSnapshot
+    );
   } 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
+    );
+    findLeakedAllocations(
+      invertedNoImportsIntervals,
+      invertedNoImportsNodeObjects,
+      rangeStart,
+      rangeEnd
+    );
   }
 
-  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){
+  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 +334,13 @@ 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 ? c : 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["high_water_mark_by_snapshot"] != 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 %}