[DO NOT COMMIT] Refactor systrace to support new clock sync design

systrace/systrace/tracing_agents/atrace_agent.py

Index: systrace/systrace/tracing_agents/atrace_agent.py
diff --git a/systrace/systrace/agents/atrace_agent.py b/systrace/systrace/tracing_agents/atrace_agent.py
similarity index 34%
rename from systrace/systrace/agents/atrace_agent.py
rename to systrace/systrace/tracing_agents/atrace_agent.py
index 3125821b22a743c0a58ad2f35620a4aead772037..c507b36536c069f14bc460c322f05acc47232758 100644
--- a/systrace/systrace/agents/atrace_agent.py
+++ b/systrace/systrace/tracing_agents/atrace_agent.py
@@ -2,16 +2,25 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import Queue
 import re
 import subprocess
 import sys
 import threading
-import time
 import zlib
 
-from systrace import systrace_agent
+from collections import namedtuple
+from devil.android import device_utils
+from devil.utils import timeout_retry
 from systrace import util
+from systrace.tracing_agents import TraceResults
+from systrace.tracing_agents import TracingAgent
+
+atrace_postproc = namedtuple('atrace_postproc', ['fix_threads',
+                                                 'fix_tgids',
+                                                 'fix_circular'])
+
+#TODO(alexandermont): add "from py_trace_event import trace_time"
+#after persistent shell CL lands
 
 # Text that ADB sends, but does not need to be displayed to the user.
 ADB_IGNORE_REGEXP = r'^capturing trace\.\.\. done|^capturing trace\.\.\.'
@@ -50,230 +59,221 @@ LEGACY_TRACE_TAG_BITS = (
     ('camera', 1 << 10),
 )
 
+def list_categories(options):
+  """List the possible trace event categories.
+
+  This function needs the tracing options since it needs to get the serial
+  number of the device to send a command to.
+
+  Args:
+      options: Tracing options.
+  """
+  devutils = device_utils.DeviceUtils(options.device_serial)
+  for line in devutils.RunShellCommand(LIST_CATEGORIES_ARGS):
+    print line
+
 
-def try_create_agent(options, categories):
+def get_available_categories(options):
+  """Gets the list of atrace categories available for tracing.
+
+  Args:
+      options: Tracing options.
+  """
+  devutils = device_utils.DeviceUtils(options.device_serial)
+  categories_output = devutils.RunShellCommand(LIST_CATEGORIES_ARGS)
+  return [c.split('-')[0].strip() for c in categories_output]
+
+
+def try_create_agent(options, _):
+  """Create an Atrace agent.
+
+  Args:
+      options: Command line options.
+      categories: Categories of trace events to capture.
+  """
   if options.target != 'android':
     return False
   if options.from_file is not None:
-    return AtraceAgent(options, categories)
+    return False
 
+  # Check for supported versions
   device_sdk_version = util.get_device_sdk_version()
-  if device_sdk_version >= 18:
-    if options.boot:
-      # atrace --async_stop, which is used by BootAgent, does not work properly
-      # on the device SDK version 22 or before.
-      if device_sdk_version <= 22:
-        print >> sys.stderr, ('--boot option does not work on the device SDK '
-                              'version 22 or before.\nYour device SDK version '
-                              'is %d.' % device_sdk_version)
-        sys.exit(1)
-      return BootAgent(options, categories)
-    else:
-      return AtraceAgent(options, categories)
-  elif device_sdk_version >= 16:
-    return AtraceLegacyAgent(options, categories)
+  if device_sdk_version <= 17:
+    print >> sys.stderr, ('Device SDK versions <= 17 not supported.\n'
+                          'Your device SDK version is %d.' % device_sdk_version)
+  elif device_sdk_version <= 22 and options.boot:
+    print >> sys.stderr, (
+        'Device SDK versions <= 22 do not support boot tracing.\n'
+        'Your device SDK version is %d.' % device_sdk_version)
 
+  if options.boot:
+    return BootAgent()
+  else:
+    return AtraceAgent()
 
