1# Copyright 2014 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import time
6
7from profile_chrome import chrome_startup_tracing_agent
8from profile_chrome import chrome_tracing_agent
9from profile_chrome import ui
10from profile_chrome import util
11from systrace import output_generator
12from systrace import tracing_controller
13
14
15def _GetResults(trace_results, controller, output, compress, write_json,
16                interval):
17  ui.PrintMessage('Downloading...')
18
19  # Wait for the trace file to get written.
20  time.sleep(1)
21
22  for agent in controller.get_child_agents:
23    if isinstance(agent, chrome_tracing_agent.ChromeTracingAgent):
24      time.sleep(interval / 4)
25
26  # Ignore the systraceController because it will not contain any results,
27  # instead being in charge of collecting results.
28  trace_results = [x for x in controller.all_results if not (x.source_name ==
29      'systraceController')]
30
31  if not trace_results:
32    ui.PrintMessage('No results')
33    return ''
34
35  result = None
36  trace_results = output_generator.MergeTraceResultsIfNeeded(trace_results)
37  if not write_json:
38    ui.PrintMessage('Writing trace HTML...')
39    html_file = trace_results[0].source_name + '.html'
40    result = output_generator.GenerateHTMLOutput(trace_results, html_file)
41    ui.PrintMessage('\nWrote file://%s' % result)
42  elif compress and len(trace_results) == 1:
43    result = output or trace_results[0].source_name + '.gz'
44    util.WriteDataToCompressedFile(trace_results[0].raw_data, result)
45  elif len(trace_results) > 1:
46    result = (output or 'chrome-combined-trace-%s.zip' %
47              util.GetTraceTimestamp())
48    util.ArchiveData(trace_results, result)
49  elif output:
50    result = output
51    with open(result, 'wb') as f:
52      f.write(trace_results[0].raw_data)
53  else:
54    result = trace_results[0].source_name
55    with open(result, 'wb') as f:
56      f.write(trace_results[0].raw_data)
57
58  return result
59
60
61def CaptureProfile(options, interval, modules, output=None,
62                   compress=False, write_json=False):
63  """Records a profiling trace saves the result to a file.
64
65  Args:
66    options: Command line options.
67    interval: Time interval to capture in seconds. An interval of None (or 0)
68        continues tracing until stopped by the user.
69    modules: The list of modules to initialize the tracing controller with.
70    output: Output file name or None to use an automatically generated name.
71    compress: If True, the result will be compressed either with gzip or zip
72        depending on the number of captured subtraces.
73    write_json: If True, prefer JSON output over HTML.
74
75  Returns:
76    Path to saved profile.
77  """
78  agents_with_config = tracing_controller.CreateAgentsWithConfig(options,
79                                                                 modules)
80  if chrome_startup_tracing_agent in modules:
81    controller_config = tracing_controller.GetChromeStartupControllerConfig(
82        options)
83  else:
84    controller_config = tracing_controller.GetControllerConfig(options)
85  controller = tracing_controller.TracingController(agents_with_config,
86                                                    controller_config)
87  try:
88    result = controller.StartTracing()
89    trace_type = controller.GetTraceType()
90    if not result:
91      ui.PrintMessage('Trace starting failed.')
92    if interval:
93      ui.PrintMessage(('Capturing %d-second %s. Press Enter to stop early...' %
94                     (interval, trace_type)), eol='')
95      ui.WaitForEnter(interval)
96    else:
97      ui.PrintMessage('Capturing %s. Press Enter to stop...' % trace_type,
98                      eol='')
99      raw_input()
100
101    ui.PrintMessage('Stopping...')
102    all_results = controller.StopTracing()
103  finally:
104    if interval:
105      ui.PrintMessage('done')
106
107  return _GetResults(all_results, controller, output, compress, write_json,
108                     interval)
109