heatmap: Fix symbolization of hot pages in heatmap tool.

This patch fixes the heatmap tool that was broken to display
hot symbols within hot pages with newly added customized hugepage
region. It also adds unit tests to avoid regression in the future.

BUG=chromium:956109
TEST=Pass unittests and tested with one image locally.

Change-Id: I5aba39b6f27582b060248d756ddb9b869345065e
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/1699094
Commit-Queue: Tiancong Wang <tcwang@google.com>
Tested-by: Tiancong Wang <tcwang@google.com>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>
Reviewed-by: George Burgess <gbiv@chromium.org>

4 files changed, 133 insertions, 58 deletions

diff --git a/heatmaps/heat_map.py b/heatmaps/heat_map.py
index b7006e3d..2fd742d2 100755
--- a/heatmaps/heat_map.py
+++ b/heatmaps/heat_map.py
@@ -82,10 +82,7 @@ class HeatMapProducer(object):
     generator.draw()
     # Analyze top N hottest symbols with the binary, if provided
     if self.binary:
-      if top_n_pages is not None:
-        generator.analyze(self.binary, top_n_pages)
-      else:
-        generator.analyze(self.binary)
+      generator.analyze(self.binary, top_n_pages)
 
   def _RemoveFiles(self):
     files = [
@@ -95,7 +92,7 @@ class HeatMapProducer(object):
       if os.path.exists(f):
         os.remove(f)
 
-  def Run(self, top_n_pages=None):
+  def Run(self, top_n_pages):
     try:
       self._EnsureFileInChroot()
       self._GeneratePerfReport()
@@ -105,6 +102,9 @@ class HeatMapProducer(object):
     msg = ('heat map and time histogram genereated in the current '
            'directory with name heat_map.png and timeline.png '
            'accordingly.')
+    if self.binary:
+      msg += ('\nThe hottest %d pages inside and outside hugepage '
+              'is symbolized and saved to addr2symbol.txt' % top_n_pages)
     if self.logger:
       self.logger.LogOutput(msg)
     else:
@@ -142,9 +142,10 @@ def main(argv):
   parser.add_argument(
       '--top_n',
       dest='top_n',
+      type=int,
+      default=10,
       help='Print out top N hottest pages within/outside huge page range. '
-      'Must be used with --hugepage and --binary.',
-      default=None)
+      'Must be used with --hugepage and --binary. (Default: %(default)s)')
   parser.add_argument(
       '--title', dest='title', help='Title of the heatmap', default='')
   parser.add_argument(
diff --git a/heatmaps/heat_map_test.py b/heatmaps/heat_map_test.py
index c6474670..21f90d41 100755
--- a/heatmaps/heat_map_test.py
+++ b/heatmaps/heat_map_test.py
@@ -84,7 +84,7 @@ class HeatmapTest(unittest.TestCase):
   @mock.patch('heatmap_generator.HeatmapGenerator')
   def test_GetHeatMap(self, mock_heatmap_generator):
     heatmap = make_heatmap()
-    heatmap._GetHeatMap(None)
+    heatmap._GetHeatMap(10)
     self.assertTrue(mock_heatmap_generator.called)
 
   @mock.patch.object(heat_map.HeatMapProducer, '_EnsureFileInChroot')
@@ -94,10 +94,10 @@ class HeatmapTest(unittest.TestCase):
   def test_Run(self, mock_remove_files, mock_get_heatmap,
                mock_generate_perf_report, mock_ensure_file_in_chroot):
     heatmap = make_heatmap()
-    heatmap.Run()
+    heatmap.Run(10)
     mock_ensure_file_in_chroot.assert_called_once_with()
     mock_generate_perf_report.assert_called_once_with()
-    mock_get_heatmap.assert_called_once_with(None)
+    mock_get_heatmap.assert_called_once_with(10)
     mock_remove_files.assert_called_once_with()
 
   @mock.patch.object(heat_map.HeatMapProducer, '_EnsureFileInChroot')
@@ -112,7 +112,7 @@ class HeatmapTest(unittest.TestCase):
                               mock_get_heatmap, mock_ensure_file_in_chroot):
     heatmap = make_heatmap()
     with self.assertRaises(Exception):
