toolchain-utils: Improve heatmap tool.

The overall goal of this patch is to improve heatmap tool in a way that
facilitates the tool to be integrated into crosperf. The patch needs to
make sure the heat_map.py is not only able to run as a standalone tool,
but also is compatible to crosperf. It also adds functionality to the
heatmap that when the address range of huge page (offset within the text
section) is provided, the heatmap displays symbols within and outside of
the region with different colors.

This patch adds following improvements:
1. It adds support in heat_map.py that skips copying perf.data into
the chroot again, if the provided path to perf.data is within a chroot.
2. It adds support in heatmap_generator.py to take in a range of huge page
as (start, end). Update corresponding unit test.
3. It adds extensive unit tests for heat_map.py. The unit test tests both
each functions and also makes sure it run as a standalone tool.

BUG=chromium:956109
TEST=Tested locally and pass unit tests.

Change-Id: I8008da51cceb0ad90df29042faf7a9d91d04f4b1
Reviewed-on: https://chromium-review.googlesource.com/1582523
Commit-Ready: Tiancong Wang <tcwang@google.com>
Tested-by: Tiancong Wang <tcwang@google.com>
Reviewed-by: Tiancong Wang <tcwang@google.com>
diff --git a/heat_map.py b/heat_map.py
deleted file mode 100755
index 1009eca..0000000
--- a/heat_map.py
+++ /dev/null
@@ -1,144 +0,0 @@
-#!/usr/bin/env python2
-# -*- coding: utf-8 -*-
-# Copyright 2015 The Chromium OS Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-"""Wrapper to generate heat maps for chrome."""
-
-from __future__ import print_function
-
-import argparse
-import shutil
-import os
-import sys
-import tempfile
-
-from cros_utils import command_executer
-import heatmap_generator
-
-
-def IsARepoRoot(directory):
-  """Returns True if directory is the root of a repo checkout."""
-  return os.path.exists(os.path.join(directory, '.repo'))
-
-
-class HeatMapProducer(object):
-  """Class to produce heat map."""
-
-  def __init__(self, chromeos_root, perf_data, page_size, binary, title):
-    self.chromeos_root = os.path.realpath(chromeos_root)
-    self.perf_data = os.path.realpath(perf_data)
-    self.page_size = page_size
-    self.dir = os.path.dirname(os.path.realpath(__file__))
-    self.binary = binary
-    self.tempDir = ''
-    self.ce = command_executer.GetCommandExecuter()
-    self.loading_address = None
-    self.temp_perf = ''
-    self.temp_perf_inchroot = ''
-    self.perf_report = ''
-    self.title = title
-
-  def copyFileToChroot(self):
-    self.tempDir = tempfile.mkdtemp(
-        prefix=os.path.join(self.chromeos_root, 'src/'))
-    self.temp_perf = os.path.join(self.tempDir, 'perf.data')
-    shutil.copy2(self.perf_data, self.temp_perf)
-    self.temp_perf_inchroot = os.path.join('~/trunk/src',
-                                           os.path.basename(self.tempDir))
-
-  def getPerfReport(self):
-    if os.path.isfile(os.path.join(self.tempDir, 'perf_report.txt')):
-      self.perf_report = os.path.join(self.tempDir, 'perf_report.txt')
-      return
-
-    cmd = ('cd %s; perf report -D -i perf.data > perf_report.txt' %
-           self.temp_perf_inchroot)
-    retval = self.ce.ChrootRunCommand(self.chromeos_root, cmd)
-    if retval:
-      raise RuntimeError('Failed to generate perf report')
-    self.perf_report = os.path.join(self.tempDir, 'perf_report.txt')
-
-  def getHeatMap(self, top_n_pages=None):
-    generator = heatmap_generator.HeatmapGenerator(self.perf_report,
-                                                   self.page_size, self.title)
-    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)
-
-  def RemoveFiles(self):
-    shutil.rmtree(self.tempDir)
-    if os.path.isfile(os.path.join(os.getcwd(), 'out.txt')):
-      os.remove(os.path.join(os.getcwd(), 'out.txt'))
-    if os.path.isfile(os.path.join(os.getcwd(), 'inst-histo.txt')):
-      os.remove(os.path.join(os.getcwd(), 'inst-histo.txt'))
-
-
-def main(argv):
-  """Parse the options.
-
-  Args:
-    argv: The options with which this script was invoked.
-
-  Returns:
-    0 unless an exception is raised.
-  """
-  parser = argparse.ArgumentParser()
-
-  parser.add_argument(
-      '--chromeos_root',
-      dest='chromeos_root',
-      required=True,
-      help='ChromeOS root to use for generate heatmaps.')
-  parser.add_argument(
-      '--perf_data', dest='perf_data', required=True, help='The raw perf data.')
-  parser.add_argument(
-      '--binary',
-      dest='binary',
-      required=False,
-      help='The path to the Chrome binary.',
-      default=None)
-  parser.add_argument(
-      '--top_n',
-      dest='top_n',
-      required=False,
-      help='Print out top N hottest pages within/outside huge page range(30MB)',
-      default=None)
-  parser.add_argument(
-      '--page_size',
-      dest='page_size',
-      required=False,
-      help='The page size for heat maps.',
-      default=4096)
-  parser.add_argument('--title', dest='title', default='')
-
-  options = parser.parse_args(argv)
-
-  if not IsARepoRoot(options.chromeos_root):
-    parser.error('% does not contain .repo dir.' % options.chromeos_root)
-
-  if not os.path.isfile(options.perf_data):
-    parser.error('Cannot find perf_data: %s.' % options.perf_data)
-
-  heatmap_producer = HeatMapProducer(options.chromeos_root, options.perf_data,
-                                     options.page_size, options.binary,
-                                     options.title)
-  try:
-    heatmap_producer.copyFileToChroot()
-    heatmap_producer.getPerfReport()
-    heatmap_producer.getHeatMap(options.top_n)
-    print('\nheat map and time histgram genereated in the current directory '
-          'with name heat_map.png and timeline.png accordingly.')
-  except RuntimeError, e:
-    print(e)
-  finally:
-    heatmap_producer.RemoveFiles()
-
-
-if __name__ == '__main__':
-  sys.exit(main(sys.argv[1:]))
diff --git a/heatmaps/heat_map.py b/heatmaps/heat_map.py
new file mode 100755
index 0000000..b7006e3
--- /dev/null
+++ b/heatmaps/heat_map.py
@@ -0,0 +1,184 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Wrapper to generate heat maps for chrome."""
+
+from __future__ import print_function
+
+import argparse
+import os
+import shutil
+import sys
+import tempfile
+
+from cros_utils import command_executer
+import heatmap_generator
+
+
+def IsARepoRoot(directory):
+  """Returns True if directory is the root of a repo checkout."""
+  return os.path.exists(
+      os.path.join(os.path.realpath(os.path.expanduser(directory)), '.repo'))
+
+
+class HeatMapProducer(object):
+  """Class to produce heat map."""
+
+  def __init__(self,
+               chromeos_root,
+               perf_data,
+               hugepage,
+               binary,
+               title,
+               logger=None):
+    self.chromeos_root = os.path.realpath(os.path.expanduser(chromeos_root))
+    self.perf_data = os.path.realpath(os.path.expanduser(perf_data))
+    self.hugepage = hugepage
+    self.dir = os.path.dirname(os.path.realpath(__file__))
+    self.binary = binary
+    self.ce = command_executer.GetCommandExecuter()
+    self.temp_dir = ''
+    self.temp_perf_inchroot = ''
+    self.temp_dir_created = False
+    self.perf_report = ''
+    self.title = title
+    self.logger = logger
+
+  def _EnsureFileInChroot(self):
+    chroot_prefix = os.path.join(self.chromeos_root, 'chroot')
+    if self.perf_data.startswith(chroot_prefix):
+      # If the path to perf_data starts with the same chromeos_root, assume
+      # it's in the chromeos_root so no need for temporary directory and copy.
+      self.temp_dir = self.perf_data.replace('perf.data', '')
+      self.temp_perf_inchroot = self.temp_dir.replace(chroot_prefix, '')
+
+    else:
+      # Otherwise, create a temporary directory and copy perf.data into chroot.
+      self.temp_dir = tempfile.mkdtemp(
+          prefix=os.path.join(self.chromeos_root, 'src/'))
+      temp_perf = os.path.join(self.temp_dir, 'perf.data')
+      shutil.copy2(self.perf_data, temp_perf)
+      self.temp_perf_inchroot = os.path.join('~/trunk/src',
+                                             os.path.basename(self.temp_dir))
+      self.temp_dir_created = True
+
+  def _GeneratePerfReport(self):
+    cmd = ('cd %s && perf report -D -i perf.data > perf_report.txt' %
+           self.temp_perf_inchroot)
+    retval = self.ce.ChrootRunCommand(self.chromeos_root, cmd)
+    if retval:
+      raise RuntimeError('Failed to generate perf report')
+    self.perf_report = os.path.join(self.temp_dir, 'perf_report.txt')
+
+  def _GetHeatMap(self, top_n_pages):
+    generator = heatmap_generator.HeatmapGenerator(
+        perf_report=self.perf_report,
+        page_size=4096,
+        hugepage=self.hugepage,
+        title=self.title)
+    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)
+
+  def _RemoveFiles(self):
+    files = [
+        'out.txt', 'inst-histo.txt', 'inst-histo-hp.txt', 'inst-histo-sp.txt'
+    ]
+    for f in files:
+      if os.path.exists(f):
+        os.remove(f)
+
+  def Run(self, top_n_pages=None):
+    try:
+      self._EnsureFileInChroot()
+      self._GeneratePerfReport()
+      self._GetHeatMap(top_n_pages)
+    finally:
+      self._RemoveFiles()
+    msg = ('heat map and time histogram genereated in the current '
+           'directory with name heat_map.png and timeline.png '
+           'accordingly.')
+    if self.logger:
+      self.logger.LogOutput(msg)
+    else:
+      print(msg)
+
+
+def main(argv):
+  """Parse the options.
+
+  Args:
+    argv: The options with which this script was invoked.
+
+  Returns:
+    0 unless an exception is raised.
+  """
+  parser = argparse.ArgumentParser()
+
+  parser.add_argument(
+      '--chromeos_root',
+      dest='chromeos_root',
+      required=True,
+      help='ChromeOS root to use for generate heatmaps.')
+  parser.add_argument(
+      '--perf_data',
+      dest='perf_data',
+      required=True,
+      help='The raw perf data. Must be collected with -e instructions while '
+      'disabling ASLR.')
+  parser.add_argument(
+      '--binary',
+      dest='binary',
+      help='The path to the Chrome binary. Only useful if want to print '
+      'symbols on hottest pages',
+      default=None)
+  parser.add_argument(
+      '--top_n',
+      dest='top_n',
+      help='Print out top N hottest pages within/outside huge page range. '
+      'Must be used with --hugepage and --binary.',
+      default=None)
+  parser.add_argument(
+      '--title', dest='title', help='Title of the heatmap', default='')
+  parser.add_argument(
+      '--hugepage',
+      dest='hugepage',
+      help='A range of addresses (start,end) where huge page starts and ends'
+      ' in text section, separated by a comma.'
+      ' Used to differentiate regions in heatmap.'
+      ' Example: --hugepage=0,4096'
+      ' If not specified, no effect on the heatmap.',
+      default=None)
+
+  options = parser.parse_args(argv)
+
+  if not IsARepoRoot(options.chromeos_root):
+    parser.error('%s does not contain .repo dir.' % options.chromeos_root)
+
+  if not os.path.isfile(options.perf_data):
+    parser.error('Cannot find perf_data: %s.' % options.perf_data)
+
+  hugepage_range = None
+  if options.hugepage:
+    hugepage_range = options.hugepage.split(',')
+    if len(hugepage_range) != 2 or \
+     int(hugepage_range[0]) > int(hugepage_range[1]):
+      parser.error('Wrong format of hugepage range: %s' % options.hugepage)
+    hugepage_range = [int(x) for x in hugepage_range]
+
+  heatmap_producer = HeatMapProducer(options.chromeos_root, options.perf_data,
+                                     hugepage_range, options.binary,
+                                     options.title)
+
+  heatmap_producer.Run(options.top_n)
+
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv[1:]))
diff --git a/heatmaps/heat_map_test.py b/heatmaps/heat_map_test.py
new file mode 100755
index 0000000..c647467
--- /dev/null
+++ b/heatmaps/heat_map_test.py
@@ -0,0 +1,157 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Tests for heat_map.py."""
+
+from __future__ import print_function
+
+import mock
+import unittest
+
+import os
+
+from cros_utils import command_executer
+
+import heat_map
+
+
+def make_heatmap(chromeos_root='/path/to/fake/chromeos_root/',
+                 perf_data='/any_path/perf.data'):
+  return heat_map.HeatMapProducer(chromeos_root, perf_data, None, None, '')
+
+
+def fake_mkdtemp(prefix):
+  """Mock tempfile.mkdtemp() by just create a pathname."""
+  return prefix + 'random_dir'
+
+
+def fake_parser_error(_, msg):
+  """Redirect parser.error() to exception."""
+  raise Exception(msg)
+
+
+def fake_generate_perf_report_exception(_):
+  raise Exception
+
+
+class HeatmapTest(unittest.TestCase):
+  """All of our tests for heat_map."""
+
+  #pylint: disable=protected-access
+  @mock.patch('shutil.copy2')
+  @mock.patch('tempfile.mkdtemp')
+  def test_EnsureFileInChrootAlreadyInside(self, mock_mkdtemp, mock_copy):
+    perf_data_inchroot = (
+        '/path/to/fake/chromeos_root/chroot/inchroot_path/perf.data')
+    heatmap = make_heatmap(perf_data=perf_data_inchroot)
+    heatmap._EnsureFileInChroot()
+    self.assertFalse(heatmap.temp_dir_created)
+    self.assertEqual(heatmap.temp_dir,
+                     '/path/to/fake/chromeos_root/chroot/inchroot_path/')
+    self.assertEqual(heatmap.temp_perf_inchroot, '/inchroot_path/')
+    mock_mkdtemp.assert_not_called()
+    mock_copy.assert_not_called()
+
+  @mock.patch('shutil.copy2')
+  @mock.patch('tempfile.mkdtemp', fake_mkdtemp)
+  def test_EnsureFileInChrootOutsideNeedCopy(self, mock_copy):
+    heatmap = make_heatmap()
+    heatmap._EnsureFileInChroot()
+    self.assertTrue(heatmap.temp_dir_created)
+    self.assertEqual(mock_copy.call_count, 1)
+    self.assertEqual(heatmap.temp_dir,
+                     '/path/to/fake/chromeos_root/src/random_dir')
+    self.assertEqual(heatmap.temp_perf_inchroot, '~/trunk/src/random_dir')
+
+  @mock.patch.object(command_executer.CommandExecuter, 'ChrootRunCommand')
+  def test_GeneratePerfReport(self, mock_ChrootRunCommand):
+    heatmap = make_heatmap()
+    heatmap.temp_dir = '/fake/chroot/inchroot_path/'
+    heatmap.temp_perf_inchroot = '/inchroot_path/'
+    mock_ChrootRunCommand.return_value = 0
+    heatmap._GeneratePerfReport()
+    cmd = ('cd %s && perf report -D -i perf.data > perf_report.txt' %
+           heatmap.temp_perf_inchroot)
+    mock_ChrootRunCommand.assert_called_with(heatmap.chromeos_root, cmd)
+    self.assertEqual(mock_ChrootRunCommand.call_count, 1)
+    self.assertEqual(heatmap.perf_report,
+                     '/fake/chroot/inchroot_path/perf_report.txt')
+
+  @mock.patch('heatmap_generator.HeatmapGenerator')
+  def test_GetHeatMap(self, mock_heatmap_generator):
+    heatmap = make_heatmap()
+    heatmap._GetHeatMap(None)
+    self.assertTrue(mock_heatmap_generator.called)
+
+  @mock.patch.object(heat_map.HeatMapProducer, '_EnsureFileInChroot')
+  @mock.patch.object(heat_map.HeatMapProducer, '_GeneratePerfReport')
+  @mock.patch.object(heat_map.HeatMapProducer, '_GetHeatMap')
+  @mock.patch.object(heat_map.HeatMapProducer, '_RemoveFiles')
+  def test_Run(self, mock_remove_files, mock_get_heatmap,
+               mock_generate_perf_report, mock_ensure_file_in_chroot):
+    heatmap = make_heatmap()
+    heatmap.Run()
+    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_remove_files.assert_called_once_with()
+
+  @mock.patch.object(heat_map.HeatMapProducer, '_EnsureFileInChroot')
+  @mock.patch.object(
+      heat_map.HeatMapProducer,
+      '_GeneratePerfReport',
+      new=fake_generate_perf_report_exception)
+  @mock.patch.object(heat_map.HeatMapProducer, '_GetHeatMap')
+  @mock.patch.object(heat_map.HeatMapProducer, '_RemoveFiles')
+  @mock.patch('__builtin__.print')
+  def test_Run_with_exception(self, mock_print, mock_remove_files,
+                              mock_get_heatmap, mock_ensure_file_in_chroot):
+    heatmap = make_heatmap()
+    with self.assertRaises(Exception):
+      heatmap.Run()
+    mock_ensure_file_in_chroot.assert_called_once_with()
+    mock_get_heatmap.assert_not_called()
+    mock_remove_files.assert_called_once_with()
+    mock_print.assert_not_called()
+
+  @mock.patch('argparse.ArgumentParser.error', fake_parser_error)
+  @mock.patch.object(os.path, 'isfile')
+  @mock.patch.object(heat_map, 'IsARepoRoot')
+  def test_main_arg_format(self, mock_IsARepoRoot, mock_isfile):
+    """Test wrong arg format are detected."""
+    args = ['--chromeos_root=/fake/chroot/', '--perf_data=/path/to/perf.data']
+
+    # Test --chromeos_root format
+    mock_IsARepoRoot.return_value = False
+    with self.assertRaises(Exception) as msg:
+      heat_map.main(args)
+    self.assertIn('does not contain .repo dir.', str(msg.exception))
+
+    # Test --perf_data format
+    mock_IsARepoRoot.return_value = True
+    mock_isfile.return_value = False
+    with self.assertRaises(Exception) as msg:
+      heat_map.main(args)
+    self.assertIn('Cannot find perf_data', str(msg.exception))
+
+    # Test --hugepage format
+    mock_isfile.return_value = True
+    args.append('--hugepage=0')
+    with self.assertRaises(Exception) as msg:
+      heat_map.main(args)
+    self.assertIn('Wrong format of hugepage range', str(msg.exception))
+
+    # Test --hugepage parse
+    args[-1] = '--hugepage=0,4096'
+    heat_map.HeatMapProducer = mock.MagicMock()
+    heat_map.main(args)
+    heat_map.HeatMapProducer.assert_called_with(
+        '/fake/chroot/', '/path/to/perf.data', [0, 4096], None, '')
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/heatmap_generator.py b/heatmaps/heatmap_generator.py
similarity index 93%
rename from heatmap_generator.py
rename to heatmaps/heatmap_generator.py
index 2b7bf92..4f76450 100644
--- a/heatmap_generator.py
+++ b/heatmaps/heatmap_generator.py
@@ -23,6 +23,8 @@
 
 from cros_utils import command_executer
 
