1#!/usr/bin/env python3
2
3# Copyright (C) 2017 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import argparse
22import atexit
23import hashlib
24import os
25import shutil
26import signal
27import subprocess
28import sys
29import tempfile
30import time
31import uuid
32import platform
33
34
35TRACE_TO_TEXT_SHAS = {
36    'linux': '7e3e10dfb324e31723efd63ac25037856e06eba0',
37    'mac': '21f0f42dd019b4f09addd404a114fbf2322ca8a4',
38}
39TRACE_TO_TEXT_PATH = tempfile.gettempdir()
40TRACE_TO_TEXT_BASE_URL = ('https://storage.googleapis.com/perfetto/')
41
42NULL = open(os.devnull)
43NOOUT = {
44    'stdout': NULL,
45    'stderr': NULL,
46}
47
48UUID = str(uuid.uuid4())[-6:]
49
50def check_hash(file_name, sha_value):
51  file_hash = hashlib.sha1()
52  with open(file_name, 'rb') as fd:
53    while True:
54      chunk = fd.read(4096)
55      if not chunk:
56        break
57      file_hash.update(chunk)
58    return file_hash.hexdigest() == sha_value
59
60
61def load_trace_to_text(os_name):
62  sha_value = TRACE_TO_TEXT_SHAS[os_name]
63  file_name = 'trace_to_text-' + os_name + '-' + sha_value
64  local_file = os.path.join(TRACE_TO_TEXT_PATH, file_name)
65
66  if os.path.exists(local_file):
67    if not check_hash(local_file, sha_value):
68      os.remove(local_file)
69    else:
70      return local_file
71
72  url = TRACE_TO_TEXT_BASE_URL + file_name
73  subprocess.check_call(['curl', '-L', '-#', '-o', local_file, url])
74  if not check_hash(local_file, sha_value):
75    os.remove(local_file)
76    raise ValueError("Invalid signature.")
77  os.chmod(local_file, 0o755)
78  return local_file
79
80
81PACKAGES_LIST_CFG = '''data_sources {
82  config {
83    name: "android.packages_list"
84  }
85}
86'''
87
88CFG_INDENT = '      '
89CFG = '''buffers {{
90  size_kb: 63488
91}}
92
93data_sources {{
94  config {{
95    name: "android.heapprofd"
96    heapprofd_config {{
97      shmem_size_bytes: {shmem_size}
98      sampling_interval_bytes: {interval}
99{target_cfg}
100    }}
101  }}
102}}
103
104duration_ms: {duration}
105write_into_file: true
106flush_timeout_ms: 30000
107flush_period_ms: 604800000
108'''
109
110# flush_period_ms of 1 week to suppress trace_processor_shell warning.
111
112CONTINUOUS_DUMP = """
113      continuous_dump_config {{
114        dump_phase_ms: 0
115        dump_interval_ms: {dump_interval}
116      }}
117"""
118
119PROFILE_LOCAL_PATH = os.path.join(tempfile.gettempdir(), UUID)
120
121IS_INTERRUPTED = False
122
123def sigint_handler(sig, frame):
124  global IS_INTERRUPTED
125  IS_INTERRUPTED = True
126
127
128def print_no_profile_error():
129  print("No profiles generated", file=sys.stderr)
130  print(
131    "If this is unexpected, check "
132    "https://perfetto.dev/docs/data-sources/native-heap-profiler#troubleshooting.",
133    file=sys.stderr)
134
135def known_issues_url(number):
136  return ('https://perfetto.dev/docs/data-sources/native-heap-profiler'
137          '#known-issues-android{}'.format(number))
138
139KNOWN_ISSUES = {
140  '10': known_issues_url(10),
141  'Q': known_issues_url(10),
142  '11': known_issues_url(11),
143  'R': known_issues_url(11),
144}
145
146def maybe_known_issues():
147  release_or_codename = subprocess.check_output(
148    ['adb', 'shell', 'getprop', 'ro.build.version.release_or_codename']
149  ).decode('utf-8').strip()
150  return KNOWN_ISSUES.get(release_or_codename, None)
151
152SDK = {
153    'R': 30,
154}
155
156def release_or_newer(release):
157  sdk = int(subprocess.check_output(
158    ['adb', 'shell', 'getprop', 'ro.system.build.version.sdk']
159  ).decode('utf-8').strip())
160  if sdk >= SDK[release]:
161    return True
162  codename = subprocess.check_output(
163    ['adb', 'shell', 'getprop', 'ro.build.version.codename']
164  ).decode('utf-8').strip()
165  return codename == release
166
167def main(argv):
168  parser = argparse.ArgumentParser()
169  parser.add_argument(
170      "-i",
171      "--interval",
172      help="Sampling interval. "
173      "Default 4096 (4KiB)",
174      type=int,
175      default=4096)
176  parser.add_argument(
177      "-d",
178      "--duration",
179      help="Duration of profile (ms). 0 to run until interrupted. "
180      "Default: until interrupted by user.",
181      type=int,
182      default=0)
183  # This flag is a no-op now. We never start heapprofd explicitly using system
184  # properties.
185  parser.add_argument(
186      "--no-start", help="Do not start heapprofd.", action='store_true')
187  parser.add_argument(
188      "-p",
189      "--pid",
190      help="Comma-separated list of PIDs to "
191      "profile.",
192      metavar="PIDS")
193  parser.add_argument(
194      "-n",
195      "--name",
196      help="Comma-separated list of process "
197      "names to profile.",
198      metavar="NAMES")
199  parser.add_argument(
200      "-c",
201      "--continuous-dump",
202      help="Dump interval in ms. 0 to disable continuous dump.",
203      type=int,
204      default=0)
205  parser.add_argument(
206      "--heaps",
207      help="Comma-separated list of heaps to collect, e.g: malloc,art. "
208      "Requires Android 12.",
209      metavar="HEAPS")
210  parser.add_argument(
211      "--all-heaps",
212      action="store_true",
213      help="Collect allocations from all heaps registered by target."
214  )
215  parser.add_argument(
216      "--no-android-tree-symbolization",
217      action="store_true",
218      help="Do not symbolize using currently lunched target in the "
219      "Android tree."
220  )
221  parser.add_argument(
222      "--disable-selinux",
223      action="store_true",
224      help="Disable SELinux enforcement for duration of "
225      "profile.")
226  parser.add_argument(
227      "--no-versions",
228      action="store_true",
229      help="Do not get version information about APKs.")
230  parser.add_argument(
231      "--no-running",
232      action="store_true",
233      help="Do not target already running processes. Requires Android 11.")
234  parser.add_argument(
235      "--no-startup",
236      action="store_true",
237      help="Do not target processes that start during "
238      "the profile. Requires Android 11.")
239  parser.add_argument(
240      "--shmem-size",
241      help="Size of buffer between client and "
242      "heapprofd. Default 8MiB. Needs to be a power of two "
243      "multiple of 4096, at least 8192.",
244      type=int,
245      default=8 * 1048576)
246  parser.add_argument(
247      "--block-client",
248      help="When buffer is full, block the "
249      "client to wait for buffer space. Use with caution as "
250      "this can significantly slow down the client. "
251      "This is the default",
252      action="store_true")
253  parser.add_argument(
254      "--block-client-timeout",
255      help="If --block-client is given, do not block any allocation for "
256      "longer than this timeout (us).",
257      type=int)
258  parser.add_argument(
259      "--no-block-client",
260      help="When buffer is full, stop the "
261      "profile early.",
262      action="store_true")
263  parser.add_argument(
264      "--idle-allocations",
265      help="Keep track of how many "
266      "bytes were unused since the last dump, per "
267      "callstack",
268      action="store_true")
269  parser.add_argument(
270      "--dump-at-max",
271      help="Dump the maximum memory usage "
272      "rather than at the time of the dump.",
273      action="store_true")
274  parser.add_argument(
275      "--disable-fork-teardown",
276      help="Do not tear down client in forks. This can be useful for programs "
277      "that use vfork. Android 11+ only.",
278      action="store_true")
279  parser.add_argument(
280      "--simpleperf",
281      action="store_true",
282      help="Get simpleperf profile of heapprofd. This is "
283      "only for heapprofd development.")
284  parser.add_argument(
285      "--trace-to-text-binary",
286      help="Path to local trace to text. For debugging.")
287  parser.add_argument(
288      "--print-config",
289      action="store_true",
290      help="Print config instead of running. For debugging.")
291  parser.add_argument(
292      "-o",
293      "--output",
294      help="Output directory.",
295      metavar="DIRECTORY",
296      default=None)
297
298  args = parser.parse_args()
299  fail = False
300  if args.block_client and args.no_block_client:
301    print(
302        "FATAL: Both block-client and no-block-client given.", file=sys.stderr)
303    fail = True
304  if args.pid is None and args.name is None:
305    print("FATAL: Neither PID nor NAME given.", file=sys.stderr)
306    fail = True
307  if args.duration is None:
308    print("FATAL: No duration given.", file=sys.stderr)
309    fail = True
310  if args.interval is None:
311    print("FATAL: No interval given.", file=sys.stderr)
312    fail = True
313  if args.shmem_size % 4096:
314    print("FATAL: shmem-size is not a multiple of 4096.", file=sys.stderr)
315    fail = True
316  if args.shmem_size < 8192:
317    print("FATAL: shmem-size is less than 8192.", file=sys.stderr)
318    fail = True
319  if args.shmem_size & (args.shmem_size - 1):
320    print("FATAL: shmem-size is not a power of two.", file=sys.stderr)
321    fail = True
322
323  target_cfg = ""
324  if not args.no_block_client:
325    target_cfg += CFG_INDENT + "block_client: true\n"
326  if args.block_client_timeout:
327    target_cfg += (
328      CFG_INDENT + "block_client_timeout_us: %s\n" % args.block_client_timeout
329    )
330  if args.no_startup:
331    target_cfg += CFG_INDENT + "no_startup: true\n"
332  if args.no_running:
333    target_cfg += CFG_INDENT + "no_running: true\n"
334  if args.dump_at_max:
335    target_cfg += CFG_INDENT + "dump_at_max: true\n"
336  if args.disable_fork_teardown:
337    target_cfg += CFG_INDENT + "disable_fork_teardown: true\n"
338  if args.all_heaps:
339    target_cfg += CFG_INDENT + "all_heaps: true\n"
340  if args.pid:
341    for pid in args.pid.split(','):
342      try:
343        pid = int(pid)
344      except ValueError:
345        print("FATAL: invalid PID %s" % pid, file=sys.stderr)
346        fail = True
347      target_cfg += CFG_INDENT + 'pid: {}\n'.format(pid)
348  if args.name:
349    for name in args.name.split(','):
350      target_cfg += CFG_INDENT + 'process_cmdline: "{}"\n'.format(name)
351  if args.heaps:
352    for heap in args.heaps.split(','):
353      target_cfg += CFG_INDENT + 'heaps: "{}"\n'.format(heap)
354
355  if fail:
356    parser.print_help()
357    return 1
358
359  trace_to_text_binary = args.trace_to_text_binary
360
361  if args.continuous_dump:
362    target_cfg += CONTINUOUS_DUMP.format(dump_interval=args.continuous_dump)
363  cfg = CFG.format(
364      interval=args.interval,
365      duration=args.duration,
366      target_cfg=target_cfg,
367      shmem_size=args.shmem_size)
368  if not args.no_versions:
369    cfg += PACKAGES_LIST_CFG
370
371  if args.print_config:
372    print(cfg)
373    return 0
374
375  # Do this AFTER print_config so we do not download trace_to_text only to
376  # print out the config.
377  has_trace_to_text = True
378  if trace_to_text_binary is None:
379    os_name = None
380    if sys.platform.startswith('linux'):
381      os_name = 'linux'
382    elif sys.platform.startswith('darwin'):
383      os_name = 'mac'
384    elif sys.platform.startswith('win32'):
385      has_trace_to_text = False
386    else:
387      print("Invalid platform: {}".format(sys.platform), file=sys.stderr)
388      return 1
389
390    arch = platform.machine()
391    if arch not in ['x86_64', 'amd64']:
392      has_trace_to_text = False
393
394    if has_trace_to_text:
395      trace_to_text_binary = load_trace_to_text(os_name)
396
397  known_issues = maybe_known_issues()
398  if known_issues:
399    print('If you are experiencing problems, please see the known issues for '
400          'your release: {}.'.format(known_issues))
401
402  # TODO(fmayer): Maybe feature detect whether we can remove traces instead of
403  # this.
404  uuid_trace = release_or_newer('R')
405  if uuid_trace:
406    profile_device_path = '/data/misc/perfetto-traces/profile-' + UUID
407  else:
408    user = subprocess.check_output(
409      ['adb', 'shell', 'whoami']).decode('utf-8').strip()
410    profile_device_path = '/data/misc/perfetto-traces/profile-' + user
411
412  perfetto_cmd = ('CFG=\'{cfg}\'; echo ${{CFG}} | '
413                  'perfetto --txt -c - -o ' + profile_device_path + ' -d')
414
415  if args.disable_selinux:
416    enforcing = subprocess.check_output(['adb', 'shell', 'getenforce'])
417    atexit.register(
418        subprocess.check_call,
419        ['adb', 'shell', 'su root setenforce %s' % enforcing])
420    subprocess.check_call(['adb', 'shell', 'su root setenforce 0'])
421
422  if args.simpleperf:
423    subprocess.check_call([
424        'adb', 'shell', 'mkdir -p /data/local/tmp/heapprofd_profile && '
425        'cd /data/local/tmp/heapprofd_profile &&'
426        '(nohup simpleperf record -g -p $(pidof heapprofd) 2>&1 &) '
427        '> /dev/null'
428    ])
429
430  profile_target = PROFILE_LOCAL_PATH
431  if args.output is not None:
432    profile_target = args.output
433  else:
434    os.mkdir(profile_target)
435
436  if not os.path.isdir(profile_target):
437    print("Output directory {} not found".format(profile_target),
438            file=sys.stderr)
439    return 1
440
441  if os.listdir(profile_target):
442    print("Output directory {} not empty".format(profile_target),
443            file=sys.stderr)
444    return 1
445
446  perfetto_pid = subprocess.check_output(
447      ['adb', 'exec-out',
448       perfetto_cmd.format(cfg=cfg)]).strip()
449  try:
450    perfetto_pid = int(perfetto_pid.strip())
451  except ValueError:
452    print("Failed to invoke perfetto: {}".format(perfetto_pid), file=sys.stderr)
453    return 1
454
455  old_handler = signal.signal(signal.SIGINT, sigint_handler)
456  print("Profiling active. Press Ctrl+C to terminate.")
457  print("You may disconnect your device.")
458  print()
459  exists = True
460  device_connected = True
461  while not device_connected or (exists and not IS_INTERRUPTED):
462    exists = subprocess.call(
463        ['adb', 'shell', '[ -d /proc/{} ]'.format(perfetto_pid)], **NOOUT) == 0
464    device_connected = subprocess.call(['adb', 'shell', 'true'], **NOOUT) == 0
465    time.sleep(1)
466  print("Waiting for profiler shutdown...")
467  signal.signal(signal.SIGINT, old_handler)
468  if IS_INTERRUPTED:
469    # Not check_call because it could have existed in the meantime.
470    subprocess.call(['adb', 'shell', 'kill', '-INT', str(perfetto_pid)])
471  if args.simpleperf:
472    subprocess.check_call(['adb', 'shell', 'killall', '-INT', 'simpleperf'])
473    print("Waiting for simpleperf to exit.")
474    while subprocess.call(
475        ['adb', 'shell', '[ -f /proc/$(pidof simpleperf)/exe ]'], **NOOUT) == 0:
476      time.sleep(1)
477    subprocess.check_call(
478        ['adb', 'pull', '/data/local/tmp/heapprofd_profile', profile_target])
479    print(
480      "Pulled simpleperf profile to " + profile_target + "/heapprofd_profile")
481
482  # Wait for perfetto cmd to return.
483  while exists:
484    exists = subprocess.call(
485        ['adb', 'shell', '[ -d /proc/{} ]'.format(perfetto_pid)]) == 0
486    time.sleep(1)
487
488  profile_host_path = os.path.join(profile_target, 'raw-trace')
489  subprocess.check_call(
490    ['adb', 'pull', profile_device_path, profile_host_path], stdout=NULL)
491  if uuid_trace:
492    subprocess.check_call(
493          ['adb', 'shell', 'rm', profile_device_path], stdout=NULL)
494
495  if not has_trace_to_text:
496    print('Wrote profile to {}'.format(profile_host_path))
497    print('This file can be opened using the Perfetto UI, https://ui.perfetto.dev')
498    return 0
499
500  binary_path = os.getenv('PERFETTO_BINARY_PATH')
501  if not args.no_android_tree_symbolization:
502    product_out = os.getenv('ANDROID_PRODUCT_OUT')
503    if product_out:
504      product_out_symbols = product_out + '/symbols'
505    else:
506      product_out_symbols = None
507
508    if binary_path is None:
509      binary_path = product_out_symbols
510    elif product_out_symbols is not None:
511      binary_path += ":" + product_out_symbols
512
513  trace_file = os.path.join(profile_target, 'raw-trace')
514  concat_files = [trace_file]
515
516  if binary_path is not None:
517    with open(os.path.join(profile_target, 'symbols'), 'w') as fd:
518      ret = subprocess.call([
519          trace_to_text_binary, 'symbolize',
520          os.path.join(profile_target, 'raw-trace')],
521          env=dict(os.environ, PERFETTO_BINARY_PATH=binary_path),
522          stdout=fd)
523    if ret == 0:
524      concat_files.append(os.path.join(profile_target, 'symbols'))
525    else:
526      print("Failed to symbolize. Continuing without symbols.",
527      file=sys.stderr)
528
529  proguard_map = os.getenv('PERFETTO_PROGUARD_MAP')
530  if proguard_map is not None:
531    with open(os.path.join(profile_target, 'deobfuscation-packets'), 'w') as fd:
532      ret = subprocess.call([
533          trace_to_text_binary, 'deobfuscate',
534          os.path.join(profile_target, 'raw-trace')],
535          env=dict(os.environ, PERFETTO_PROGUARD_MAP=proguard_map),
536          stdout=fd)
537    if ret == 0:
538      concat_files.append(
539        os.path.join(profile_target, 'deobfuscation-packets'))
540    else:
541      print("Failed to deobfuscate. Continuing without deobfuscated.",
542      file=sys.stderr)
543
544  if len(concat_files) > 1:
545    with open(os.path.join(profile_target, 'symbolized-trace'), 'wb') as out:
546      for fn in concat_files:
547        with open(fn, 'rb') as inp:
548          while True:
549            buf = inp.read(4096)
550            if not buf:
551              break
552            out.write(buf)
553    trace_file = os.path.join(profile_target, 'symbolized-trace')
554
555  trace_to_text_output = subprocess.check_output(
556      [trace_to_text_binary, 'profile', trace_file])
557  profile_path = None
558  for word in trace_to_text_output.decode('utf-8').split():
559    if 'heap_profile-' in word:
560      profile_path = word
561  if profile_path is None:
562    print_no_profile_error()
563    return 1
564
565  profile_files = os.listdir(profile_path)
566  if not profile_files:
567    print_no_profile_error()
568    return 1
569
570  for profile_file in profile_files:
571    shutil.copy(os.path.join(profile_path, profile_file), profile_target)
572
573  subprocess.check_call(
574      ['gzip'] +
575      [os.path.join(profile_target, x) for x in profile_files])
576
577  symlink_path = None
578  if args.output is None:
579    symlink_path = os.path.join(
580      os.path.dirname(profile_target), "heap_profile-latest")
581    if os.path.lexists(symlink_path):
582      os.unlink(symlink_path)
583    os.symlink(profile_target, symlink_path)
584
585  if symlink_path is not None:
586    print("Wrote profiles to {} (symlink {})".format(
587        profile_target, symlink_path))
588  else:
589    print("Wrote profiles to {}".format(profile_target))
590
591  print("These can be viewed using pprof. Googlers: head to pprof/ and "
592        "upload them.")
593
594
595if __name__ == '__main__':
596  sys.exit(main(sys.argv))
597