-class AtraceAgent(systrace_agent.SystraceAgent):
+def _construct_extra_trace_args(options, categories):
+  """Construct extra arguments (-a and -k) for trace command.
 
-  def __init__(self, options, categories):
-    super(AtraceAgent, self).__init__(options, categories)
-    self._expect_trace = False
-    self._adb = None
+  Args:
+      options: Tracing options.
+      categories: Categories of trace events to capture.
+  """
+  extra_args = []
+  if options.app_name is not None:
+    extra_args.extend(['-a', options.app_name])
+
+  if options.kfuncs is not None:
+    extra_args.extend(['-k', options.kfuncs])
+
+  avail_cats = get_available_categories(options)
+  extra_args.extend([x for x in categories if x in avail_cats])
+  return extra_args
+
+
+def _construct_trace_args(options, categories):
+  """Builds a command-line used to invoke a trace process.
+
+  Returns:
+    A tuple where the first element is an array of command-line arguments, and
+    the second element is a boolean which will be true if the command will
+    stream trace data.
+  """
+  atrace_args = ATRACE_BASE_ARGS[:]
+
+  if options.compress_trace_data:
+    atrace_args.extend(['-z'])
+
+  if ((options.trace_time is not None)
+      and (options.trace_time > 0)):
+    atrace_args.extend(['-t', str(options.trace_time)])
+
+  if ((options.trace_buf_size is not None)
+      and (options.trace_buf_size > 0)):
+    atrace_args.extend(['-b', str(options.trace_buf_size)])
+  elif 'sched' in categories:
+    # 'sched' is a high-volume tag, double the default buffer size
+    # to accommodate that
+    atrace_args.extend(['-b', '4096'])
+  extra_args = _construct_extra_trace_args(options, categories)
+  atrace_args.extend(extra_args)
+
+  return atrace_args
+
+class AtraceAgent(TracingAgent):
+
+  def __init__(self):
+    super(AtraceAgent, self).__init__()
     self._trace_data = None
     self._tracer_args = None
-    if not self._categories:
-      self._categories = get_default_categories(self._options.device_serial)
+    self._adb = None
+    self._collection_process = None
+    self._device_utils = None
+    self._device_serial = None
+    self._postproc = None
 
-  def start(self):
-    self._tracer_args = self._construct_trace_command()
+  def __str__(self):
+    return 'Atrace Agent'
 
-    self._adb = do_popen(self._tracer_args)
+  def _StartAgentTracingImpl(self, options, categories):
+    """Starts tracing.
 
-  def collect_result(self):
+    Args:
+        options: Tracing options.
+        categories: Categories of trace events to capture.
+    """
+    if not categories:
+      categories = DEFAULT_CATEGORIES
+    self._device_utils = device_utils.DeviceUtils(options.device_serial)
+    self._device_serial = options.device_serial
+    self._tracer_args = _construct_trace_args(options, categories)
+    self._device_utils.RunShellCommand(self._tracer_args + ['--async_start'])
+    self._postproc = atrace_postproc(options.fix_threads,
+                                     options.fix_tgids,
+                                     options.fix_circular)
+    return True
+
+  def StartAgentTracing(self, options, categories, timeout=10):
+    return timeout_retry.Run(self._StartAgentTracingImpl,
+                             timeout, 1, args=[options, categories])
+
+
+  def _collect_and_preprocess(self):
+    """Collects and preprocesses trace data.
+
+    Stores results in self._trace_data.
+    """
     trace_data = self._collect_trace_data()
-    if self._expect_trace:
-      self._trace_data = self._preprocess_trace_data(trace_data)
+    self._trace_data = self._preprocess_trace_data(trace_data)
 
-  def expect_trace(self):
-    return self._expect_trace
+  def _StopAgentTracingImpl(self):
+    """Stops tracing and collect results asynchronously.
 
-  def get_trace_data(self):
-    return self._trace_data
+    Creates a new process that stops tracing, collects results, and
+    records them in self._trace_data. Returns immediately, without
+    waiting for data collection to finish.
+    """
+    self._collection_process = threading.Thread(
+        target=self._collect_and_preprocess)
+    self._collection_process.start()
+    return True
 
-  def get_class_name(self):
-    return 'trace-data'
+  def StopAgentTracing(self, timeout=10):
+    return timeout_retry.Run(self._StopAgentTracingImpl,
+                             timeout, 1)
 
-  def _construct_list_categories_command(self):
-    return util.construct_adb_shell_command(
-          LIST_CATEGORIES_ARGS, self._options.device_serial)
+  def _GetResultsImpl(self):
+    """Returns trace results."""
+    self._collection_process.join()
+    self._collection_process = None
+    return TraceResults('systemTraceEvents', self._trace_data)
 