+HugepageRange = collections.namedtuple('HugepageRange', ['start', 'end'])
+
 
 class MMap(object):
   """Class to store mmap information in perf report.
@@ -78,7 +80,12 @@
   chrome binary
   """
 
-  def __init__(self, perf_report, page_size, title, log_level='verbose'):
+  def __init__(self,
+               perf_report,
+               page_size,
+               hugepage,
+               title,
+               log_level='verbose'):
     self.perf_report = perf_report
     # Pick 1G as a relatively large number. All addresses less than it will
     # be recorded. The actual heatmap will show up to a boundary of the
@@ -93,13 +100,14 @@
     self.processes = {}
     self.deleted_processes = {}
     self.count = 0
-    self.page_size = page_size
+    if hugepage:
+      self.hugepage = HugepageRange(start=hugepage[0], end=hugepage[1])
+    else:
+      self.hugepage = None
     self.title = title
     self.symbol_addresses = []
     self.symbol_names = []
-    # Set huge page region of Chrome to be first 30MB, only used when printing
-    # out hottest pages
-    self.huge_page_size = 30 * 1024 * 1024
+    self.page_size = page_size
 
   def _parse_perf_sample(self, line):
     # In a perf report, generated with -D, a PERF_RECORD_SAMPLE command should
@@ -268,9 +276,14 @@
     ' the starting address of Chrome'
     if address < self.max_addr:
       self.count += 1
