Pass in custom options to Systrace agents

systrace/profile_chrome/chrome_tracing_agent.py

Index: systrace/profile_chrome/chrome_tracing_agent.py
diff --git a/systrace/profile_chrome/chrome_tracing_agent.py b/systrace/profile_chrome/chrome_tracing_agent.py
index b13f5ea2a187eeb56b9993771f811ee1e4a50659..18f1bb499f81212d39441faa283782c972095112 100644
--- a/systrace/profile_chrome/chrome_tracing_agent.py
+++ b/systrace/profile_chrome/chrome_tracing_agent.py
@@ -3,6 +3,7 @@
 # found in the LICENSE file.
 
 import json
+import optparse
 import os
 import py_utils
 import re
@@ -14,16 +15,15 @@ from systrace import trace_result
 from systrace import tracing_agents
 
 
+_DEFAULT_CHROME_CATEGORIES = '_DEFAULT_CHROME_CATEGORIES'
 _HEAP_PROFILE_MMAP_PROPERTY = 'heapprof.mmap'
 
 
 class ChromeTracingAgent(tracing_agents.TracingAgent):
-  def __init__(self, device, package_info,
-               categories, ring_buffer, trace_memory=False):
+  def __init__(self, device, package_info, ring_buffer, trace_memory=False):
     tracing_agents.TracingAgent.__init__(self)
     self._device = device
     self._package_info = package_info
-    self._categories = categories
     self._ring_buffer = ring_buffer
     self._logcat_monitor = self._device.GetLogcatMonitor()
     self._trace_file = None
@@ -33,6 +33,7 @@ class ChromeTracingAgent(tracing_agents.TracingAgent):
        re.compile(r'Logging performance trace to file')
     self._trace_finish_re = \
        re.compile(r'Profiler finished[.] Results are in (.*)[.]')
+    self._categories = None
 
   def __repr__(self):
     return 'chrome trace'
@@ -62,7 +63,8 @@ class ChromeTracingAgent(tracing_agents.TracingAgent):
     return list(record_categories), list(disabled_by_default_categories)
 
   @py_utils.Timeout(tracing_agents.START_STOP_TIMEOUT)
-  def StartAgentTracing(self, options, categories, timeout=None):
+  def StartAgentTracing(self, config, timeout=None):
+    self._categories = _ComputeChromeCategories(config)
     self._logcat_monitor.Start()
     start_extras = {'categories': ','.join(self._categories)}
     if self._ring_buffer:
@@ -125,3 +127,74 @@ class ChromeTracingAgent(tracing_agents.TracingAgent):
   def RecordClockSyncMarker(self, sync_id, did_record_sync_marker_callback):
     assert self.SupportsExplicitClockSync(), ('Clock sync marker cannot be '
         'recorded since explicit clock sync is not supported.')
+
+class ChromeConfig(tracing_agents.TracingConfig):
+  def __init__(self, chrome_categories, trace_cc, trace_frame_viewer,
+               trace_ubercompositor, trace_gpu, trace_flow, trace_memory,
+               trace_scheduler):
+    tracing_agents.TracingConfig.__init__(self)
+    self.chrome_categories = chrome_categories
+    self.trace_cc = trace_cc
+    self.trace_frame_viewer = trace_frame_viewer
+    self.trace_ubercompositor = trace_ubercompositor
+    self.trace_gpu = trace_gpu
+    self.trace_flow = trace_flow
+    self.trace_memory = trace_memory
+    self.trace_scheduler = trace_scheduler
+
+def add_options(parser):
+  chrome_opts = optparse.OptionGroup(parser, 'Chrome tracing options')
+  chrome_opts.add_option('-c', '--categories', help='Select Chrome tracing '
+                         'categories with comma-delimited wildcards, '
+                         'e.g., "*", "cat1*,-cat1a". Omit this option to trace '
+                         'Chrome\'s default categories. Chrome tracing can be '
+                         'disabled with "--categories=\'\'". Use "list" to '
+                         'see the available categories.',
+                         metavar='CHROME_CATEGORIES', dest='chrome_categories',
+                         default=_DEFAULT_CHROME_CATEGORIES)
+  chrome_opts.add_option('--trace-cc',
+                         help='Deprecated, use --trace-frame-viewer.',
+                         action='store_true')
+  chrome_opts.add_option('--trace-frame-viewer',
+                         help='Enable enough trace categories for '
+                         'compositor frame viewing.', action='store_true')
+  chrome_opts.add_option('--trace-ubercompositor',
+                         help='Enable enough trace categories for '
+                         'ubercompositor frame data.', action='store_true')
+  chrome_opts.add_option('--trace-gpu', help='Enable extra trace categories '
+                         'for GPU data.', action='store_true')
+  chrome_opts.add_option('--trace-flow', help='Enable extra trace categories '
+                         'for IPC message flows.', action='store_true')
+  chrome_opts.add_option('--trace-memory', help='Enable extra trace categories '
+                         'for memory profile. (tcmalloc required)',
+                         action='store_true')
+  chrome_opts.add_option('--trace-scheduler', help='Enable extra trace '
+                         'categories for scheduler state',
+                         action='store_true')
+  return chrome_opts
+
+def get_config(options):
+  return ChromeConfig(options.chrome_categories, options.trace_cc,
+                      options.trace_frame_viewer, options.trace_ubercompositor,
+                      options.trace_gpu, options.trace_flow,
+                      options.trace_memory, options.trace_scheduler)
+
+def _ComputeChromeCategories(options):
+  categories = []
+  if options.trace_frame_viewer:
+    categories.append('disabled-by-default-cc.debug')
+  if options.trace_ubercompositor:
+    categories.append('disabled-by-default-cc.debug*')
+  if options.trace_gpu:
+    categories.append('disabled-by-default-gpu.debug*')
+  if options.trace_flow:
+    categories.append('disabled-by-default-toplevel.flow')
+  if options.trace_memory:
+    categories.append('disabled-by-default-memory')
+  if options.trace_scheduler:
+    categories.append('disabled-by-default-blink.scheduler')
+    categories.append('disabled-by-default-cc.debug.scheduler')
+    categories.append('disabled-by-default-renderer.scheduler')
+  if options.chrome_categories:
+    categories += options.chrome_categories.split(',')
+  return categories