python/tools/record_android_trace.py

#!/usr/bin/env python3
# Copyright (C) 2021 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import atexit
import argparse
import datetime
import hashlib
import http.server
import json
import os
import re
import shutil
import signal
import socketserver
import subprocess
import sys
import time
import urllib.parse
import webbrowser

from perfetto.prebuilts.manifests.tracebox import *
from perfetto.prebuilts.perfetto_prebuilts import *
from perfetto.common.repo_utils import *

# This is not required. It's only used as a fallback if no adb is found on the
# PATH. It's fine if it doesn't exist so this script can be copied elsewhere.
HERMETIC_ADB_PATH = repo_dir('buildtools/android_sdk/platform-tools/adb')

# Translates the Android ro.product.cpu.abi into the GN's target_cpu.
ABI_TO_ARCH = {
    'armeabi-v7a': 'arm',
    'arm64-v8a': 'arm64',
    'x86': 'x86',
    'x86_64': 'x64',
}

MAX_ADB_FAILURES = 15  # 2 seconds between retries, 30 seconds total.

devnull = open(os.devnull, 'rb')
adb_path = None
procs = []


class ANSI:
  END = '\033[0m'
  BOLD = '\033[1m'
  RED = '\033[91m'
  BLACK = '\033[30m'
  BLUE = '\033[94m'
  BG_YELLOW = '\033[43m'
  BG_BLUE = '\033[44m'


# HTTP Server used to open the trace in the browser.
class HttpHandler(http.server.SimpleHTTPRequestHandler):

  def end_headers(self):
    self.send_header('Access-Control-Allow-Origin', self.server.allow_origin)
    self.send_header('Cache-Control', 'no-cache')
    super().end_headers()

  def do_GET(self):
    if self.path != '/' + self.server.expected_fname:
      self.send_error(404, "File not found")
      return

    self.server.fname_get_completed = True
    super().do_GET()

  def do_POST(self):
    self.send_error(404, "File not found")