-      print(('%d/%d: %d %d' % (pid[0], pid[1], self.count,
-                               address / self.page_size * self.page_size)),
-            file=self.hist_temp_output)
+      line = '%d/%d: %d %d' % (pid[0], pid[1], self.count,
+                               address / self.page_size * self.page_size)
+      if self.hugepage:
+        if self.hugepage.start <= address < self.hugepage.end:
+          line += ' hugepage'
+        else:
+          line += ' smallpage'
+      print(line, file=self.hist_temp_output)
 
   def _read_perf_report(self):
     # Serve as main function to read perf report, generated by -D
@@ -292,7 +305,12 @@
     # of results written in out.txt and also generate pngs for
     # heat maps.
     heatmap_script = os.path.join(self.dir, 'perf-to-inst-page.sh')
-    cmd = '{0} {1}'.format(heatmap_script, pipes.quote(self.title))
+    if self.hugepage:
+      hp_arg = 'hugepage'
+    else:
+      hp_arg = 'none'
+
+    cmd = '{0} {1} {2}'.format(heatmap_script, pipes.quote(self.title), hp_arg)
     retval = self.ce.RunCommand(cmd)
     if retval:
       raise RuntimeError('Failed to run script to generate heatmap')
@@ -366,11 +384,11 @@
       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 k < self.huge_page_size
