Pass in custom options to Systrace agents

systrace/profile_chrome/atrace_tracing_agent.py

Index: systrace/profile_chrome/atrace_tracing_agent.py
diff --git a/systrace/profile_chrome/atrace_tracing_agent.py b/systrace/profile_chrome/atrace_tracing_agent.py
index c9ce8e446d17a9b5f20a770ba4827fcf0e6655f5..8d9e941e968fddc2052270f5e6cd555362948803 100644
--- a/systrace/profile_chrome/atrace_tracing_agent.py
+++ b/systrace/profile_chrome/atrace_tracing_agent.py
@@ -2,6 +2,7 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import optparse

[Sami, 2016/08/26 15:53:43]

Not something for this patch, but we should consider migrating from optparse to
argparse. Optparse has been deprecated and argparse is available since Python
2.7. It also has a nice subparser concept which seems to match what we need
here.

[washingtonp, 2016/08/26 18:50:54]

I agree with this. The subcommand support, better support for
variable numbers of arguments, and support for positional
arguments have me sold. I will do this in a future CL, and have
created a new Catapult issue for this here:

https://github.com/catapult-project/catapult/issues/2725

 import py_utils
 import threading
 import zlib
@@ -26,14 +27,14 @@ _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'
@@ -43,7 +44,8 @@ class AtraceAgent(tracing_agents.TracingAgent):
     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()
 
@@ -117,3 +119,29 @@ class AtraceAgent(tracing_agents.TracingAgent):
 
     # Skip the initial newline.
     return trace_data[1:]
+

[Sami, 2016/08/26 15:53:43]

nit: two blank lines around top level entries please (here and elsewhere).

[washingtonp, 2016/08/26 18:50:54]

Done.

+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(options):
+  if not options.atrace_categories:
+    return []
+  return options.atrace_categories.split(',')