def setup_arguments():
  atexit.register(kill_all_subprocs_on_exit)
  default_out_dir_str = '~/traces/'
  default_out_dir = os.path.expanduser(default_out_dir_str)

  examples = '\n'.join([
      ANSI.BOLD + 'Examples' + ANSI.END, '  -t 10s -b 32mb sched gfx wm -a*',
      '  -t 5s sched/sched_switch raw_syscalls/sys_enter raw_syscalls/sys_exit',
      '  -c /path/to/full-textual-trace.config', '',
      ANSI.BOLD + 'Long traces' + ANSI.END,
      'If you want to record a hours long trace and stream it into a file ',
      'you need to pass a full trace config and set write_into_file = true.',
      'See https://perfetto.dev/docs/concepts/config#long-traces .'
  ])
  parser = argparse.ArgumentParser(
      epilog=examples, formatter_class=argparse.RawTextHelpFormatter)

  help = 'Output file or directory (default: %s)' % default_out_dir_str
  parser.add_argument('-o', '--out', default=default_out_dir, help=help)

  help = 'Don\'t open or serve the trace'
  parser.add_argument('-n', '--no-open', action='store_true', help=help)

  help = 'Don\'t open in browser, but still serve trace (good for remote use)'
  parser.add_argument('--no-open-browser', action='store_true', help=help)

  help = 'The web address used to open trace files'
  parser.add_argument('--origin', default='https://ui.perfetto.dev', help=help)

  help = (
      'Advanced: The URL Params to pass to browser when open the trace file. '
      'For UI startup commands, use --ui-startup-commands instead')
  parser.add_argument('--ui-url-params', nargs='+', default=None, help=help)

  help = (
      'Commands to execute on UI startup in the browser (JSON array of command objects). '
      'Example: \'[{"id":"dev.perfetto.PinTracksByRegex","args":[".*sched.*"]}]\''
  )
  parser.add_argument('--ui-startup-commands', default=None, help=help)

  help = 'Force the use of the sideloaded binaries rather than system daemons'
  parser.add_argument('--sideload', action='store_true', help=help)

  help = ('Sideload the given binary rather than downloading it. ' +
          'Implies --sideload')
  parser.add_argument('--sideload-path', default=None, help=help)

  help = 'Ignores any tracing guardrails which might be used'
  parser.add_argument('--no-guardrails', action='store_true', help=help)

  help = 'Don\'t run `adb root` run as user (only when sideloading)'
  parser.add_argument('-u', '--user', action='store_true', help=help)

  help = 'Specify the ADB device serial'
  parser.add_argument('--serial', '-s', default=None, help=help)

  grp = parser.add_argument_group(
      'Short options: (only when not using -c/--config)')

  help = 'Trace duration N[s,m,h] (default: trace until stopped)'
  grp.add_argument('-t', '--time', default='0s', help=help)

  help = 'Ring buffer size N[mb,gb] (default: 32mb)'
  grp.add_argument('-b', '--buffer', default='32mb', help=help)

  help = ('Android (atrace) app names. Can be specified multiple times.\n-a*' +
          'for all apps (without space between a and * or bash will expand it)')
  grp.add_argument(
      '-a',
      '--app',
      metavar='com.myapp',
      action='append',
      default=[],
      help=help)

  help = 'sched, gfx, am, wm (see --list)'
  grp.add_argument('events', metavar='Atrace events', nargs='*', help=help)

  help = 'sched/sched_switch kmem/kmem (see --list-ftrace)'
  grp.add_argument('_', metavar='Ftrace events', nargs='*', help=help)

  help = 'Lists all the categories available'
  grp.add_argument('--list', action='store_true', help=help)

  help = 'Lists all the ftrace events available'
  grp.add_argument('--list-ftrace', action='store_true', help=help)

  section = ('Full trace config (only when not using short options)')
  grp = parser.add_argument_group(section)

  help = 'Can be generated with https://ui.perfetto.dev/#!/record'
  grp.add_argument('-c', '--config', default=None, help=help)

  help = 'Parse input from --config as binary proto (default: parse as text)'
  grp.add_argument('--bin', action='store_true', help=help)

  help = ('Pass the trace through the trace reporter API. Only works when '
          'using the full trace config (-c) with the reporter package name '
          "'android.perfetto.cts.reporter' and the reporter class name "
          "'android.perfetto.cts.reporter.PerfettoReportService' with the "
          'reporter installed on the device (see '
          'tools/install_test_reporter_app.py).')
  grp.add_argument('--reporter-api', action='store_true', help=help)

  args = parser.parse_args()
  args.sideload = args.sideload or args.sideload_path is not None

  if args.serial:
    os.environ["ANDROID_SERIAL"] = args.serial

  find_adb()

  if args.list:
    adb('shell', 'atrace', '--list_categories').wait()
    sys.exit(0)

  if args.list_ftrace:
    adb('shell', 'cat /d/tracing/available_events | tr : /').wait()
    sys.exit(0)

  if args.config is not None and not os.path.exists(args.config):
    prt('Config file not found: %s' % args.config, ANSI.RED)
    sys.exit(1)

  if len(args.events) == 0 and args.config is None:
    prt('Must either pass short options (e.g. -t 10s sched) or a --config file',
        ANSI.RED)
    parser.print_help()
    sys.exit(1)

  if args.config is None and args.events and os.path.exists(args.events[0]):
    prt(('The passed event name "%s" is a local file. ' % args.events[0] +
         'Did you mean to pass -c / --config ?'), ANSI.RED)
    sys.exit(1)

  if args.reporter_api and not args.config:
    prt('Must pass --config when using --reporter-api', ANSI.RED)
    parser.print_help()
    sys.exit(1)

  return args


class SignalException(Exception):
  pass


def signal_handler(sig, frame):
  raise SignalException('Received signal ' + str(sig))


signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)