-      ][:int(top_n)]
+          (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.huge_page_size
-      ][:int(top_n)]
+          (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
diff --git a/heatmap_generator_test.py b/heatmaps/heatmap_generator_test.py
similarity index 79%
rename from heatmap_generator_test.py
rename to heatmaps/heatmap_generator_test.py
index 53f7dd2..271f9f4 100755
--- a/heatmap_generator_test.py
+++ b/heatmaps/heatmap_generator_test.py
@@ -45,18 +45,21 @@
   print(' ...... dso: /opt/google/chrome/chrome\n', file=fp)
 
 
-def _heatmap(file_name, page_size=4096):
+def _heatmap(file_name, page_size=4096, hugepage=None):
   generator = heatmap_generator.HeatmapGenerator(
-      file_name, page_size, '', log_level='none')  # Don't log to stdout
+      file_name, page_size, hugepage, '',
+      log_level='none')  # Don't log to stdout
   generator.draw()
 
 
 def _cleanup(file_name):
-  os.remove(file_name)
-  os.remove('out.txt')
-  os.remove('inst-histo.txt')
-  os.remove('heat_map.png')
-  os.remove('timeline.png')
+  files = [
+      file_name, 'out.txt', 'inst-histo.txt', 'inst-histo-hp.txt',
+      'inst-histo-sp.txt', 'heat_map.png', 'timeline.png'
+  ]
+  for f in files:
+    if os.path.exists(f):
+      os.remove(f)
 
 
 class Tests(unittest.TestCase):
