android: Add systrace and combined trace support to adb_profile_chrome

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 b328e47ca979fc6ecd5fc817f4ed5c6394617caf..c8cea89ceb2406430293fed702220fae239ea433 100755
--- a/build/android/adb_profile_chrome.py
+++ b/build/android/adb_profile_chrome.py
@@ -10,13 +10,156 @@ import optparse
 import os
 import re
 import sys
+import threading
 import time
+import zipfile
+import zlib
 
 from pylib import android_commands
+from pylib import cmd_helper
 from pylib import constants
 from pylib import pexpect
 
 
+def _GetTraceTimestamp():
+ return time.strftime('%Y-%m-%d-%H%M%S', time.localtime())
+
+
+class ChromeTracingController(object):
+  def __init__(self, adb, package_info, categories, ring_buffer):
+    self._adb = adb
+    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._adb.StartMonitoringLogcat(clear=False)
+
+  def __str__(self):
+    return 'chrome trace'
+
+  def StartTracing(self, interval):
+    self._trace_interval = interval
+    self._adb.SyncLogCat()
+    self._adb.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._trace_file = self._adb.WaitForLogMatch(self._trace_start_re,
+                                                   None,
+                                                   timeout=5).group(1)
+    except pexpect.TIMEOUT:
+      raise RuntimeError('Trace start marker not found. Is the correct version '
+                         'of the browser running?')
+
+  def StopTracing(self):
+    if not self._trace_file:
+      return
+    self._adb.BroadcastIntent(self._package_info.package, 'GPU_PROFILER_STOP')
+    self._adb.WaitForLogMatch(self._trace_finish_re, None, timeout=120)
+
+  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._adb.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, adb, categories, ring_buffer):
+    self._adb = adb
+    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(adb):
+    return adb.RunShellCommand('atrace --list_categories')
+
+  def StartTracing(self, interval):
+    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):
+    # We use a separate interface to adb because the one from AndroidCommands
+    # isn't re-entrant.
+    device = ['-s', self._adb.GetDevice()] if self._adb.GetDevice() else []
+    cmd = ['adb'] + device + ['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 = {
@@ -32,111 +175,93 @@ def _GetSupportedBrowsers():
   return supported_browsers
 
 
-def _StartTracing(adb, package_info, categories, continuous):
-  adb.BroadcastIntent(package_info.package, 'GPU_PROFILER_START',
-                      '-e categories "%s"' % categories,
-                      '-e continuous' if continuous else '')
-
-
-def _StopTracing(adb, package_info):
-  adb.BroadcastIntent(package_info.package, 'GPU_PROFILER_STOP')
+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 _GetLatestTraceFileName(adb, check_for_multiple_traces=True):
-  # 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 downloaded.
-  #
-  # We have to look for both of these messages to make sure we get the results
-  # from the latest trace and that the trace file is complete.
-  trace_start_re = re.compile(r'Logging performance trace to file')
-  trace_finish_re = re.compile(r'Profiler finished[.] Results are in (.*)[.]')
+def _ArchiveFiles(host_files, output):
+  with zipfile.ZipFile(output, 'w') as z:
+    for host_file in host_files:
+      z.write(host_file)
+      os.unlink(host_file)
 
-  start_timeout = 5
-  start_match = None
-  finish_match = None
 
-  while True:
-    try:
-      start_match = adb.WaitForLogMatch(trace_start_re, None,
-                                        timeout=start_timeout)
-    except pexpect.TIMEOUT:
-      if start_match:
-        break
-      raise RuntimeError('Trace start marker not found. Is the correct version '
-                         'of the browser running?')
-    finish_match = adb.WaitForLogMatch(trace_finish_re, None, timeout=120)
-    if not check_for_multiple_traces:
-      break
-    # Now that we've found one trace file, switch to polling for the rest of the
-    # log to see if there are more.
-    start_timeout = 1
-  return finish_match.group(1)
+def _PrintMessage(heading, eol='\n'):
+  sys.stdout.write('%s%s' % (heading, eol))
+  sys.stdout.flush()
 
 
-def _DownloadTrace(adb, trace_file):
-  trace_file = trace_file.replace('/storage/emulated/0/', '/sdcard/')
-  host_file = os.path.abspath(os.path.basename(trace_file))
-  adb.PullFileFromDevice(trace_file, host_file)
-  return host_file
+def _StartTracing(controllers, interval):
+  for controller in controllers:
+    controller.StartTracing(interval)
 
 
-def _CompressFile(host_file):
-  compressed_file = host_file + '.gz'
-  with gzip.open(compressed_file, 'wb') as out:
-    with open(host_file, 'rb') as input_file:
-      out.write(input_file.read())
-  os.unlink(host_file)
-  return compressed_file
+def _StopTracing(controllers):
+  for controller in controllers:
+    controller.StopTracing()
 
 
-def _PrintMessage(heading, eol='\n'):
-  sys.stdout.write(heading + eol)
-  sys.stdout.flush()
+def _PullTraces(controllers, output, compress):
+  _PrintMessage('Downloading...', eol='')
+  trace_files = []
+  for controller in controllers:
+    trace_files.append(controller.PullTrace())
 
+  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
+    os.rename(trace_files[0], result)
+  else:
+    result = trace_files[0]
 
-def _DownloadLatestTrace(adb, compress, check_for_multiple_traces=True):
-  _PrintMessage('Downloading trace...', eol='')
-  trace_file = _GetLatestTraceFileName(adb, check_for_multiple_traces)
-  host_file = _DownloadTrace(adb, trace_file)
-  if compress:
-    host_file = _CompressFile(host_file)
   _PrintMessage('done')
-  _PrintMessage('Trace written to %s' % host_file)
-
+  _PrintMessage('Trace written to %s' % os.path.abspath(result))
 
-def _CaptureAndDownloadTimedTrace(adb, package_info, categories, interval,
-                                 continuous, compress):
-  adb.StartMonitoringLogcat(clear=False)
-  adb.SyncLogCat()
 
+def _CaptureAndPullTrace(controllers, interval, output, compress):
+  trace_type = ' + '.join(map(str, controllers))
   try:
-    _PrintMessage('Capturing %d-second trace. Press Ctrl-C to stop early...' \
-        % interval, eol='')
-    _StartTracing(adb, package_info, categories, continuous)
-    time.sleep(interval)
+    _StartTracing(controllers, interval)
+    if interval:
+      _PrintMessage('Capturing %d-second %s. Press Ctrl-C to stop early...' % \
+          (interval, trace_type), eol='')
+      time.sleep(interval)
+    else:
+      _PrintMessage('Capturing %s. Press Enter to stop...' % trace_type, eol='')
+      raw_input()
   except KeyboardInterrupt:
-    _PrintMessage('\nInterrupted, stopping...', eol='')
-  _StopTracing(adb, package_info)
-  _PrintMessage('done')
+    _PrintMessage('\nInterrupted...', eol='')
+  finally:
+    _StopTracing(controllers)
+  if interval:
+    _PrintMessage('done')
 
-  # Wait a bit for the browser to finish writing the trace file.
-  time.sleep(interval / 4 + 1)
+  _PullTraces(controllers, output, compress)
 
-  _DownloadLatestTrace(adb, compress, check_for_multiple_traces=False)
 
-
-def _ComputeCategories(options):
-  categories = [options.categories]
+def _ComputeChromeCategories(options):
+  categories = []
   if options.trace_cc:
     categories.append('disabled-by-default-cc.debug*')
   if options.trace_gpu:
     categories.append('disabled-by-default-gpu.debug*')
-  return ",".join(categories)
+  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():
@@ -145,36 +270,40 @@ def main():
                                  'chromium.org/developers/how-tos/trace-event-'
                                  'profiling-tool for detailed instructions for '
                                  'profiling.')