def start_trace(args, print_log=True):
  perfetto_cmd = 'perfetto'
  device_dir = '/data/misc/perfetto-traces/'

  # Check the version of android. If too old (< Q) sideload tracebox. Also use
  # use /data/local/tmp as /data/misc/perfetto-traces was introduced only later.
  probe_cmd = 'getprop ro.build.version.sdk; getprop ro.product.cpu.abi; whoami'
  probe = adb('shell', probe_cmd, stdout=subprocess.PIPE)
  lines = probe.communicate()[0].decode().strip().split('\n')
  lines = [x.strip() for x in lines]  # To strip \r(s) on Windows.
  if probe.returncode != 0:
    prt('ADB connection failed', ANSI.RED)
    sys.exit(1)
  api_level = int(lines[0])
  abi = lines[1]
  arch = ABI_TO_ARCH.get(abi)
  if arch is None:
    prt('Unsupported ABI: ' + abi)
    sys.exit(1)
  shell_user = lines[2]
  if api_level < 29 or args.sideload:  # 29: Android Q.
    tracebox_bin = args.sideload_path
    if tracebox_bin is None:
      tracebox_bin = get_perfetto_prebuilt(
          TRACEBOX_MANIFEST, arch='android-' + arch)
    perfetto_cmd = '/data/local/tmp/tracebox'
    exit_code = adb('push', '--sync', tracebox_bin, perfetto_cmd).wait()
    exit_code |= adb('shell', 'chmod 755 ' + perfetto_cmd).wait()
    if exit_code != 0:
      prt('ADB push failed', ANSI.RED)
      sys.exit(1)
    device_dir = '/data/local/tmp/'
    if shell_user != 'root' and not args.user:
      # Run as root if possible as that will give access to more tracing
      # capabilities. Non-root still works, but some ftrace events might not be
      # available.
      adb('root').wait()

  tstamp = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')
  fname = '%s-%s.pftrace' % (tstamp, os.urandom(3).hex())
  device_file = device_dir + fname

  cmd = [perfetto_cmd, '--background']
  if not args.bin:
    cmd.append('--txt')

  if args.no_guardrails:
    cmd.append('--no-guardrails')

  if args.reporter_api:
    # Remove all old reporter files to avoid polluting the file we will extract
    # later.
    adb('shell',
        'rm /sdcard/Android/data/android.perfetto.cts.reporter/files/*').wait()
    cmd.append('--upload')
  else:
    cmd.extend(['-o', device_file])

  on_device_config = None
  on_host_config = None
  if args.config is not None:
    cmd += ['-c', '-']
    if api_level < 24:
      # adb shell does not redirect stdin. Push the config on a temporary file
      # on the device.
      mktmp = adb(
          'shell',
          'mktemp',
          '--tmpdir',
          '/data/local/tmp',
          stdout=subprocess.PIPE)
      on_device_config = mktmp.communicate()[0].decode().strip().strip()
      if mktmp.returncode != 0:
        prt('Failed to create config on device', ANSI.RED)
        sys.exit(1)
      exit_code = adb('push', '--sync', args.config, on_device_config).wait()
      if exit_code != 0:
        prt('Failed to push config on device', ANSI.RED)
        sys.exit(1)
      cmd = ['cat', on_device_config, '|'] + cmd
    else:
      on_host_config = args.config
  else:
    cmd += ['-t', args.time, '-b', args.buffer]
    for app in args.app:
      cmd += ['--app', '\'' + app + '\'']
    cmd += args.events

  # Work out the output file or directory.
  if args.out.endswith('/') or os.path.isdir(args.out):
    host_dir = args.out
    host_file = os.path.join(args.out, fname)
  else:
    host_file = args.out
    host_dir = os.path.dirname(host_file)
    if host_dir == '':
      host_dir = '.'
      host_file = './' + host_file
  if not os.path.exists(host_dir):
    shutil.os.makedirs(host_dir)

  with open(on_host_config or os.devnull, 'rb') as f:
    if print_log:
      print('Running ' + ' '.join(cmd))
    proc = adb('shell', *cmd, stdin=f, stdout=subprocess.PIPE)
    proc_out = proc.communicate()[0].decode().strip()
    if on_device_config is not None:
      adb('shell', 'rm', on_device_config).wait()
    # On older versions of Android (x86_64 emulator running API 22) the output
    # looks like:
    #   WARNING: linker: /data/local/tmp/tracebox: unused DT entry: ...
    #   WARNING: ... (other 2 WARNING: linker: lines)
    #   1234  <-- The actual pid we want.
    match = re.search(r'^(\d+)$', proc_out, re.M)
    if match is None:
      prt('Failed to read the pid from perfetto --background', ANSI.RED)
      prt(proc_out)
      sys.exit(1)
    bg_pid = match.group(1)
    exit_code = proc.wait()

  if exit_code != 0:
    prt('Perfetto invocation failed', ANSI.RED)
    sys.exit(1)

  prt('Trace started. Press CTRL+C to stop', ANSI.BLACK + ANSI.BG_BLUE)
  log_level = "-v"
  if not print_log:
    log_level = "-e"
  logcat = adb('logcat', log_level, 'brief', '-s', 'perfetto', '-b', 'main',
               '-T', '1')

  ctrl_c_count = 0
  adb_failure_count = 0
  while ctrl_c_count < 2:
    try:
      # On older Android devices adbd doesn't propagate the exit code. Hence
      # the RUN/TERM parts.
      poll = adb(
          'shell',
          'test -d /proc/%s && echo RUN || echo TERM' % bg_pid,
          stdout=subprocess.PIPE)
      poll_res = poll.communicate()[0].decode().strip()
      if poll_res == 'TERM':
        break  # Process terminated
      if poll_res == 'RUN':
        # The 'perfetto' cmdline client is still running. If previously we had
        # an ADB error, tell the user now it's all right again.
        if adb_failure_count > 0:
          adb_failure_count = 0
          prt('ADB connection re-established, the trace is still ongoing',
              ANSI.BLUE)
        time.sleep(0.5)
        continue
      # Some ADB error happened. This can happen when tracing soon after boot,
      # before logging in, when adb gets restarted.
      adb_failure_count += 1
      if adb_failure_count >= MAX_ADB_FAILURES:
        prt('Too many unrecoverable ADB failures, bailing out', ANSI.RED)
        sys.exit(1)
      time.sleep(2)
    except (KeyboardInterrupt, SignalException):
      sig = 'TERM' if ctrl_c_count == 0 else 'KILL'
      ctrl_c_count += 1
      if print_log:
        prt('Stopping the trace (SIG%s)' % sig, ANSI.BLACK + ANSI.BG_YELLOW)
      adb('shell', 'kill -%s %s' % (sig, bg_pid)).wait()

  logcat.kill()
  logcat.wait()

  if args.reporter_api:
    if print_log:
      prt('Waiting a few seconds to allow reporter to copy trace')
    time.sleep(5)

    ret = adb(
        'shell',
        'cp /sdcard/Android/data/android.perfetto.cts.reporter/files/* ' +
        device_file).wait()
    if ret != 0:
      prt('Failed to extract reporter trace', ANSI.RED)
      sys.exit(1)

  if print_log:
    prt('\n')
    prt('Pulling into %s' % host_file, ANSI.BOLD)
  adb('pull', device_file, host_file).wait()
  adb('shell', 'rm -f ' + device_file).wait()

  if not args.no_open:
    if print_log:
      prt('\n')
      prt('Opening the trace (%s) in the browser' % host_file)
    open_browser = not args.no_open_browser
    open_trace_in_browser(host_file, open_browser, args.origin,
                          args.ui_url_params, args.ui_startup_commands)

  return host_file