-      heatmap.Run()
+      heatmap.Run(10)
     mock_ensure_file_in_chroot.assert_called_once_with()
     mock_get_heatmap.assert_not_called()
     mock_remove_files.assert_called_once_with()
diff --git a/heatmaps/heatmap_generator.py b/heatmaps/heatmap_generator.py
index 4f76450c..42fd6352 100644
--- a/heatmaps/heatmap_generator.py
+++ b/heatmaps/heatmap_generator.py
@@ -315,12 +315,17 @@ class HeatmapGenerator(object):
     if retval:
       raise RuntimeError('Failed to run script to generate heatmap')
 
-  def _restore_histogram(self, name):
+  def _restore_histogram(self):
+    # When hugepage is used, there are two files inst-histo-{hp,sp}.txt
+    # So we need to read in all the files.
+    names = [x for x in os.listdir('.') if 'inst-histo' in x and '.txt' in x]
     hist = {}
-    with open(name) as f:
-      for l in f.readlines():
-        num, addr = l.strip().split(' ')
-        hist[int(addr)] = int(num)
+    for n in names:
+      with open(n) as f:
+        for l in f.readlines():
+          num, addr = l.strip().split(' ')
+          assert int(addr) not in hist
+          hist[int(addr)] = int(num)
     return hist
 
   def _read_symbols_from_binary(self, binary):
@@ -367,32 +372,6 @@ class HeatmapGenerator(object):
           self.symbol_addresses.append(addr - text_section_start)
           self.symbol_names.append(name)
 
-  def _get_list_of_pages_to_show(self, hist, top_n):
-    sorted_hist = sorted(
-        hist.iteritems(), key=lambda value: value[1], reverse=True)
-    _, max_count = sorted_hist[0]
-
-    # Depending on the configuration of top_n, select pages in the list
-    # if % is in top_n, e.g. top_n = '20%', we will select the pages that has
-    # sample count larger than 20% of the peak amount
-    # otherwise, if top_n is a number, e.g. top_n = 5, we will select top 5
-    # hottest pages within the 30MB region and top 5 hottest pages outside of
-    # the 30MB region
-
-    if '%' in top_n:
-      count_threshold = max_count * int(top_n[:-1]) / 100
-      list_to_show = [(k, v) for (k, v) in sorted_hist if v >= count_threshold]
-    else:
-      in_huge_page = [
-          (k, v) for (k, v) in sorted_hist \
-          if self.hugepage.start <= k < self.hugepage.end][:int(top_n)]
-      outside_huge_page = [
-          (k, v) for (k, v) in sorted_hist \
-          if k < self.hugepage.start or k >= self.hugepage.end][:int(top_n)]
-
-      list_to_show = in_huge_page + outside_huge_page
-    return list_to_show, max_count
-
   def _map_addr_to_symbol(self, addr):
     # Find out the symbol name
     assert len(self.symbol_addresses) > 0
@@ -402,14 +381,15 @@ class HeatmapGenerator(object):
         index, len(self.symbol_names))
     return self.symbol_names[index - 1]
 
-  def _get_symbols_in_hot_pages(self, fp, pages_to_show, max_count):
+  def _print_symbols_in_hot_pages(self, fp, pages_to_show):
     # Print symbols in all the pages of interest
     for page_num, sample_num in pages_to_show:
       print(
           '----------------------------------------------------------', file=fp)
-      print(('Page Offset: %d MB, Count: %d (%.1f%%)' % (
-          page_num / 1024 / 1024, sample_num, 100.0 * sample_num / max_count)),
-            file=fp)
+      print(
+          'Page Offset: %d MB, Count: %d' % (page_num / 1024 / 1024,
+                                             sample_num),
+          file=fp)
 
       symbol_counts = collections.Counter()
       # Read Sample File and find out the occurance of symbols in the page