-  manual_options = optparse.OptionGroup(parser, 'Manual tracing')
-  manual_options.add_option('--start', help='Start tracing.',
-                            action='store_true')
-  manual_options.add_option('--stop', help='Stop tracing.',
-                            action='store_true')
-  manual_options.add_option('-d', '--download', help='Download latest trace.',
-                            action='store_true')
-  parser.add_option_group(manual_options)
-
-  auto_options = optparse.OptionGroup(parser, 'Automated tracing')
-  auto_options.add_option('-t', '--time', help='Profile for N seconds and '
+
+  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(auto_options)
+                           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 categories to trace '
                         'with comma-delimited wildcards, e.g., '
-                        '"*", "cat1*,-cat1a". Default is "*".', default='*')
+                        '"*", "cat1*,-cat1a". Default is "*".', default='*',
+                        dest='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.', metavar='SYS_CATEGORIES',
+                        dest='systrace_categories', default='')
   categories.add_option('--trace-cc', help='Enable extra trace categories for '
                         'compositor frame viewer data.', action='store_true')
   categories.add_option('--trace-gpu', help='Enable extra trace categories for '
                         'GPU data.', action='store_true')
   parser.add_option_group(categories)
 
-  parser.add_option('-o', '--output', help='Save profile output to file. '
-                    'Default is "/sdcard/Download/chrome-profile-results-*".')
-  parser.add_option('--continuous', help='Using the trace buffer as a ring '
-                    'buffer, continuously profile until stopped.',
-                    action='store_true')
+  parser.add_option('-o', '--output', help='Save profile output to file.')
   browsers = sorted(_GetSupportedBrowsers().keys())
   parser.add_option('-b', '--browser', help='Select among installed browsers. '
                     'One of ' + ', '.join(browsers) + ', "stable" is used by '
@@ -186,26 +315,45 @@ def main():
                     'with gzip. ', action='store_true')
   options, args = parser.parse_args()
 
-  if not any([options.start, options.stop, options.time, options.download]):
-    _PrintMessage('One of start/stop/download/time should be specified.')
-    return 1
-
   if options.verbose:
     logging.getLogger().setLevel(logging.DEBUG)
 
   adb = android_commands.AndroidCommands()
-  categories = _ComputeCategories(options)
+  if options.systrace_categories in ['list', 'help']:
+    _PrintMessage('\n'.join(SystraceController.GetCategories(adb)))
+    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)
   package_info = _GetSupportedBrowsers()[options.browser]
 
-  if options.start:
-    _StartTracing(adb, package_info, categories, options.continuous)
-  elif options.stop:
-    _StopTracing(adb, package_info)
-  elif options.download:
-    _DownloadLatestTrace(adb, options.compress)
-  elif options.time:
-    _CaptureAndDownloadTimedTrace(adb, package_info, categories, options.time,
-                                  options.continuous, options.compress)
+  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(adb,
+                                               package_info,
+                                               chrome_categories,
+                                               options.ring_buffer))
+  if systrace_categories:
+    controllers.append(SystraceController(adb,
+                                          systrace_categories,
+                                          options.ring_buffer))
+
+  if not controllers:
+    _PrintMessage('No trace categories enabled.')
+    return 1
+
+  _CaptureAndPullTrace(controllers,
+                       options.time if not options.continuous else 0,
+                       options.output,
+                       options.compress)
 
 
 if __name__ == '__main__':