@@ -68,13 +71,13 @@
     with open(fname, 'w') as f:
       _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
       _write_perf_sample(101, 101, 0xABCD101, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
       lines = f.readlines()
       self.assertEqual(len(lines), 1)
       self.assertIn('101/101: 1 0', lines[0])
-    _cleanup(fname)
 
   def test_with_one_mmap_multiple_samples(self):
     """Tests one perf record and three samples."""
@@ -84,6 +87,7 @@
       _write_perf_sample(101, 101, 0xABCD101, f)
       _write_perf_sample(101, 101, 0xABCD102, f)
       _write_perf_sample(101, 101, 0xABCE102, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
@@ -92,7 +96,6 @@
       self.assertIn('101/101: 1 0', lines[0])
       self.assertIn('101/101: 2 0', lines[1])
       self.assertIn('101/101: 3 4096', lines[2])
-    _cleanup(fname)
 
   def test_with_fork_and_exit(self):
     """Tests perf fork and perf exit."""
@@ -103,6 +106,7 @@
       _write_perf_sample(101, 101, 0xABCD101, f)
       _write_perf_sample(202, 202, 0xABCE101, f)
       _write_perf_exit(202, 202, 202, 202, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
@@ -110,7 +114,6 @@
       self.assertEqual(len(lines), 2)
       self.assertIn('101/101: 1 0', lines[0])
       self.assertIn('202/202: 2 4096', lines[1])
-    _cleanup(fname)
 
   def test_hugepage_creates_two_chrome_mmaps(self):
     """Test two chrome mmaps for the same process."""
@@ -122,6 +125,7 @@
       _write_perf_mmap(202, 202, 0xABCD300, 0xD00, f)
       _write_perf_sample(101, 101, 0xABCD102, f)
       _write_perf_sample(202, 202, 0xABCD102, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
@@ -129,7 +133,6 @@
       self.assertEqual(len(lines), 2)
       self.assertIn('101/101: 1 0', lines[0])
       self.assertIn('202/202: 2 0', lines[1])
-    _cleanup(fname)
 
   def test_hugepage_creates_two_chrome_mmaps_fail(self):
     """Test two chrome mmaps for the same process."""
@@ -139,6 +142,7 @@
       _write_perf_mmap(101, 101, 0xABCD000, 0x1000, f)
       _write_perf_fork(101, 101, 202, 202, f)
       _write_perf_mmap(202, 202, 0xABCD000, 0x10000, f)
+    self.addCleanup(_cleanup, fname)
     with self.assertRaises(AssertionError) as msg:
       _heatmap(fname)
     self.assertIn('Original MMAP size', str(msg.exception))
@@ -173,11 +177,12 @@
     fname = 'test_histo.txt'
     with open(fname, 'w') as f:
       _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
-      for i in range(0, 100):
+      for i in range(100):
         _write_perf_sample(101, 101, 0xABCD000 + i, f)
         _write_perf_sample(101, 101, 0xABCE000 + i, f)
         _write_perf_sample(101, 101, 0xABFD000 + i, f)
         _write_perf_sample(101, 101, 0xAFCD000 + i, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('inst-histo.txt', os.listdir('.'))
     with open('inst-histo.txt') as f:
@@ -187,18 +192,18 @@
       self.assertIn('100 4096', lines[1])
       self.assertIn('100 196608', lines[2])
       self.assertIn('100 4194304', lines[3])
-    _cleanup(fname)
 
   def test_histogram_two_mb_page(self):
     """Tests handling of 2MB page."""
     fname = 'test_histo.txt'
     with open(fname, 'w') as f:
       _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
-      for i in range(0, 100):
+      for i in range(100):
         _write_perf_sample(101, 101, 0xABCD000 + i, f)
         _write_perf_sample(101, 101, 0xABCE000 + i, f)
         _write_perf_sample(101, 101, 0xABFD000 + i, f)
         _write_perf_sample(101, 101, 0xAFCD000 + i, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname, page_size=2 * 1024 * 1024)
     self.assertIn('inst-histo.txt', os.listdir('.'))
     with open('inst-histo.txt') as f:
@@ -206,7 +211,33 @@
       self.assertEqual(len(lines), 2)
       self.assertIn('300 0', lines[0])
       self.assertIn('100 4194304', lines[1])
-    _cleanup(fname)
+
+  def test_histogram_in_and_out_hugepage(self):
+    """Tests handling the case of separating samples in and out huge page."""
+    fname = 'test_histo.txt'
+    with open(fname, 'w') as f:
+      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
+      for i in range(100):
+        _write_perf_sample(101, 101, 0xABCD000 + i, f)
+        _write_perf_sample(101, 101, 0xABCE000 + i, f)
+        _write_perf_sample(101, 101, 0xABFD000 + i, f)
+        _write_perf_sample(101, 101, 0xAFCD000 + i, f)
+    self.addCleanup(_cleanup, fname)
+    _heatmap(fname, hugepage=[0, 8192])
+    file_list = os.listdir('.')
+    self.assertNotIn('inst-histo.txt', file_list)
+    self.assertIn('inst-histo-hp.txt', file_list)
+    self.assertIn('inst-histo-sp.txt', file_list)
+    with open('inst-histo-hp.txt') as f:
+      lines = f.readlines()
+      self.assertEqual(len(lines), 2)
+      self.assertIn('100 0', lines[0])
+      self.assertIn('100 4096', lines[1])
+    with open('inst-histo-sp.txt') as f:
+      lines = f.readlines()
+      self.assertEqual(len(lines), 2)
+      self.assertIn('100 196608', lines[0])
+      self.assertIn('100 4194304', lines[1])
 
 
 if __name__ == '__main__':
diff --git a/heatmaps/perf-to-inst-page.sh b/heatmaps/perf-to-inst-page.sh
new file mode 100755
index 0000000..d6acd5e
--- /dev/null
+++ b/heatmaps/perf-to-inst-page.sh
@@ -0,0 +1,68 @@
+#! /bin/bash -u
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+# This script takes the out.txt, generated by heatmap_generator.py
+# and sorted into a heatmap data file (inst-histo.txt) and then
+# call gnu plot to draw the heat map and the time map.
+# A heat map shows the overall hotness of instructions being executed
+# while the time map shows the hotness of instruction at different time.
+#
+# Usage:
+# ./perf-to-inst-page.sh HEATMAP_TITLE
+# HEATMAP_TITLE: the title to display on the heatmap
+
+HEAT_PNG="heat_map.png"
+TIMELINE_PNG="timeline.png"
+HEATMAP_TITLE=$1
+ENABLE_HUGEPAGE=$2
+
+heatmap_command="
+set terminal png size 600,450
+set xlabel \"Instruction Virtual Address (MB)\"
+set ylabel \"Sample Occurance\"
+set grid
+
+set output \"${HEAT_PNG}\"
+set title \"${HEATMAP_TITLE}\"
+"
+
+if [[ "${ENABLE_HUGEPAGE}" = "hugepage" ]]; then
+  hugepage_hist="inst-histo-hp.txt"
+  smallpage_hist="inst-histo-sp.txt"
+  cat out.txt | grep hugepage | awk '{print $3}' \
+    | sort -n | uniq -c > "${hugepage_hist}"
+  cat out.txt | grep smallpage | awk '{print $3}' \
+    | sort -n | uniq -c > "${smallpage_hist}"
+  # generate inst heat map
+  heatmap_in_hugepage=("'${hugepage_hist}' using \
+(\$2/1024/1024):1 with impulses notitle lt rgb 'red'")
+  heatmap_outside_hugepage=("'${smallpage_hist}' using \
+(\$2/1024/1024):1 with impulses notitle lt rgb 'blue'")
+  echo "
+  ${heatmap_command}
+  plot ${heatmap_in_hugepage}, ${heatmap_outside_hugepage}
+  " | gnuplot
+else
+  awk '{print $3}' out.txt | sort -n | uniq -c > inst-histo.txt
+  # generate inst heat map
+  echo "
+  ${heatmap_command}
+  plot 'inst-histo.txt' using (\$2/1024/1024):1 with impulses notitle
+  " | gnuplot
+fi
+
+# generate instruction page access timeline
+num=$(awk 'END {print NR+1}' out.txt)
+
+echo "
+set terminal png size 600,450
+set xlabel \"time (sec)\"
+set ylabel \"Instruction Virtual Address (MB)\"
+
+set output \"${TIMELINE_PNG}\"
+set title \"instruction page accessd timeline\"
+
+plot 'out.txt' using (\$0/$num*10):(\$3/1024/1024) with dots notitle
+" | gnuplot
diff --git a/perf-to-inst-page.sh b/perf-to-inst-page.sh
deleted file mode 100755
index ea6e43c..0000000
--- a/perf-to-inst-page.sh
+++ /dev/null
@@ -1,47 +0,0 @@
-#! /bin/bash -u
-# Copyright 2015 The Chromium OS Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-# This script takes the out.txt, generated by heatmap_generator.py
-# and sorted into a heatmap data file (inst-histo.txt) and then
-# call gnu plot to draw the heat map and the time map.
-# A heat map shows the overall hotness of instructions being executed
-# while the time map shows the hotness of instruction at different time.
-#
-# Usage:
-# ./perf-to-inst-page.sh HEATMAP_TITLE
-# HEATMAP_TITLE: the title to display on the heatmap
-
-HEAT_PNG="heat_map.png"
-TIMELINE_PNG="timeline.png"
-HEATMAP_TITLE=$1
-
-awk '{print $3}' out.txt | sort -n | uniq -c > inst-histo.txt
-
-# generate inst heat map
-echo "
-set terminal png size 600,450
-set xlabel \"Instruction Virtual Address (MB)\"
-set ylabel \"Sample Occurance\"
-set grid
-
-set output \"${HEAT_PNG}\"
-set title \"${HEATMAP_TITLE}\"
-
-plot 'inst-histo.txt' using (\$2/1024/1024):1 with impulses notitle
-" | gnuplot
-
-# generate instruction page access timeline
-num=$(awk 'END {print NR+1}' out.txt)
-
-echo "
-set terminal png size 600,450
-set xlabel \"time (sec)\"
-set ylabel \"Instruction Virtual Address (MB)\"
-
-set output \"${TIMELINE_PNG}\"
-set title \"instruction page accessd timeline\"
-
-plot 'out.txt' using (\$0/$num*10):(\$3/1024/1024) with dots notitle
-" | gnuplot