Pass in custom options to Systrace agents

systrace/profile_chrome/atrace_tracing_agent.py

 # 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 optparse
 import py_utils
 import threading
 import zlib
 
 from devil.utils import cmd_helper
 
 from systrace import trace_result
 from systrace import tracing_agents
 
 
 _ATRACE_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 atrace data.
 _ATRACE_INTERVAL = 15
 
+# If a custom list of categories is not specified, traces will include
+# these categories (if available on the device).
+_DEFAULT_CATEGORIES = 'sched gfx view dalvik webview input disk am wm'.split()
+
 _TRACING_ON_PATH = '/sys/kernel/debug/tracing/tracing_on'
 
 
 class AtraceAgent(tracing_agents.TracingAgent):
-  def __init__(self, device, categories, ring_buffer):
+  def __init__(self, device, ring_buffer):
     tracing_agents.TracingAgent.__init__(self)
     self._device = device
-    self._categories = categories
     self._ring_buffer = ring_buffer
     self._done = threading.Event()
     self._thread = None
     self._trace_data = None
+    self._categories = None
 
   def __repr__(self):
     return 'atrace'
 
   @staticmethod
   def GetCategories(device):
     return device.RunShellCommand('atrace --list_categories')
 
   @py_utils.Timeout(tracing_agents.START_STOP_TIMEOUT)
-  def StartAgentTracing(self, options, categories, timeout=None):
+  def StartAgentTracing(self, config, timeout=None):
+    self._categories = _ComputeAtraceCategories(config)
     self._thread = threading.Thread(target=self._CollectData)
     self._thread.start()
 
   @py_utils.Timeout(tracing_agents.START_STOP_TIMEOUT)
   def StopAgentTracing(self, timeout=None):
     self._done.set()
 
   @py_utils.Timeout(tracing_agents.GET_RESULTS_TIMEOUT)
   def GetResults(self, timeout=None):
     self._thread.join()
@@ skipping 53 matching lines @@
     except ValueError:
       raise RuntimeError('Atrace 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:]
+
+
+class AtraceConfig(tracing_agents.TracingConfig):
+  def __init__(self, atrace_categories):
+    tracing_agents.TracingConfig.__init__(self)
+    self.atrace_categories = atrace_categories
+
+
+def add_options(parser):
+  atrace_opts = optparse.OptionGroup(parser, 'Atrace tracing options')
+  atrace_opts.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. Note that in this case, Systrace is '
+                           'synonymous with Atrace.',
+                           metavar='ATRACE_CATEGORIES',
+                           dest='atrace_categories', default='')
+  return atrace_opts
+
+def get_config(options):
+  return AtraceConfig(options.atrace_categories)
+
+def _ComputeAtraceCategories(config):
+  if not config.atrace_categories:
+    return _DEFAULT_CATEGORIES
+  return config.atrace_categories.split(',')