def main():
  args = setup_arguments()
  start_trace(args)


def prt(msg, colors=ANSI.END):
  print(colors + msg + ANSI.END)


def find_adb():
  """ Locate the "right" adb path

  If adb is in the PATH use that (likely what the user wants) otherwise use the
  hermetic one in our SDK copy.
  """
  global adb_path
  for path in ['adb', HERMETIC_ADB_PATH]:
    try:
      subprocess.call([path, '--version'], stdout=devnull, stderr=devnull)
      adb_path = path
      break
    except OSError:
      continue
  if adb_path is None:
    sdk_url = 'https://developer.android.com/studio/releases/platform-tools'
    prt('Could not find a suitable adb binary in the PATH. ', ANSI.RED)
    prt('You can download adb from %s' % sdk_url, ANSI.RED)
    sys.exit(1)


def open_trace_in_browser(path,
                          open_browser,
                          origin,
                          url_params,
                          startup_commands=None):
  # We reuse the HTTP+RPC port because it's the only one allowed by the CSP.
  PORT = 9001
  path = os.path.abspath(path)
  os.chdir(os.path.dirname(path))
  fname = os.path.basename(path)
  socketserver.TCPServer.allow_reuse_address = True
  with socketserver.TCPServer(('127.0.0.1', PORT), HttpHandler) as httpd:
    address = f'{origin}/#!/?url=http://127.0.0.1:{PORT}/{fname}&referrer=record_android_trace'

    # Build URL parameters list
    params = []
    if url_params:
      params.extend(url_params)

    # Add startup commands if provided
    if startup_commands:
      try:
        # Validate that startup_commands is valid JSON
        json.loads(startup_commands)  # Validate JSON format

        # URL-encode the startup commands parameter
        encoded_commands = urllib.parse.quote(startup_commands)
        params.append(f'startupCommands={encoded_commands}')
      except (json.JSONDecodeError, TypeError) as e:
        prt(f'Warning: Invalid startup commands JSON, ignoring. Error: {e}',
            ANSI.RED)

    if params:
      address += '&' + '&'.join(params)

    if open_browser:
      webbrowser.open_new_tab(address)
    else:
      print(f'Open URL in browser: {address}')

    httpd.expected_fname = fname
    httpd.fname_get_completed = None
    httpd.allow_origin = origin
    while httpd.fname_get_completed is None:
      httpd.handle_request()


def adb(*args, stdin=devnull, stdout=None, stderr=None):
  cmd = [adb_path, *args]
  setpgrp = None
  if os.name != 'nt':
    # On Linux/Mac, start a new process group so all child processes are killed
    # on exit. Unsupported on Windows.
    setpgrp = lambda: os.setpgrp()
  proc = subprocess.Popen(
      cmd, stdin=stdin, stdout=stdout, stderr=stderr, preexec_fn=setpgrp)
  procs.append(proc)
  return proc


def kill_all_subprocs_on_exit():
  for p in [p for p in procs if p.poll() is None]:
    p.kill()


def check_hash(file_name, sha_value):
  with open(file_name, 'rb') as fd:
    file_hash = hashlib.sha1(fd.read()).hexdigest()
    return file_hash == sha_value


if __name__ == '__main__':
  sys.exit(main())