@@ -417,8 +397,6 @@ class HeatmapGenerator(object):
       for line in lines:
         if 'PERF_RECORD_SAMPLE' in line:
           pid, addr = self._parse_perf_sample(line + next(lines) + next(lines))
-          lines.next()
-          lines.next()
           if pid is None:
             # The sampling is not on Chrome
             continue
@@ -452,14 +430,39 @@ class HeatmapGenerator(object):
     # Then use gnu plot to draw heat map
     self._draw_heat_map()
 
-  def analyze(self, binary, top_n='10'):
+  def analyze(self, binary, top_n):
     # Read histogram from histo.txt
-    hist = self._restore_histogram('inst-histo.txt')
-    # Generate Symbol Names and save it to nm.txt
+    hist = self._restore_histogram()
+    # Sort the pages in histogram
+    sorted_hist = sorted(
+        hist.iteritems(), key=lambda value: value[1], reverse=True)
+
+    # Generate symbolizations
     self._read_symbols_from_binary(binary)
-    # Sort the pages according to the hotness
-    pages_to_show, max_count = self._get_list_of_pages_to_show(hist, top_n)
 
     # Write hottest pages
     with open('addr2symbol.txt', 'w') as fp:
-      self._get_symbols_in_hot_pages(fp, pages_to_show, max_count)
+      if self.hugepage:
+        # Print hugepage region first
+        print(
+            'Hugepage top %d hot pages (%d MB - %d MB):' %
+            (top_n, self.hugepage.start / 1024 / 1024,
+             self.hugepage.end / 1024 / 1024),
+            file=fp)
+        pages_to_print = [(k, v)
+                          for k, v in sorted_hist
+                          if self.hugepage.start <= k < self.hugepage.end
+                         ][:top_n]
+        self._print_symbols_in_hot_pages(fp, pages_to_print)
+        print('==========================================', file=fp)
+        print('Top %d hot pages landed outside of hugepage:' % top_n, file=fp)
+        # Then print outside pages
+        pages_to_print = [(k, v)
+                          for k, v in sorted_hist
+                          if k < self.hugepage.start or k >= self.hugepage.end
+                         ][:top_n]
+        self._print_symbols_in_hot_pages(fp, pages_to_print)
+      else:
+        # Print top_n hottest pages.
+        pages_to_print = [(k, v) for k, v in sorted_hist][:top_n]
+        self._print_symbols_in_hot_pages(fp, pages_to_print)
diff --git a/heatmaps/heatmap_generator_test.py b/heatmaps/heatmap_generator_test.py
index 271f9f42..0c0bbfc8 100755
--- a/heatmaps/heatmap_generator_test.py
+++ b/heatmaps/heatmap_generator_test.py
@@ -8,6 +8,7 @@
 
 from __future__ import division, print_function
 
+import mock
 import os
 import unittest
 
@@ -45,25 +46,27 @@ def _write_perf_sample(pid, tid, addr, fp):
   print(' ...... dso: /opt/google/chrome/chrome\n', file=fp)
 
 
-def _heatmap(file_name, page_size=4096, hugepage=None):
+def _heatmap(file_name, page_size=4096, hugepage=None, analyze=False, top_n=10):
   generator = heatmap_generator.HeatmapGenerator(
       file_name, page_size, hugepage, '',
       log_level='none')  # Don't log to stdout
   generator.draw()
+  if analyze:
+    generator.analyze('/path/to/chrome', top_n)
 
 
 def _cleanup(file_name):
   files = [
       file_name, 'out.txt', 'inst-histo.txt', 'inst-histo-hp.txt',
-      'inst-histo-sp.txt', 'heat_map.png', 'timeline.png'
+      'inst-histo-sp.txt', 'heat_map.png', 'timeline.png', 'addr2symbol.txt'
   ]
   for f in files:
     if os.path.exists(f):
       os.remove(f)
 
 