-  def _construct_extra_trace_command(self):
-    extra_args = []
-    if self._options.app_name is not None:
-      extra_args.extend(['-a', self._options.app_name])
+  def GetResults(self, timeout=30):
+    return timeout_retry.Run(self._GetResultsImpl,
+                             timeout, 1)
 
-    if self._options.kfuncs is not None:
-      extra_args.extend(['-k', self._options.kfuncs])
 
-    extra_args.extend(self._categories)
-    return extra_args
+  def SupportsExplicitClockSync(self):
+    """Returns whether or not this supports explicit clock sync."""
+    return False
+    #TODO(alexandermont): return True after persistent shell CL lands
 
-  def _construct_trace_command(self):
-    """Builds a command-line used to invoke a trace process.
+  def RecordClockSyncMarker(self, sync_id, did_record_sync_marker_callback):
+    """Records a clock sync marker.
 
-    Returns:
-      A tuple where the first element is an array of command-line arguments, and
-      the second element is a boolean which will be true if the commend will
-      stream trace data.
+    Args:
+        sync_id: ID string for clock sync marker.
     """
-    if self._options.list_categories:
-      tracer_args = self._construct_list_categories_command()
-      self._expect_trace = False
-    elif self._options.from_file is not None:
-      tracer_args = ['cat', self._options.from_file]
-      self._expect_trace = True
-    else:
-      atrace_args = ATRACE_BASE_ARGS[:]
-      self._expect_trace = True
-      if self._options.compress_trace_data:
-        atrace_args.extend(['-z'])
-
-      if ((self._options.trace_time is not None)
-          and (self._options.trace_time > 0)):
-        atrace_args.extend(['-t', str(self._options.trace_time)])
-
-      if ((self._options.trace_buf_size is not None)
-          and (self._options.trace_buf_size > 0)):
-        atrace_args.extend(['-b', str(self._options.trace_buf_size)])
-      elif 'sched' in self._categories:
-        # 'sched' is a high-volume tag, double the default buffer size
-        # to accommodate that
-        atrace_args.extend(['-b', '4096'])
-      extra_args = self._construct_extra_trace_command()
-      atrace_args.extend(extra_args)
-
-      tracer_args = util.construct_adb_shell_command(
-          atrace_args, self._options.device_serial)
-
-    return tracer_args
+    return
+    #TODO(alexandermont): put the functoin back after persistent shell CL lands
 
-  def _collect_trace_data(self):
-    # Read the output from ADB in a worker thread.  This allows us to monitor
-    # the progress of ADB and bail if ADB becomes unresponsive for any reason.
-
-    # Limit the stdout_queue to 128 entries because we will initially be reading
-    # one byte at a time.  When the queue fills up, the reader thread will
-    # block until there is room in the queue.  Once we start downloading the
-    # trace data, we will switch to reading data in larger chunks, and 128
-    # entries should be plenty for that purpose.
-    stdout_queue = Queue.Queue(maxsize=128)
-    stderr_queue = Queue.Queue()
-
-    if self._expect_trace:
-      # Use stdout.write() (here and for the rest of this function) instead
-      # of print() to avoid extra newlines.
-      sys.stdout.write('Capturing trace...')
-
-    # Use a chunk_size of 1 for stdout so we can display the output to
-    # the user without waiting for a full line to be sent.
-    stdout_thread = FileReaderThread(self._adb.stdout, stdout_queue,
-                                     text_file=False, chunk_size=1)
-    stderr_thread = FileReaderThread(self._adb.stderr, stderr_queue,
-                                     text_file=True)
-    stdout_thread.start()
-    stderr_thread.start()
-
-    # Holds the trace data returned by ADB.
-    trace_data = []
-    # Keep track of the current line so we can find the TRACE_START_REGEXP.
-    current_line = ''
-    # Set to True once we've received the TRACE_START_REGEXP.
-    reading_trace_data = False
-
-    last_status_update_time = time.time()
-
-    while (stdout_thread.isAlive() or stderr_thread.isAlive() or
-           not stdout_queue.empty() or not stderr_queue.empty()):
-      if self._expect_trace:
-        last_status_update_time = status_update(last_status_update_time)
-
-      while not stderr_queue.empty():
-        # Pass along errors from adb.
-        line = stderr_queue.get()
-        sys.stderr.write(line)
-
-      # Read stdout from adb.  The loop exits if we don't get any data for
-      # ADB_STDOUT_READ_TIMEOUT seconds.
-      while True:
-        try:
-          chunk = stdout_queue.get(True, ADB_STDOUT_READ_TIMEOUT)
-        except Queue.Empty:
-          # Didn't get any data, so exit the loop to check that ADB is still
-          # alive and print anything sent to stderr.
-          break
-
-        if reading_trace_data:
-          # Save, but don't print, the trace data.
-          trace_data.append(chunk)
-        else:
-          if not self._expect_trace:
-            sys.stdout.write(chunk)
-          else:
-            # Buffer the output from ADB so we can remove some strings that
-            # don't need to be shown to the user.
-            current_line += chunk
-            if re.match(TRACE_START_REGEXP, current_line):
-              # We are done capturing the trace.
-              sys.stdout.write('Done.\n')
-              # Now we start downloading the trace data.
-              sys.stdout.write('Downloading trace...')
-
-              current_line = ''
-              # Use a larger chunk size for efficiency since we no longer
-              # need to worry about parsing the stream.
-              stdout_thread.set_chunk_size(4096)
-              reading_trace_data = True
-            elif chunk == '\n' or chunk == '\r':
-              # Remove ADB output that we don't care about.
-              current_line = re.sub(ADB_IGNORE_REGEXP, '', current_line)
-              if len(current_line) > 1:
-                # ADB printed something that we didn't understand, so show it
-                # it to the user (might be helpful for debugging).
-                sys.stdout.write(current_line)
-              # Reset our current line.
-              current_line = ''
-
-    if self._expect_trace:
-      if reading_trace_data:
-        # Indicate to the user that the data download is complete.
-        sys.stdout.write('Done.\n')
-      else:
-        # We didn't receive the trace start tag, so something went wrong.
-        sys.stdout.write('ERROR.\n')
-        # Show any buffered ADB output to the user.
-        current_line = re.sub(ADB_IGNORE_REGEXP, '', current_line)
-        if current_line:
-          sys.stdout.write(current_line)
-          sys.stdout.write('\n')
-
-    # The threads should already have stopped, so this is just for cleanup.
-    stdout_thread.join()
-    stderr_thread.join()
-
-    self._adb.stdout.close()
-    self._adb.stderr.close()
-
-    # The adb process should be done since it's io pipes are closed.  Call
-    # poll() to set the returncode.
-    self._adb.poll()
-
-    if self._adb.returncode != 0:
-      print >> sys.stderr, ('The command "%s" returned error code %d.' %
-                            (' '.join(self._tracer_args), self._adb.returncode))
-      sys.exit(1)
+  def _dump_trace(self):
+    """Dumps the atrace buffer and returns the dumped buffer."""
+    dump_cmd = self._tracer_args + ['--async_dump']
+    return self._device_utils.RunShellCommand(dump_cmd, split_lines=False)
 
-    return trace_data
+  def _stop_trace(self):
+    """Stops the atrace trace.
+
+    Tries to stop the atrace trace asynchronously and uses the fallback
+    method of running a zero-length synchronous trace if that fails.
+    """
+    self._device_utils.RunShellCommand(self._tracer_args + ['--async_stop'])
+    cmd = ['cat', '/sys/kernel/debug/tracing/tracing_on']
+    trace_on = int(self._device_utils.RunShellCommand(cmd)[0])
+    if trace_on:
+      self._device_utils.RunShellCommand(self._tracer_args + ['-t 0'])
+
+  def _collect_trace_data(self):
+    """Reads the output from atrace and stops the trace."""
+    result = self._dump_trace()
+    data_start = re.search(TRACE_START_REGEXP, result)
+    if data_start:
+      data_start = data_start.end(0)
+    else:
+      raise IOError('Unable to get atrace data. Did you forget adb root?')
+    output = re.sub(ADB_IGNORE_REGEXP, '', result[data_start:])
+    self._stop_trace()
+    return output
 
   def _preprocess_trace_data(self, trace_data):
     """Performs various processing on atrace data.
