1#!/usr/bin/env python3
2
3import atexit
4import argparse
5import datetime
6import http.server
7import os
8import shutil
9import socketserver
10import subprocess
11import sys
12import time
13import webbrowser
14
15ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
16
17# This is not required. It's only used as a fallback if no adb is found on the
18# PATH. It's fine if it doesn't exist so this script can be copied elsewhere.
19HERMETIC_ADB_PATH = ROOT_DIR + '/buildtools/android_sdk/platform-tools/adb'
20
21devnull = open(os.devnull, 'rb')
22adb_path = None
23procs = []
24
25
26class ANSI:
27  END = '\033[0m'
28  BOLD = '\033[1m'
29  RED = '\033[91m'
30  BLACK = '\033[30m'
31  BLUE = '\033[94m'
32  BG_YELLOW = '\033[43m'
33  BG_BLUE = '\033[44m'
34
35
36# HTTP Server used to open the trace in the browser.
37class HttpHandler(http.server.SimpleHTTPRequestHandler):
38
39  def end_headers(self):
40    self.send_header('Access-Control-Allow-Origin', '*')
41    return super().end_headers()
42
43  def do_GET(self):
44    self.server.last_request = self.path
45    return super().do_GET()
46
47  def do_POST(self):
48    self.send_error(404, "File not found")
49
50
51def main():
52  atexit.register(kill_all_subprocs_on_exit)
53  default_out_dir_str = '~/traces/'
54  default_out_dir = os.path.expanduser(default_out_dir_str)
55
56  examples = '\n'.join([
57      ANSI.BOLD + 'Examples' + ANSI.END, '  -t 10s -b 32mb sched gfx wm',
58      '  -t 5s sched/sched_switch raw_syscalls/sys_enter raw_syscalls/sys_exit',
59      '  -c /path/to/full-textual-trace.config', '',
60      ANSI.BOLD + 'Long traces' + ANSI.END,
61      'If you want to record a hours long trace and stream it into a file ',
62      'you need to pass a full trace config and set write_into_file = true.',
63      'See https://perfetto.dev/docs/concepts/config#long-traces .'
64  ])
65  parser = argparse.ArgumentParser(
66      epilog=examples, formatter_class=argparse.RawTextHelpFormatter)
67
68  help = 'Output file or directory (default: %s)' % default_out_dir_str
69  parser.add_argument('-o', '--out', default=default_out_dir, help=help)
70
71  help = 'Don\'t open in the browser'
72  parser.add_argument('-n', '--no-open', action='store_true', help=help)
73
74  grp = parser.add_argument_group(
75      'Short options: (only when not using -c/--config)')
76
77  help = 'Trace duration N[s,m,h] (default: trace until stopped)'
78  grp.add_argument('-t', '--time', default='0s', help=help)
79
80  help = 'Ring buffer size N[mb,gb] (default: 32mb)'
81  grp.add_argument('-b', '--buffer', default='32mb', help=help)
82
83  help = 'Android (atrace) app names (can be specified multiple times)'
84  grp.add_argument(
85      '-a',
86      '--app',
87      metavar='Atrace apps',
88      action='append',
89      default=[],
90      help=help)
91
92  help = 'sched, gfx, am, wm (see --list)'
93  grp.add_argument('events', metavar='Atrace events', nargs='*', help=help)
94
95  help = 'sched/sched_switch kmem/kmem (see --list-ftrace)'
96  grp.add_argument('_', metavar='Ftrace events', nargs='*', help=help)
97
98  help = 'Lists all the categories available'
99  grp.add_argument('--list', action='store_true', help=help)
100
101  help = 'Lists all the ftrace events available'
102  grp.add_argument('--list-ftrace', action='store_true', help=help)
103
104  section = ('Full trace config (only when not using short options)')
105  grp = parser.add_argument_group(section)
106
107  help = 'Can be generated with https://ui.perfetto.dev/#!/record'
108  grp.add_argument('-c', '--config', default=None, help=help)
109  args = parser.parse_args()
110
111  tstamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')
112  fname = '%s.pftrace' % tstamp
113  device_file = '/data/misc/perfetto-traces/' + fname
114
115  find_adb()
116
117  if args.list:
118    adb('shell', 'atrace', '--list_categories').wait()
119    sys.exit(0)
120
121  if args.list_ftrace:
122    adb('shell', 'cat /d/tracing/available_events | tr : /').wait()
123    sys.exit(0)
124
125  if args.config is not None and not os.path.exists(args.config):
126    prt('Config file not found: %s' % args.config, ANSI.RED)
127    sys.exit(1)
128
129  if len(args.events) == 0 and args.config is None:
130    prt('Must either pass short options (e.g. -t 10s sched) or a --config file',
131        ANSI.RED)
132    parser.print_help()
133    sys.exit(1)
134
135  if args.config is None and args.events and os.path.exists(args.events[0]):
136    prt(('The passed event name "%s" is a local file. ' % args.events[0] +
137         'Did you mean to pass -c / --config ?'), ANSI.RED)
138    sys.exit(1)
139
140  cmd = ['perfetto', '--background', '--txt', '-o', device_file]
141  if args.config is not None:
142    cmd += ['-c', '-']
143  else:
144    cmd += ['-t', args.time, '-b', args.buffer]
145    for app in args.app:
146      cmd += ['--app', app]
147    cmd += args.events
148
149  # Perfetto will error out with a proper message if both a config file and
150  # short options are specified. No need to replicate that logic.
151
152  # Work out the output file or directory.
153  if args.out.endswith('/') or os.path.isdir(args.out):
154    host_dir = args.out
155    host_file = os.path.join(args.out, fname)
156  else:
157    host_file = args.out
158    host_dir = os.path.dirname(host_file)
159    if host_dir == '':
160      host_dir = '.'
161      host_file = './' + host_file
162  if not os.path.exists(host_dir):
163    shutil.os.makedirs(host_dir)
164
165  with open(args.config or os.devnull, 'rb') as f:
166    print('Running ' + ' '.join(cmd))
167    proc = adb('shell', *cmd, stdin=f, stdout=subprocess.PIPE)
168    bg_pid = proc.communicate()[0].decode().strip()
169    exit_code = proc.wait()
170
171  if exit_code != 0:
172    prt('Perfetto invocation failed', ANSI.RED)
173    sys.exit(1)
174
175  prt('Trace started. Press CTRL+C to stop', ANSI.BLACK + ANSI.BG_BLUE)
176  logcat = adb('logcat', '-v', 'brief', '-s', 'perfetto', '-b', 'main', '-T',
177               '1')
178
179  ctrl_c_count = 0
180  while ctrl_c_count < 2:
181    try:
182      poll = adb('shell', 'test -d /proc/' + bg_pid)
183      if poll.wait() != 0:
184        break
185      time.sleep(0.5)
186    except KeyboardInterrupt:
187      sig = 'TERM' if ctrl_c_count == 0 else 'KILL'
188      ctrl_c_count += 1
189      prt('Stopping the trace (SIG%s)' % sig, ANSI.BLACK + ANSI.BG_YELLOW)
190      res = adb('shell', 'kill -%s %s' % (sig, bg_pid)).wait()
191
192  logcat.kill()
193  logcat.wait()
194
195  prt('\n')
196  prt('Pulling into %s' % host_file, ANSI.BOLD)
197  adb('pull', device_file, host_file).wait()
198
199  if not args.no_open:
200    prt('\n')
201    prt('Opening the trace (%s) in the browser' % host_file)
202    open_trace_in_browser(host_file)
203
204
205def prt(msg, colors=ANSI.END):
206  print(colors + msg + ANSI.END)
207
208
209def find_adb():
210  """ Locate the "right" adb path
211
212  If adb is in the PATH use that (likely what the user wants) otherwise use the
213  hermetic one in our SDK copy.
214  """
215  global adb_path
216  for path in ['adb', HERMETIC_ADB_PATH]:
217    try:
218      subprocess.call([path, '--version'], stdout=devnull, stderr=devnull)
219      adb_path = path
220      break
221    except OSError:
222      continue
223  if adb_path is None:
224    sdk_url = 'https://developer.android.com/studio/releases/platform-tools'
225    prt('Could not find a suitable adb binary in the PATH. ', ANSI.RED)
226    prt('You can download adb from %s' % sdk_url, ANSI.RED)
227    sys.exit(1)
228
229
230def open_trace_in_browser(path):
231  # We reuse the HTTP+RPC port because it's the only one allowed by the CSP.
232  PORT = 9001
233  os.chdir(os.path.dirname(path))
234  fname = os.path.basename(path)
235  socketserver.TCPServer.allow_reuse_address = True
236  with socketserver.TCPServer(('127.0.0.1', PORT), HttpHandler) as httpd:
237    webbrowser.open_new_tab(
238            'https://ui.perfetto.dev/#!/?url=http://127.0.0.1:%d/%s' %
239        (PORT, fname))
240    while httpd.__dict__.get('last_request') != '/' + fname:
241      httpd.handle_request()
242
243
244def adb(*args, stdin=devnull, stdout=None):
245  cmd = [adb_path, *args]
246  setpgrp = None
247  if os.name != 'nt':
248    # On Linux/Mac, start a new process group so all child processes are killed
249    # on exit. Unsupported on Windows.
250    setpgrp = lambda: os.setpgrp()
251  proc = subprocess.Popen(cmd, stdin=stdin, stdout=stdout, preexec_fn=setpgrp)
252  procs.append(proc)
253  return proc
254
255
256def kill_all_subprocs_on_exit():
257  for p in [p for p in procs if p.poll() is None]:
258    p.kill()
259
260
261if __name__ == '__main__':
262  sys.exit(main())
263