1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# Copyright 2018 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""Tests for heatmap_generator.py."""
8
9from __future__ import division, print_function
10
11import unittest.mock as mock
12import unittest
13
14import os
15
16from heatmaps import heatmap_generator
17
18
19def _write_perf_mmap(pid, tid, addr, size, fp):
20  print(
21      '0 0 0 0 PERF_RECORD_MMAP2 %d/%d: '
22      '[%x(%x) @ 0x0 0:0 0 0] '
23      'r-xp /opt/google/chrome/chrome\n' % (pid, tid, addr, size),
24      file=fp)
25
26
27def _write_perf_fork(pid_from, tid_from, pid_to, tid_to, fp):
28  print(
29      '0 0 0 0 PERF_RECORD_FORK(%d:%d):(%d:%d)\n' % (pid_to, tid_to, pid_from,
30                                                     tid_from),
31      file=fp)
32
33
34def _write_perf_exit(pid_from, tid_from, pid_to, tid_to, fp):
35  print(
36      '0 0 0 0 PERF_RECORD_EXIT(%d:%d):(%d:%d)\n' % (pid_to, tid_to, pid_from,
37                                                     tid_from),
38      file=fp)
39
40
41def _write_perf_sample(pid, tid, addr, fp):
42  print(
43      '0 0 0 0 PERF_RECORD_SAMPLE(IP, 0x2): '
44      '%d/%d: %x period: 100000 addr: 0' % (pid, tid, addr),
45      file=fp)
46  print(' ... thread: chrome:%d' % tid, file=fp)
47  print(' ...... dso: /opt/google/chrome/chrome\n', file=fp)
48
49
50def _heatmap(file_name, page_size=4096, hugepage=None, analyze=False, top_n=10):
51  generator = heatmap_generator.HeatmapGenerator(
52      file_name, page_size, hugepage, '',
53      log_level='none')  # Don't log to stdout
54  generator.draw()
55  if analyze:
56    generator.analyze('/path/to/chrome', top_n)
57
58
59def _cleanup(file_name):
60  files = [
61      file_name, 'out.txt', 'inst-histo.txt', 'inst-histo-hp.txt',
62      'inst-histo-sp.txt', 'heat_map.png', 'timeline.png', 'addr2symbol.txt'
63  ]
64  for f in files:
65    if os.path.exists(f):
66      os.remove(f)
67
68
69class HeatmapGeneratorDrawTests(unittest.TestCase):
70  """All of our tests for heatmap_generator.draw() and related."""
71
72  def test_with_one_mmap_one_sample(self):
73    """Tests one perf record and one sample."""
74    fname = 'test.txt'
75    with open(fname, 'w') as f:
76      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
77      _write_perf_sample(101, 101, 0xABCD101, f)
78    self.addCleanup(_cleanup, fname)
79    _heatmap(fname)
80    self.assertIn('out.txt', os.listdir('.'))
81    with open('out.txt') as f:
82      lines = f.readlines()
83      self.assertEqual(len(lines), 1)
84      self.assertIn('101/101: 1 0', lines[0])
85
86  def test_with_one_mmap_multiple_samples(self):
87    """Tests one perf record and three samples."""
88    fname = 'test.txt'
89    with open(fname, 'w') as f:
90      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
91      _write_perf_sample(101, 101, 0xABCD101, f)
92      _write_perf_sample(101, 101, 0xABCD102, f)
93      _write_perf_sample(101, 101, 0xABCE102, f)
94    self.addCleanup(_cleanup, fname)
95    _heatmap(fname)
96    self.assertIn('out.txt', os.listdir('.'))
97    with open('out.txt') as f:
98      lines = f.readlines()
99      self.assertEqual(len(lines), 3)
100      self.assertIn('101/101: 1 0', lines[0])
101      self.assertIn('101/101: 2 0', lines[1])
102      self.assertIn('101/101: 3 4096', lines[2])
103
104  def test_with_fork_and_exit(self):
105    """Tests perf fork and perf exit."""
106    fname = 'test_fork.txt'
107    with open(fname, 'w') as f:
108      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
109      _write_perf_fork(101, 101, 202, 202, f)
110      _write_perf_sample(101, 101, 0xABCD101, f)
111      _write_perf_sample(202, 202, 0xABCE101, f)
112      _write_perf_exit(202, 202, 202, 202, f)
113    self.addCleanup(_cleanup, fname)
114    _heatmap(fname)
115    self.assertIn('out.txt', os.listdir('.'))
116    with open('out.txt') as f:
117      lines = f.readlines()
118      self.assertEqual(len(lines), 2)
119      self.assertIn('101/101: 1 0', lines[0])
120      self.assertIn('202/202: 2 4096', lines[1])
121
122  def test_hugepage_creates_two_chrome_mmaps(self):
123    """Test two chrome mmaps for the same process."""
124    fname = 'test_hugepage.txt'
125    with open(fname, 'w') as f:
126      _write_perf_mmap(101, 101, 0xABCD000, 0x1000, f)
127      _write_perf_fork(101, 101, 202, 202, f)
128      _write_perf_mmap(202, 202, 0xABCD000, 0x100, f)
129      _write_perf_mmap(202, 202, 0xABCD300, 0xD00, f)
130      _write_perf_sample(101, 101, 0xABCD102, f)
131      _write_perf_sample(202, 202, 0xABCD102, f)
132    self.addCleanup(_cleanup, fname)
133    _heatmap(fname)
134    self.assertIn('out.txt', os.listdir('.'))
135    with open('out.txt') as f:
136      lines = f.readlines()
137      self.assertEqual(len(lines), 2)
138      self.assertIn('101/101: 1 0', lines[0])
139      self.assertIn('202/202: 2 0', lines[1])
140
141  def test_hugepage_creates_two_chrome_mmaps_fail(self):
142    """Test two chrome mmaps for the same process."""
143    fname = 'test_hugepage.txt'
144    # Cases where first_mmap.size < second_mmap.size
145    with open(fname, 'w') as f:
146      _write_perf_mmap(101, 101, 0xABCD000, 0x1000, f)
147      _write_perf_fork(101, 101, 202, 202, f)
148      _write_perf_mmap(202, 202, 0xABCD000, 0x10000, f)
149    self.addCleanup(_cleanup, fname)
150    with self.assertRaises(AssertionError) as msg:
151      _heatmap(fname)
152    self.assertIn('Original MMAP size', str(msg.exception))
153
154    # Cases where first_mmap.address > second_mmap.address
155    with open(fname, 'w') as f:
156      _write_perf_mmap(101, 101, 0xABCD000, 0x1000, f)
157      _write_perf_fork(101, 101, 202, 202, f)
158      _write_perf_mmap(202, 202, 0xABCC000, 0x10000, f)
159    with self.assertRaises(AssertionError) as msg:
160      _heatmap(fname)
161    self.assertIn('Original MMAP starting address', str(msg.exception))
162
163    # Cases where first_mmap.address + size <
164    # second_mmap.address + second_mmap.size
165    with open(fname, 'w') as f:
166      _write_perf_mmap(101, 101, 0xABCD000, 0x1000, f)
167      _write_perf_fork(101, 101, 202, 202, f)
168      _write_perf_mmap(202, 202, 0xABCD100, 0x10000, f)
169    with self.assertRaises(AssertionError) as msg:
170      _heatmap(fname)
171    self.assertIn('exceeds the end of original MMAP', str(msg.exception))
172
173  def test_histogram(self):
174    """Tests if the tool can generate correct histogram.
175
176    In the tool, histogram is generated from statistics
177    of perf samples (saved to out.txt). The histogram is
178    generated by perf-to-inst-page.sh and saved to
179    inst-histo.txt. It will be used to draw heat maps.
180    """
181    fname = 'test_histo.txt'
182    with open(fname, 'w') as f:
183      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
184      for i in range(100):
185        _write_perf_sample(101, 101, 0xABCD000 + i, f)
186        _write_perf_sample(101, 101, 0xABCE000 + i, f)
187        _write_perf_sample(101, 101, 0xABFD000 + i, f)
188        _write_perf_sample(101, 101, 0xAFCD000 + i, f)
189    self.addCleanup(_cleanup, fname)
190    _heatmap(fname)
191    self.assertIn('inst-histo.txt', os.listdir('.'))
192    with open('inst-histo.txt') as f:
193      lines = f.readlines()
194      self.assertEqual(len(lines), 4)
195      self.assertIn('100 0', lines[0])
196      self.assertIn('100 4096', lines[1])
197      self.assertIn('100 196608', lines[2])
198      self.assertIn('100 4194304', lines[3])
199
200  def test_histogram_two_mb_page(self):
201    """Tests handling of 2MB page."""
202    fname = 'test_histo.txt'
203    with open(fname, 'w') as f:
204      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
205      for i in range(100):
206        _write_perf_sample(101, 101, 0xABCD000 + i, f)
207        _write_perf_sample(101, 101, 0xABCE000 + i, f)
208        _write_perf_sample(101, 101, 0xABFD000 + i, f)
209        _write_perf_sample(101, 101, 0xAFCD000 + i, f)
210    self.addCleanup(_cleanup, fname)
211    _heatmap(fname, page_size=2 * 1024 * 1024)
212    self.assertIn('inst-histo.txt', os.listdir('.'))
213    with open('inst-histo.txt') as f:
214      lines = f.readlines()
215      self.assertEqual(len(lines), 2)
216      self.assertIn('300 0', lines[0])
217      self.assertIn('100 4194304', lines[1])
218
219  def test_histogram_in_and_out_hugepage(self):
220    """Tests handling the case of separating samples in and out huge page."""
221    fname = 'test_histo.txt'
222    with open(fname, 'w') as f:
223      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
224      for i in range(100):
225        _write_perf_sample(101, 101, 0xABCD000 + i, f)
226        _write_perf_sample(101, 101, 0xABCE000 + i, f)
227        _write_perf_sample(101, 101, 0xABFD000 + i, f)
228        _write_perf_sample(101, 101, 0xAFCD000 + i, f)
229    self.addCleanup(_cleanup, fname)
230    _heatmap(fname, hugepage=[0, 8192])
231    file_list = os.listdir('.')
232    self.assertNotIn('inst-histo.txt', file_list)
233    self.assertIn('inst-histo-hp.txt', file_list)
234    self.assertIn('inst-histo-sp.txt', file_list)
235    with open('inst-histo-hp.txt') as f:
236      lines = f.readlines()
237      self.assertEqual(len(lines), 2)
238      self.assertIn('100 0', lines[0])
239      self.assertIn('100 4096', lines[1])
240    with open('inst-histo-sp.txt') as f:
241      lines = f.readlines()
242      self.assertEqual(len(lines), 2)
243      self.assertIn('100 196608', lines[0])
244      self.assertIn('100 4194304', lines[1])
245
246
247class HeatmapGeneratorAnalyzeTests(unittest.TestCase):
248  """All of our tests for heatmap_generator.analyze() and related."""
249
250  def setUp(self):
251    # Use the same perf report for testing
252    self.fname = 'test_histo.txt'
253    with open(self.fname, 'w') as f:
254      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
255      for i in range(10):
256        _write_perf_sample(101, 101, 0xABCD000 + i, f)
257        _write_perf_sample(101, 101, 0xABCE000 + i, f)
258        _write_perf_sample(101, 101, 0xABFD000 + i, f)
259    self.nm = ('000000000abcd000 t Func1@Page1\n'
260               '000000000abcd001 t Func2@Page1\n'
261               '000000000abcd0a0 t Func3@Page1andFunc1@Page2\n'
262               '000000000abce010 t Func2@Page2\n'
263               '000000000abfd000 t Func1@Page3\n')
264
265  def tearDown(self):
266    _cleanup(self.fname)
267
268  @mock.patch('subprocess.check_output')
269  def test_analyze_hot_pages_with_hp_top(self, mock_nm):
270    """Test if the analyze() can print the top page with hugepage."""
271    mock_nm.return_value = self.nm
272    _heatmap(self.fname, hugepage=[0, 8192], analyze=True, top_n=1)
273    file_list = os.listdir('.')
274    self.assertIn('addr2symbol.txt', file_list)
275    with open('addr2symbol.txt') as f:
276      contents = f.read()
277      self.assertIn('Func2@Page1 : 9', contents)
278      self.assertIn('Func1@Page1 : 1', contents)
279      self.assertIn('Func1@Page3 : 10', contents)
280      # Only displaying one page in hugepage
281      self.assertNotIn('Func3@Page1andFunc1@Page2 : 10', contents)
282
283  @mock.patch('subprocess.check_output')
284  def test_analyze_hot_pages_without_hp_top(self, mock_nm):
285    """Test if the analyze() can print the top page without hugepage."""
286    mock_nm.return_value = self.nm
287    _heatmap(self.fname, analyze=True, top_n=1)
288    file_list = os.listdir('.')
289    self.assertIn('addr2symbol.txt', file_list)
290    with open('addr2symbol.txt') as f:
291      contents = f.read()
292      self.assertIn('Func2@Page1 : 9', contents)
293      self.assertIn('Func1@Page1 : 1', contents)
294      # Only displaying one page
295      self.assertNotIn('Func3@Page1andFunc1@Page2 : 10', contents)
296      self.assertNotIn('Func1@Page3 : 10', contents)
297
298  @mock.patch('subprocess.check_output')
299  def test_analyze_hot_pages_with_hp_top10(self, mock_nm):
300    """Test if the analyze() can print with default top 10."""
301    mock_nm.return_value = self.nm
302    _heatmap(self.fname, analyze=True)
303    # Make sure nm command is called correctly.
304    mock_nm.assert_called_with(['nm', '-n', '/path/to/chrome'])
305    file_list = os.listdir('.')
306    self.assertIn('addr2symbol.txt', file_list)
307    with open('addr2symbol.txt') as f:
308      contents = f.read()
309      self.assertIn('Func2@Page1 : 9', contents)
310      self.assertIn('Func1@Page1 : 1', contents)
311      self.assertIn('Func3@Page1andFunc1@Page2 : 10', contents)
312      self.assertIn('Func1@Page3 : 10', contents)
313
314
315if __name__ == '__main__':
316  unittest.main()
317