-class Tests(unittest.TestCase):
-  """All of our tests for heatmap_generator."""
+class HeatmapGeneratorDrawTests(unittest.TestCase):
+  """All of our tests for heatmap_generator.draw() and related."""
 
   def test_with_one_mmap_one_sample(self):
     """Tests one perf record and one sample."""
@@ -240,5 +243,73 @@ class Tests(unittest.TestCase):
       self.assertIn('100 4194304', lines[1])
 
 
+class HeatmapGeneratorAnalyzeTests(unittest.TestCase):
+  """All of our tests for heatmap_generator.analyze() and related."""
+
+  def setUp(self):
+    # Use the same perf report for testing
+    self.fname = 'test_histo.txt'
+    with open(self.fname, 'w') as f:
+      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
+      for i in range(10):
+        _write_perf_sample(101, 101, 0xABCD000 + i, f)
+        _write_perf_sample(101, 101, 0xABCE000 + i, f)
+        _write_perf_sample(101, 101, 0xABFD000 + i, f)
+    self.nm = ('000000000abcd000 t Func1@Page1\n'
+               '000000000abcd001 t Func2@Page1\n'
+               '000000000abcd0a0 t Func3@Page1andFunc1@Page2\n'
+               '000000000abce010 t Func2@Page2\n'
+               '000000000abfd000 t Func1@Page3\n')
+
+  def tearDown(self):
+    _cleanup(self.fname)
+
+  @mock.patch('subprocess.check_output')
+  def test_analyze_hot_pages_with_hp_top(self, mock_nm):
+    """Test if the analyze() can print the top page with hugepage."""
+    mock_nm.return_value = self.nm
+    _heatmap(self.fname, hugepage=[0, 8192], analyze=True, top_n=1)
+    file_list = os.listdir('.')
+    self.assertIn('addr2symbol.txt', file_list)
+    with open('addr2symbol.txt') as f:
+      contents = f.read()
+      self.assertIn('Func2@Page1 : 9', contents)
+      self.assertIn('Func1@Page1 : 1', contents)
+      self.assertIn('Func1@Page3 : 10', contents)
+      # Only displaying one page in hugepage
+      self.assertNotIn('Func3@Page1andFunc1@Page2 : 10', contents)
+
+  @mock.patch('subprocess.check_output')
+  def test_analyze_hot_pages_without_hp_top(self, mock_nm):
+    """Test if the analyze() can print the top page without hugepage."""
+    mock_nm.return_value = self.nm
+    _heatmap(self.fname, analyze=True, top_n=1)
+    file_list = os.listdir('.')
+    self.assertIn('addr2symbol.txt', file_list)
+    with open('addr2symbol.txt') as f:
+      contents = f.read()
+      self.assertIn('Func2@Page1 : 9', contents)
+      self.assertIn('Func1@Page1 : 1', contents)
+      # Only displaying one page
+      self.assertNotIn('Func3@Page1andFunc1@Page2 : 10', contents)
+      self.assertNotIn('Func1@Page3 : 10', contents)
+
+  @mock.patch('subprocess.check_output')
+  def test_analyze_hot_pages_with_hp_top10(self, mock_nm):
+    """Test if the analyze() can print with default top 10."""
+    mock_nm.return_value = self.nm
+    _heatmap(self.fname, analyze=True)
+    # Make sure nm command is called correctly.
+    mock_nm.assert_called_with(['nm', '-n', '/path/to/chrome'])
+    file_list = os.listdir('.')
+    self.assertIn('addr2symbol.txt', file_list)
+    with open('addr2symbol.txt') as f:
+      contents = f.read()
+      self.assertIn('Func2@Page1 : 9', contents)
+      self.assertIn('Func1@Page1 : 1', contents)
+      self.assertIn('Func3@Page1andFunc1@Page2 : 10', contents)
+      self.assertIn('Func1@Page3 : 10', contents)
+
+
 if __name__ == '__main__':
   unittest.main()