Use Systrace tracing controller in profile_chrome

systrace/profile_chrome/profiler.py

Index: systrace/profile_chrome/profiler.py
diff --git a/systrace/profile_chrome/profiler.py b/systrace/profile_chrome/profiler.py
index 4c126783728a983e9310e8aef3948d6754c792e6..3a3bf42168d1b9cadfe5fb4ee537fd4804480bae 100644
--- a/systrace/profile_chrome/profiler.py
+++ b/systrace/profile_chrome/profiler.py
@@ -6,50 +6,34 @@ import time
 
 from devil.android.constants import chrome
 from profile_chrome import atrace_tracing_agent
-from profile_chrome import chrome_startup_tracing_agent
 from profile_chrome import chrome_tracing_agent
 from profile_chrome import ddms_tracing_agent
 from profile_chrome import perf_tracing_agent
 from profile_chrome import ui
 from profile_chrome import util
 from systrace import output_generator
+from systrace import tracing_controller
 
 
-# TODO(washingtonp): This mapping is temporarily in place because
-# profile_chrome does not currently follow Systrace's controller API. This
-# mapping will be removed in the CL that makes profile_chrome follow the
-# Systrace API.
-AGENT_TYPE_TO_MODULE = {'chrome trace': chrome_tracing_agent,
-                        'Browser Startup Trace': chrome_startup_tracing_agent,
-                        'ddms profile': ddms_tracing_agent,
-                        'perf profile': perf_tracing_agent,
-                        'atrace': atrace_tracing_agent}
+AGENT_MODULES = [atrace_tracing_agent, chrome_tracing_agent,
+                 ddms_tracing_agent, perf_tracing_agent]
 
 
-def _StartTracing(agents, options):
-  for agent in agents:
-    if repr(agent) not in AGENT_TYPE_TO_MODULE:
-      continue
-    agent.StartAgentTracing(AGENT_TYPE_TO_MODULE[repr(agent)].
-                            get_config(options))
-
-
-def _StopTracing(agents):
-  for agent in agents:
-    agent.StopAgentTracing()
-
-
-def _GetResults(agents, output, compress, write_json, interval):
+def _GetResults(trace_results, controller, output, compress, write_json,
+                interval):
   ui.PrintMessage('Downloading...', eol='')
 
   # Wait for the trace file to get written.
   time.sleep(1)
 
-  trace_results = []
-  for agent in agents:
+  for agent in controller.child_agents:
     if isinstance(agent, chrome_tracing_agent.ChromeTracingAgent):
       time.sleep(interval / 4)
-    trace_results.append(agent.GetResults())
+
+  # Ignore the systraceController because it will not contain any results,
+  # instead being in charge of collecting results.
+  trace_results = [x for x in controller.all_results if not (x.source_name ==
+      'systraceController')]
 
   if not trace_results:
     ui.PrintMessage('No results')
@@ -94,13 +78,12 @@ def GetSupportedBrowsers():
   return supported_browsers
 
 
-def CaptureProfile(options, agents, interval, output=None, compress=False,
-                   write_json=False):
+def CaptureProfile(options, interval, output=None,
+                   compress=False, write_json=False):
   """Records a profiling trace saves the result to a file.
 
   Args:
     options: Command line options.
-    agents: List of tracing agents.
     interval: Time interval to capture in seconds. An interval of None (or 0)
         continues tracing until stopped by the user.
     output: Output file name or None to use an automatically generated name.
@@ -111,20 +94,28 @@ def CaptureProfile(options, agents, interval, output=None, compress=False,
   Returns:
     Path to saved profile.
   """
-  trace_type = ' + '.join(map(str, agents))
+  agents_with_config = tracing_controller.CreateAgentsWithConfig(options,
+                                                                 AGENT_MODULES)
+  controller_config = tracing_controller.GetControllerConfig(options)
+  controller = tracing_controller.TracingController(agents_with_config,
+                                                    controller_config)
   try:
-    _StartTracing(agents, options)
+    result = controller.StartTracing()
+    trace_type = controller.GetTraceType()
+    if not result:
+      print 'Trace starting failed.'
     if interval:
       ui.PrintMessage(('Capturing %d-second %s. Press Enter to stop early...' %
-                      (interval, trace_type)), eol='')
+                     (interval, trace_type)), eol='')
       ui.WaitForEnter(interval)
     else:
       ui.PrintMessage('Capturing %s. Press Enter to stop...' % trace_type,
-                      eol='')
+                     eol='')
       raw_input()
+    all_results = controller.StopTracing()
   finally:
-    _StopTracing(agents)
-  if interval:
-    ui.PrintMessage('done')
+    if interval:
+      ui.PrintMessage('done')
 
-  return _GetResults(agents, output, compress, write_json, interval)
+  return _GetResults(all_results, controller, output, compress, write_json,
+                     interval)