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 threading
6import zlib
7
8from devil.utils import cmd_helper
9
10from profile_chrome import controllers
11from profile_chrome import util
12
13
14_SYSTRACE_OPTIONS = [
15    # Compress the trace before sending it over USB.
16    '-z',
17    # Use a large trace buffer to increase the polling interval.
18    '-b', '16384'
19]
20
21# Interval in seconds for sampling systrace data.
22_SYSTRACE_INTERVAL = 15
23
24_TRACING_ON_PATH = '/sys/kernel/debug/tracing/tracing_on'
25
26
27class SystraceController(controllers.BaseController):
28  def __init__(self, device, categories, ring_buffer):
29    controllers.BaseController.__init__(self)
30    self._device = device
31    self._categories = categories
32    self._ring_buffer = ring_buffer
33    self._done = threading.Event()
34    self._thread = None
35    self._trace_data = None
36
37  def __repr__(self):
38    return 'systrace'
39
40  @staticmethod
41  def GetCategories(device):
42    return device.RunShellCommand('atrace --list_categories')
43
44  def StartTracing(self, _):
45    self._thread = threading.Thread(target=self._CollectData)
46    self._thread.start()
47
48  def StopTracing(self):
49    self._done.set()
50
51  def PullTrace(self):
52    self._thread.join()
53    self._thread = None
54    if self._trace_data:
55      output_name = 'systrace-%s' % util.GetTraceTimestamp()
56      with open(output_name, 'w') as out:
57        out.write(self._trace_data)
58      return output_name
59
60  def IsTracingOn(self):
61    result = self._RunAdbShellCommand(['cat', _TRACING_ON_PATH])
62    return result.strip() == '1'
63
64  def _RunAdbShellCommand(self, command):
65    # We use a separate interface to adb because the one from AndroidCommands
66    # isn't re-entrant.
67    # TODO(jbudorick) Look at providing a way to unhandroll this once the
68    #                 adb rewrite has fully landed.
69    device_param = (['-s', str(self._device)] if str(self._device) else [])
70    cmd = ['adb'] + device_param + ['shell'] + command
71    return cmd_helper.GetCmdOutput(cmd)
72
73  def _RunATraceCommand(self, command):
74    cmd = ['atrace', '--%s' % command] + _SYSTRACE_OPTIONS + self._categories
75    return self._RunAdbShellCommand(cmd)
76
77  def _ForceStopAtrace(self):
78    # atrace on pre-M Android devices cannot be stopped asynchronously
79    # correctly. Use synchronous mode to force stop.
80    cmd = ['atrace', '-t', '0']
81    return self._RunAdbShellCommand(cmd)
82
83  def _CollectData(self):
84    trace_data = []
85    self._RunATraceCommand('async_start')
86    try:
87      while not self._done.is_set():
88        self._done.wait(_SYSTRACE_INTERVAL)
89        if not self._ring_buffer or self._done.is_set():
90          trace_data.append(
91              self._DecodeTraceData(self._RunATraceCommand('async_dump')))
92    finally:
93      trace_data.append(
94          self._DecodeTraceData(self._RunATraceCommand('async_stop')))
95      if self.IsTracingOn():
96        self._ForceStopAtrace()
97    self._trace_data = ''.join([zlib.decompress(d) for d in trace_data])
98
99  @staticmethod
100  def _DecodeTraceData(trace_data):
101    try:
102      trace_start = trace_data.index('TRACE:')
103    except ValueError:
104      raise RuntimeError('Systrace start marker not found')
105    trace_data = trace_data[trace_start + 6:]
106
107    # Collapse CRLFs that are added by adb shell.
108    if trace_data.startswith('\r\n'):
109      trace_data = trace_data.replace('\r\n', '\n')
110
111    # Skip the initial newline.
112    return trace_data[1:]
113