@@ -283,7 +283,6 @@ class AtraceAgent(systrace_agent.SystraceAgent):
     Returns:
       The processed trace data.
     """
-    trace_data = ''.join(trace_data)
     if trace_data:
       trace_data = strip_and_decompress_trace(trace_data)
 
@@ -292,225 +291,30 @@ class AtraceAgent(systrace_agent.SystraceAgent):
                             'written.')
       sys.exit(1)
 
-    if self._options.fix_threads:
+    if self._postproc.fix_threads:
       # Issue ps command to device and patch thread names
-      ps_dump = do_preprocess_adb_cmd('ps -t', self._options.device_serial)
+      ps_dump = do_preprocess_adb_cmd('ps -t', self._device_serial)
       if ps_dump is not None:
         thread_names = extract_thread_list(ps_dump)
         trace_data = fix_thread_names(trace_data, thread_names)
 
-    if self._options.fix_tgids:
+    if self._postproc.fix_tgids:
       # Issue printf command to device and patch tgids
-      procfs_dump = do_preprocess_adb_cmd('printf "%s\n" ' +
-                                          '/proc/[0-9]*/task/[0-9]*',
-                                          self._options.device_serial)
+      procfs_dump = do_preprocess_adb_cmd(
+          'printf "%s\n" ' + '/proc/[0-9]*/task/[0-9]*', self._device_serial)
       if procfs_dump is not None:
         pid2_tgid = extract_tgids(procfs_dump)
         trace_data = fix_missing_tgids(trace_data, pid2_tgid)
 
-    if self._options.fix_circular:
+    if self._postproc.fix_circular:
       trace_data = fix_circular_traces(trace_data)
 
     return trace_data
 
 
-class AtraceLegacyAgent(AtraceAgent):
-
-  def _construct_list_categories_command(self):
-    LEGACY_CATEGORIES = """       sched - CPU Scheduling
-        freq - CPU Frequency
-        idle - CPU Idle
-        load - CPU Load
-        disk - Disk I/O (requires root)
-         bus - Bus utilization (requires root)
-   workqueue - Kernel workqueues (requires root)"""
-    return ["echo", LEGACY_CATEGORIES]
-
-  def start(self):
-    super(AtraceLegacyAgent, self).start()
-    if self.expect_trace():
-      SHELL_ARGS = ['getprop', 'debug.atrace.tags.enableflags']
-      output, return_code = util.run_adb_shell(SHELL_ARGS,
-                                               self._options.device_serial)
-      if return_code != 0:
-        print >> sys.stderr, (
-            '\nThe command "%s" failed with the following message:'
-            % ' '.join(SHELL_ARGS))
-        print >> sys.stderr, str(output)
-        sys.exit(1)
-
-      flags = 0
-      try:
-        if output.startswith('0x'):
-          flags = int(output, 16)
-        elif output.startswith('0'):
-          flags = int(output, 8)
-        else:
-          flags = int(output)
-      except ValueError:
-        pass
-
-      if flags:
-        tags = []
-        for desc, bit in LEGACY_TRACE_TAG_BITS:
-          if bit & flags:
-            tags.append(desc)
-        categories = tags + self._categories
-        print 'Collecting data with following categories:', ' '.join(categories)
-
-  def _construct_extra_trace_command(self):
-    extra_args = []
-    if not self._categories:
-      self._categories = ['sched', ]
-    if 'sched' in self._categories:
-      extra_args.append('-s')
-    if 'freq' in self._categories:
-      extra_args.append('-f')
-    if 'idle' in self._categories:
-      extra_args.append('-i')
-    if 'load' in self._categories:
-      extra_args.append('-l')
-    if 'disk' in self._categories:
-      extra_args.append('-d')
-    if 'bus' in self._categories:
-      extra_args.append('-u')
-    if 'workqueue' in self._categories:
-      extra_args.append('-w')
-
-    return extra_args
-
-
-class BootAgent(AtraceAgent):
-  """AtraceAgent that specializes in tracing the boot sequence."""
-
-  def __init__(self, options, categories):
-    super(BootAgent, self).__init__(options, categories)
-
-  def start(self):
-    try:
-      setup_args = self._construct_setup_command()
-      try:
-        subprocess.check_call(setup_args)
-        print 'Hit Ctrl+C once the device has booted up.'
-        while True:
-          time.sleep(1)
-      except KeyboardInterrupt:
-        pass
-      tracer_args = self._construct_trace_command()
-      self._adb = subprocess.Popen(tracer_args, stdout=subprocess.PIPE,
-                                   stderr=subprocess.PIPE)
-    except OSError as error:
-      print >> sys.stderr, (
-          'The command "%s" failed with the following error:' %
-          ' '.join(tracer_args))
-      print >> sys.stderr, '    ', error
-      sys.exit(1)
-
-  def _construct_setup_command(self):
-    echo_args = ['echo'] + self._categories + ['>', BOOTTRACE_CATEGORIES]
-    setprop_args = ['setprop', BOOTTRACE_PROP, '1']
-    reboot_args = ['reboot']
-    return util.construct_adb_shell_command(
-        echo_args + ['&&'] + setprop_args + ['&&'] + reboot_args,
-        self._options.device_serial)
-
-  def _construct_trace_command(self):
-    self._expect_trace = True
-    atrace_args = ['atrace', '--async_stop']
-    setprop_args = ['setprop', BOOTTRACE_PROP, '0']
-    rm_args = ['rm', BOOTTRACE_CATEGORIES]
-    return util.construct_adb_shell_command(
-          atrace_args + ['&&'] + setprop_args + ['&&'] + rm_args,
-          self._options.device_serial)
-
-
-class FileReaderThread(threading.Thread):
-  """Reads data from a file/pipe on a worker thread.
-
-  Use the standard threading. Thread object API to start and interact with the
-  thread (start(), join(), etc.).
-  """
-
-  def __init__(self, file_object, output_queue, text_file, chunk_size=-1):
-    """Initializes a FileReaderThread.
-
-    Args:
-      file_object: The file or pipe to read from.
-      output_queue: A Queue.Queue object that will receive the data
-      text_file: If True, the file will be read one line at a time, and
-          chunk_size will be ignored.  If False, line breaks are ignored and
-          chunk_size must be set to a positive integer.
-      chunk_size: When processing a non-text file (text_file = False),
-          chunk_size is the amount of data to copy into the queue with each
-          read operation.  For text files, this parameter is ignored.
-    """
-    threading.Thread.__init__(self)
-    self._file_object = file_object
-    self._output_queue = output_queue
-    self._text_file = text_file
-    self._chunk_size = chunk_size
-    assert text_file or chunk_size > 0
-
-  def run(self):
-    """Overrides Thread's run() function.
-
-    Returns when an EOF is encountered.
-    """
-    if self._text_file:
-      # Read a text file one line at a time.
-      for line in self._file_object:
-        self._output_queue.put(line)
-    else:
-      # Read binary or text data until we get to EOF.
-      while True:
-        chunk = self._file_object.read(self._chunk_size)
-        if not chunk:
-          break
-        self._output_queue.put(chunk)
-
-  def set_chunk_size(self, chunk_size):
-    """Change the read chunk size.
-
-    This function can only be called if the FileReaderThread object was
-    created with an initial chunk_size > 0.
-    Args:
-      chunk_size: the new chunk size for this file.  Must be > 0.
-    """
-    # The chunk size can be changed asynchronously while a file is being read
-    # in a worker thread.  However, type of file can not be changed after the
-    # the FileReaderThread has been created.  These asserts verify that we are
-    # only changing the chunk size, and not the type of file.
-    assert not self._text_file
-    assert chunk_size > 0
-    self._chunk_size = chunk_size
-
-
-def get_default_categories(device_serial):
-  categories_output, return_code = util.run_adb_shell(LIST_CATEGORIES_ARGS,
-                                                    device_serial)
-
-  if return_code == 0 and categories_output:
-    categories = [c.split('-')[0].strip()
-                  for c in categories_output.splitlines()]
-    return [c for c in categories if c in DEFAULT_CATEGORIES]
-
-  return []
-
-
-def status_update(last_update_time):
-  current_time = time.time()
-  if (current_time - last_update_time) >= MIN_TIME_BETWEEN_STATUS_UPDATES:
-    # Gathering a trace may take a while.  Keep printing something so users
-    # don't think the script has hung.
-    sys.stdout.write('.')
-    sys.stdout.flush()
-    return current_time
-
-  return last_update_time
-
-
 def extract_thread_list(trace_text):
   """Removes the thread list from the given trace data.
