adb_profile_chrome: Refactor into multiple modules and add tests

build/android/chrome_profiler/controllers.py

Index: build/android/chrome_profiler/controllers.py
diff --git a/build/android/chrome_profiler/controllers.py b/build/android/chrome_profiler/controllers.py
new file mode 100644
index 0000000000000000000000000000000000000000..b00409eb7ae9856c498200b5da03b699d5ff15c9
--- /dev/null
+++ b/build/android/chrome_profiler/controllers.py
@@ -0,0 +1,175 @@
+# Copyright 2014 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import json
+import os
+import re
+import threading
+import time
+import zlib
+
+from chrome_profiler import util
+
+from pylib import cmd_helper
+from pylib import pexpect
+
+
+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 __repr__(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 __repr__(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' % util.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:]