adb_profile_chrome: Refactor into multiple modules and add tests

build/android/adb_profile_chrome.py

Index: build/android/adb_profile_chrome.py
diff --git a/build/android/adb_profile_chrome.py b/build/android/adb_profile_chrome.py
index 2e25d9c40779229b1e7e0db57168981e0c4dd2b7..0270a62f153a33621bc2ac408c6743cf715f702a 100755
--- a/build/android/adb_profile_chrome.py
+++ b/build/android/adb_profile_chrome.py
@@ -4,472 +4,10 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import gzip
-import json
-import logging
-import optparse
-import os
-import re
-import select
-import shutil
 import sys
-import threading
-import time
-import webbrowser
-import zipfile
-import zlib
 
-from pylib import android_commands
-from pylib import cmd_helper
-from pylib import constants
-from pylib import pexpect
-from pylib.device import device_utils
-
-_TRACE_VIEWER_ROOT = os.path.join(constants.DIR_SOURCE_ROOT,
-                                  'third_party', 'trace-viewer')
-sys.path.append(_TRACE_VIEWER_ROOT)
-from trace_viewer.build import trace2html # pylint: disable=F0401
-
-_DEFAULT_CHROME_CATEGORIES = '_DEFAULT_CHROME_CATEGORIES'
-
-
-def _GetTraceTimestamp():
-  return time.strftime('%Y-%m-%d-%H%M%S', time.localtime())
-
-
-class ChromeTracingController(object):
-  def __init__(self, device, package_info, categories, ring_buffer):
-    self._device = device
-    self._package_info = package_info
-    self._categories = categories
-    self._ring_buffer = ring_buffer
-    self._trace_file = None
-    self._trace_interval = None
-    self._trace_start_re = \
-       re.compile(r'Logging performance trace to file')
-    self._trace_finish_re = \
-       re.compile(r'Profiler finished[.] Results are in (.*)[.]')
-    self._device.old_interface.StartMonitoringLogcat(clear=False)
-
-  def __str__(self):
-    return 'chrome trace'
-
-  @staticmethod
-  def GetCategories(device, package_info):
-    device.old_interface.BroadcastIntent(
-        package_info.package, 'GPU_PROFILER_LIST_CATEGORIES')
-    try:
-      json_category_list = device.old_interface.WaitForLogMatch(
-          re.compile(r'{"traceCategoriesList(.*)'), None, timeout=5).group(0)
-    except pexpect.TIMEOUT:
-      raise RuntimeError('Performance trace category list marker not found. '
-                         'Is the correct version of the browser running?')
-
-    record_categories = []
-    disabled_by_default_categories = []
-    json_data = json.loads(json_category_list)['traceCategoriesList']
-    for item in json_data:
-      if item.startswith('disabled-by-default'):
-        disabled_by_default_categories.append(item)
-      else:
-        record_categories.append(item)
-
-    return record_categories, disabled_by_default_categories
-
-  def StartTracing(self, interval):
-    self._trace_interval = interval
-    self._device.old_interface.SyncLogCat()
-    self._device.old_interface.BroadcastIntent(
-        self._package_info.package, 'GPU_PROFILER_START',
-        '-e categories "%s"' % ','.join(self._categories),
-        '-e continuous' if self._ring_buffer else '')
-    # Chrome logs two different messages related to tracing:
-    #
-    # 1. "Logging performance trace to file"
-    # 2. "Profiler finished. Results are in [...]"
-    #
-    # The first one is printed when tracing starts and the second one indicates
-    # that the trace file is ready to be pulled.
-    try:
-      self._device.old_interface.WaitForLogMatch(
-          self._trace_start_re, None, timeout=5)
-    except pexpect.TIMEOUT:
-      raise RuntimeError('Trace start marker not found. Is the correct version '
-                         'of the browser running?')
-
-  def StopTracing(self):
-    self._device.old_interface.BroadcastIntent(
-        self._package_info.package,
-        'GPU_PROFILER_STOP')
-    self._trace_file = self._device.old_interface.WaitForLogMatch(
-        self._trace_finish_re, None, timeout=120).group(1)
-
-  def PullTrace(self):
-    # Wait a bit for the browser to finish writing the trace file.
-    time.sleep(self._trace_interval / 4 + 1)
-
-    trace_file = self._trace_file.replace('/storage/emulated/0/', '/sdcard/')
-    host_file = os.path.join(os.path.curdir, os.path.basename(trace_file))
-    self._device.old_interface.PullFileFromDevice(trace_file, host_file)
-    return host_file
-
-
-_SYSTRACE_OPTIONS = [
-    # Compress the trace before sending it over USB.
-    '-z',
-    # Use a large trace buffer to increase the polling interval.
-    '-b', '16384'
-]
-
-# Interval in seconds for sampling systrace data.
-_SYSTRACE_INTERVAL = 15
-
-
-class SystraceController(object):
-  def __init__(self, device, categories, ring_buffer):
-    self._device = device
-    self._categories = categories
-    self._ring_buffer = ring_buffer
-    self._done = threading.Event()
-    self._thread = None
-    self._trace_data = None
-
-  def __str__(self):
-    return 'systrace'
-
-  @staticmethod
-  def GetCategories(device):
-    return device.old_interface.RunShellCommand('atrace --list_categories')
-
-  def StartTracing(self, _):
-    self._thread = threading.Thread(target=self._CollectData)
-    self._thread.start()
-
-  def StopTracing(self):
-    self._done.set()
-
-  def PullTrace(self):
-    self._thread.join()
-    self._thread = None
-    if self._trace_data:
-      output_name = 'systrace-%s' % _GetTraceTimestamp()
-      with open(output_name, 'w') as out:
-        out.write(self._trace_data)
-      return output_name
-
-  def _RunATraceCommand(self, command):
-    # TODO(jbudorick) can this be made work with DeviceUtils?
-    # We use a separate interface to adb because the one from AndroidCommands
-    # isn't re-entrant.
-    device_param = (['-s', self._device.old_interface.GetDevice()]
-                    if self._device.old_interface.GetDevice() else [])
-    cmd = ['adb'] + device_param + ['shell', 'atrace', '--%s' % command] + \
-        _SYSTRACE_OPTIONS + self._categories
-    return cmd_helper.GetCmdOutput(cmd)
-
-  def _CollectData(self):
-    trace_data = []
-    self._RunATraceCommand('async_start')
-    try:
-      while not self._done.is_set():
-        self._done.wait(_SYSTRACE_INTERVAL)
-        if not self._ring_buffer or self._done.is_set():
-          trace_data.append(
-              self._DecodeTraceData(self._RunATraceCommand('async_dump')))
-    finally:
-      trace_data.append(
-          self._DecodeTraceData(self._RunATraceCommand('async_stop')))
-    self._trace_data = ''.join([zlib.decompress(d) for d in trace_data])
-
-  @staticmethod
-  def _DecodeTraceData(trace_data):
-    try:
-      trace_start = trace_data.index('TRACE:')
-    except ValueError:
-      raise RuntimeError('Systrace start marker not found')
-    trace_data = trace_data[trace_start + 6:]
-
-    # Collapse CRLFs that are added by adb shell.
-    if trace_data.startswith('\r\n'):
-      trace_data = trace_data.replace('\r\n', '\n')
-
-    # Skip the initial newline.
-    return trace_data[1:]
-
-
-def _GetSupportedBrowsers():
-  # Add aliases for backwards compatibility.
-  supported_browsers = {
-    'stable': constants.PACKAGE_INFO['chrome_stable'],
-    'beta': constants.PACKAGE_INFO['chrome_beta'],
-    'dev': constants.PACKAGE_INFO['chrome_dev'],
-    'build': constants.PACKAGE_INFO['chrome'],
-  }
-  supported_browsers.update(constants.PACKAGE_INFO)
-  unsupported_browsers = ['content_browsertests', 'gtest', 'legacy_browser']
-  for browser in unsupported_browsers:
-    del supported_browsers[browser]
-  return supported_browsers
-
-
-def _CompressFile(host_file, output):
-  with gzip.open(output, 'wb') as out:
-    with open(host_file, 'rb') as input_file:
-      out.write(input_file.read())
-  os.unlink(host_file)
-
-
-def _ArchiveFiles(host_files, output):
-  with zipfile.ZipFile(output, 'w', zipfile.ZIP_DEFLATED) as z:
-    for host_file in host_files:
-      z.write(host_file)
-      os.unlink(host_file)
-
-
-def _PackageTracesAsHtml(trace_files, html_file):
-  with open(html_file, 'w') as f:
-    trace2html.WriteHTMLForTracesToFile(trace_files, f)
-  for trace_file in trace_files:
-    os.unlink(trace_file)
-
-
-def _PrintMessage(heading, eol='\n'):
-  sys.stdout.write('%s%s' % (heading, eol))
-  sys.stdout.flush()
-
-
-def _WaitForEnter(timeout):
-  select.select([sys.stdin], [], [], timeout)
-
-
-def _StartTracing(controllers, interval):
-  for controller in controllers:
-    controller.StartTracing(interval)
-
-
-def _StopTracing(controllers):
-  for controller in controllers:
-    controller.StopTracing()
-
-
-def _PullTraces(controllers, output, compress, write_json):
-  _PrintMessage('Downloading...', eol='')
-  trace_files = []
-  for controller in controllers:
-    trace_files.append(controller.PullTrace())
-
-  if not write_json:
-    html_file = os.path.splitext(trace_files[0])[0] + '.html'
-    _PackageTracesAsHtml(trace_files, html_file)
-    trace_files = [html_file]
-
-  if compress and len(trace_files) == 1:
-    result = output or trace_files[0] + '.gz'
-    _CompressFile(trace_files[0], result)
-  elif len(trace_files) > 1:
-    result = output or 'chrome-combined-trace-%s.zip' % _GetTraceTimestamp()
-    _ArchiveFiles(trace_files, result)
-  elif output:
-    result = output
-    shutil.move(trace_files[0], result)
-  else:
-    result = trace_files[0]
-
-  _PrintMessage('done')
-  _PrintMessage('Trace written to file://%s' % os.path.abspath(result))
-  return result
-
-
-def _CaptureAndPullTrace(controllers, interval, output, compress, write_json):
-  trace_type = ' + '.join(map(str, controllers))
-  try:
-    _StartTracing(controllers, interval)
-    if interval:
-      _PrintMessage('Capturing %d-second %s. Press Enter to stop early...' % \
-          (interval, trace_type), eol='')
-      _WaitForEnter(interval)
-    else:
-      _PrintMessage('Capturing %s. Press Enter to stop...' % trace_type, eol='')
-      raw_input()
-  finally:
-    _StopTracing(controllers)
-  if interval:
-    _PrintMessage('done')
-
-  return _PullTraces(controllers, output, compress, write_json)
-
-
-def _ComputeChromeCategories(options):
-  categories = []
-  if options.trace_frame_viewer:
-    categories.append('disabled-by-default-cc.debug')
-  if options.trace_ubercompositor:
-    categories.append('disabled-by-default-cc.debug*')
-  if options.trace_gpu:
-    categories.append('disabled-by-default-gpu.debug*')
-  if options.trace_flow:
-    categories.append('disabled-by-default-toplevel.flow')
-  if options.chrome_categories:
-    categories += options.chrome_categories.split(',')
-  return categories
-
-
-def _ComputeSystraceCategories(options):
-  if not options.systrace_categories:
-    return []
-  return options.systrace_categories.split(',')
-
-
-def main():
-  parser = optparse.OptionParser(description='Record about://tracing profiles '
-                                 'from Android browsers. See http://dev.'
-                                 'chromium.org/developers/how-tos/trace-event-'
-                                 'profiling-tool for detailed instructions for '
-                                 'profiling.')
-
-  timed_options = optparse.OptionGroup(parser, 'Timed tracing')
-  timed_options.add_option('-t', '--time', help='Profile for N seconds and '
-                          'download the resulting trace.', metavar='N',
-                           type='float')
-  parser.add_option_group(timed_options)
-
-  cont_options = optparse.OptionGroup(parser, 'Continuous tracing')
-  cont_options.add_option('--continuous', help='Profile continuously until '
-                          'stopped.', action='store_true')
-  cont_options.add_option('--ring-buffer', help='Use the trace buffer as a '
-                          'ring buffer and save its contents when stopping '
-                          'instead of appending events into one long trace.',
-                          action='store_true')
-  parser.add_option_group(cont_options)
-
-  categories = optparse.OptionGroup(parser, 'Trace categories')
-  categories.add_option('-c', '--categories', help='Select Chrome tracing '
-                        'categories with comma-delimited wildcards, '
-                        'e.g., "*", "cat1*,-cat1a". Omit this option to trace '
-                        'Chrome\'s default categories. Chrome tracing can be '
-                        'disabled with "--categories=\'\'". Use "list" to see '
-                        'the available categories.',
-                        metavar='CHROME_CATEGORIES', dest='chrome_categories',
-                        default=_DEFAULT_CHROME_CATEGORIES)
-  categories.add_option('-s', '--systrace', help='Capture a systrace with the '
-                        'chosen comma-delimited systrace categories. You can '
-                        'also capture a combined Chrome + systrace by enabling '
-                        'both types of categories. Use "list" to see the '
-                        'available categories. Systrace is disabled by '
-                        'default.', metavar='SYS_CATEGORIES',
-                        dest='systrace_categories', default='')
-  categories.add_option('--trace-cc',
-                        help='Deprecated, use --trace-frame-viewer.',
-                        action='store_true')
-  categories.add_option('--trace-frame-viewer',
-                        help='Enable enough trace categories for '
-                        'compositor frame viewing.', action='store_true')
-  categories.add_option('--trace-ubercompositor',
-                        help='Enable enough trace categories for '
-                        'ubercompositor frame data.', action='store_true')
-  categories.add_option('--trace-gpu', help='Enable extra trace categories for '
-                        'GPU data.', action='store_true')
-  categories.add_option('--trace-flow', help='Enable extra trace categories '
-                        'for IPC message flows.', action='store_true')
-  parser.add_option_group(categories)
-
-  output_options = optparse.OptionGroup(parser, 'Output options')
-  output_options.add_option('-o', '--output', help='Save trace output to file.')
-  output_options.add_option('--json', help='Save trace as raw JSON instead of '
-                            'HTML.', action='store_true')
-  output_options.add_option('--view', help='Open resulting trace file in a '
-                            'browser.', action='store_true')
-  parser.add_option_group(output_options)
-
-  browsers = sorted(_GetSupportedBrowsers().keys())
-  parser.add_option('-b', '--browser', help='Select among installed browsers. '
-                    'One of ' + ', '.join(browsers) + ', "stable" is used by '
-                    'default.', type='choice', choices=browsers,
-                    default='stable')
-  parser.add_option('-v', '--verbose', help='Verbose logging.',
-                    action='store_true')
-  parser.add_option('-z', '--compress', help='Compress the resulting trace '
-                    'with gzip. ', action='store_true')
-  options, _args = parser.parse_args()
-  if options.trace_cc:
-    parser.parse_error("""--trace-cc is deprecated.
-
-For basic jank busting uses, use  --trace-frame-viewer
-For detailed study of ubercompositor, pass --trace-ubercompositor.
-
-When in doubt, just try out --trace-frame-viewer.
-""")
-
-  if options.verbose:
-    logging.getLogger().setLevel(logging.DEBUG)
-
-  devices = android_commands.GetAttachedDevices()
-  if len(devices) != 1:
-    parser.error('Exactly 1 device much be attached.')
-  device = device_utils.DeviceUtils(devices[0])
-  package_info = _GetSupportedBrowsers()[options.browser]
-
-  if options.chrome_categories in ['list', 'help']:
-    _PrintMessage('Collecting record categories list...', eol='')
-    record_categories = []
-    disabled_by_default_categories = []
-    record_categories, disabled_by_default_categories = \
-        ChromeTracingController.GetCategories(device, package_info)
-
-    _PrintMessage('done')
-    _PrintMessage('Record Categories:')
-    _PrintMessage('\n'.join('\t%s' % item \
-        for item in sorted(record_categories)))
-
-    _PrintMessage('\nDisabled by Default Categories:')
-    _PrintMessage('\n'.join('\t%s' % item \
-        for item in sorted(disabled_by_default_categories)))
-
-    return 0
-
-  if options.systrace_categories in ['list', 'help']:
-    _PrintMessage('\n'.join(SystraceController.GetCategories(device)))
-    return 0
-
-  if not options.time and not options.continuous:
-    _PrintMessage('Time interval or continuous tracing should be specified.')
-    return 1
-
-  chrome_categories = _ComputeChromeCategories(options)
-  systrace_categories = _ComputeSystraceCategories(options)
-
-  if chrome_categories and 'webview' in systrace_categories:
-    logging.warning('Using the "webview" category in systrace together with '
-                    'Chrome tracing results in duplicate trace events.')
-
-  controllers = []
-  if chrome_categories:
-    controllers.append(ChromeTracingController(device,
-                                               package_info,
-                                               chrome_categories,
-                                               options.ring_buffer))
-  if systrace_categories:
-    controllers.append(SystraceController(device,
-                                          systrace_categories,
-                                          options.ring_buffer))
-
-  if not controllers:
-    _PrintMessage('No trace categories enabled.')
-    return 1
-
-  if options.output:
-    options.output = os.path.expanduser(options.output)
-  result = _CaptureAndPullTrace(controllers,
-                                options.time if not options.continuous else 0,
-                                options.output,
-                                options.compress,
-                                options.json)
-  if options.view:
-    if sys.platform == 'darwin':
-      os.system('/usr/bin/open %s' % os.path.abspath(result))
-    else:
-      webbrowser.open(result)
+from chrome_profiler import main
 
 
 if __name__ == '__main__':
-  sys.exit(main())
+  sys.exit(main.main())