+
   Args:
     trace_text: The text portion of the trace
   Returns:
@@ -532,6 +336,7 @@ def extract_thread_list(trace_text):
 
 def extract_tgids(trace_text):
   """Removes the procfs dump from the given trace text
+
   Args:
     trace_text: The text portion of the trace
   Returns:
@@ -564,7 +369,8 @@ def strip_and_decompress_trace(trace_data):
     trace_data = trace_data.replace('\r\r\n', '\n')
 
   # Skip the initial newline.
-  trace_data = trace_data[1:]
+  if trace_data[0] == '\n':
+    trace_data = trace_data[1:]
 
   if not trace_data.startswith(TRACE_TEXT_HEADER):
     # No header found, so assume the data is compressed.
@@ -612,6 +418,7 @@ def fix_thread_names(trace_data, thread_names):
 
 def fix_missing_tgids(trace_data, pid2_tgid):
   """Replaces missing TGIDs from the trace data with those found in procfs
+
   Args:
     trace_data: the atrace data
   Returns:
@@ -636,7 +443,7 @@ def fix_missing_tgids(trace_data, pid2_tgid):
 
 
 def fix_circular_traces(out):
-  """Fix inconsistentcies in traces due to circular buffering.
+  """Fix inconsistencies in traces due to circular buffering.
 
   The circular buffers are kept per CPU, so it is not guaranteed that the
   beginning of a slice is overwritten before the end. To work around this, we
