Pass in custom options to Systrace agents

systrace/profile_chrome/perf_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 logging
+import optparse
 import os
 import py_utils
 import signal
 import subprocess
 import sys
 import tempfile
 
 from devil.android import device_temp_file
 from devil.android.perf import perf_control
 
@@ skipping 70 matching lines @@
     if not os.stat(perf_profile).st_size:
       os.remove(perf_profile)
       self._FailWithLog('Perf recorded a zero-sized file')
 
     self._log_file.close()
     self._output_file.close()
     return perf_profile
 
 
 class PerfProfilerAgent(tracing_agents.TracingAgent):
-  def __init__(self, device, categories):
+  def __init__(self, device):
     tracing_agents.TracingAgent.__init__(self)
     self._device = device
-    self._categories = categories
     self._perf_binary = self._PrepareDevice(device)
     self._perf_instance = None
+    self._categories = None
 
   def __repr__(self):
     return 'perf profile'
 
   @staticmethod
   def IsSupported():
     return bool(android_profiling_helper)
 
   @staticmethod
   def _PrepareDevice(device):
     if not 'BUILDTYPE' in os.environ:
       os.environ['BUILDTYPE'] = 'Release'
     if binary_manager.NeedsInit():
       binary_manager.InitDependencyManager(None)
     return android_profiling_helper.PrepareDeviceForPerf(device)
 
   @classmethod
   def GetCategories(cls, device):
     perf_binary = cls._PrepareDevice(device)
     return device.RunShellCommand('%s list' % perf_binary)
 
   @py_utils.Timeout(tracing_agents.START_STOP_TIMEOUT)
-  def StartAgentTracing(self, options, categories, timeout=None):
+  def StartAgentTracing(self, config, timeout=None):
+    self._categories = _ComputePerfCategories(config)
     self._perf_instance = _PerfProfiler(self._device,
                                         self._perf_binary,
                                         self._categories)
 
   @py_utils.Timeout(tracing_agents.START_STOP_TIMEOUT)
   def StopAgentTracing(self, timeout=None):
     if not self._perf_instance:
       return
     self._perf_instance.SignalAndWait()
 
@@ skipping 62 matching lines @@
         return None
 
     return json_file_name
 
   def SupportsExplicitClockSync(self):
     return False
 
   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.')
+
+def _OptionalValueCallback(default_value):
+  def callback(option, _, __, parser):  # pylint: disable=unused-argument
+    value = default_value
+    if parser.rargs and not parser.rargs[0].startswith('-'):
+      value = parser.rargs.pop(0)
+    setattr(parser.values, option.dest, value)
+  return callback
+
+
+class PerfConfig(tracing_agents.TracingConfig):
+  def __init__(self, perf_categories):
+    tracing_agents.TracingConfig.__init__(self)
+    self.perf_categories = perf_categories
+
+
+def add_options(parser):
+  options = optparse.OptionGroup(parser, 'Perf profiling options')
+  options.add_option('-p', '--perf', help='Capture a perf profile with '
+                     'the chosen comma-delimited event categories. '
+                     'Samples CPU cycles by default. Use "list" to see '
+                     'the available sample types.', action='callback',
+                     default='', callback=_OptionalValueCallback('cycles'),
+                     metavar='PERF_CATEGORIES', dest='perf_categories')
+  parser.add_option_group(options)
+  return options
+
+def get_config(options):
+  return PerfConfig(options.perf_categories)
+
+def _ComputePerfCategories(config):
+  if not PerfProfilerAgent.IsSupported():
+    return []
+  if not config.perf_categories:
+    return []
+  return config.perf_categories.split(',')