@@ -671,22 +478,16 @@ def fix_circular_traces(out):
     out = out[:end_of_header] + out[start_of_full_trace:]
   return out
 
+def do_preprocess_adb_cmd(command, serial):
+  """Run an ADB command for preprocessing of output.
 
-def do_popen(args):
-  try:
-    adb = subprocess.Popen(args, stdout=subprocess.PIPE,
-                              stderr=subprocess.PIPE)
-  except OSError as error:
-    print >> sys.stderr, (
-      'The command "%s" failed with the following error:' %
-      ' '.join(args))
-    print >> sys.stderr, '    ', error
-    sys.exit(1)
-
-  return adb
-
+  Run an ADB command and get the result. This function is used
+  for running commands relating to preprocessing of output data.
 
-def do_preprocess_adb_cmd(command, serial):
+  Args:
+      command: Command to run.
+      serial: Serial number of device.
+  """
   args = [command]
   dump, ret_code = util.run_adb_shell(args, serial)
   if ret_code != 0:
@@ -694,3 +495,50 @@ def do_preprocess_adb_cmd(command, serial):
 
   dump = ''.join(dump)
   return dump
+
+
+class BootAgent(AtraceAgent):
+  """AtraceAgent that specializes in tracing the boot sequence."""
+
+  def __init__(self):
+    super(BootAgent, self).__init__()
+
+  def _StartAgentTracingImpl(self, options, categories):
+    try:
+      setup_args = _construct_setup_command(options, categories)
+      subprocess.check_call(setup_args)
+    except OSError as error:
+      print >> sys.stderr, (
+          'The command "%s" failed with the following error:' %
+          ' '.join(setup_args))
+      print >> sys.stderr, '    ', error
+      sys.exit(1)
+
+  def StartAgentTracing(self, options, categories, timeout=10):
+    return timeout_retry.Run(self._StartAgentTracingImpl,
+                             timeout, 1, args=[options, categories])
+
+  def _dump_trace(self): #called by StopAgentTracing
+    """Dumps the running trace asynchronously and returns the dumped trace."""
+    dump_cmd = get_boot_dump_cmd()
+    return self._device_utils.RunShellCommand(dump_cmd, split_lines=False)
+
+  def _stop_trace(self): # called by _collect_trace_data via StopAgentTracing
+    # pylint: disable=no-self-use
+    # This is a member fucntion for consistency with AtraceAgent
+    pass # don't need to stop separately; already done in dump_trace
+
+def get_boot_dump_cmd():
+  """Constructs the command that is used to dump the boot trace."""
+  return ['atrace', '--async_stop', '&&',
+          'setprop', BOOTTRACE_PROP, '0', '&&',
+          'rm', BOOTTRACE_CATEGORIES]
+
+def _construct_setup_command(options, categories):
+  """Constructs the command that is used to start up the boot trace."""
+  echo_args = ['echo'] + categories + ['>', BOOTTRACE_CATEGORIES]
+  setprop_args = ['setprop', BOOTTRACE_PROP, '1']
+  reboot_args = ['reboot']
+  return util.construct_adb_shell_command(
+      echo_args + ['&&'] + setprop_args + ['&&'] + reboot_args,
+      options.device_serial)