1# Copyright (c) 2015 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 Queue
6import re
7import subprocess
8import sys
9import threading
10import time
11import zlib
12
13import systrace_agent
14import util
15
16# Text that ADB sends, but does not need to be displayed to the user.
17ADB_IGNORE_REGEXP = r'^capturing trace\.\.\. done|^capturing trace\.\.\.'
18# The number of seconds to wait on output from ADB.
19ADB_STDOUT_READ_TIMEOUT = 0.2
20# The adb shell command to initiate a trace.
21ATRACE_BASE_ARGS = ['atrace']
22# If a custom list of categories is not specified, traces will include
23# these categories (if available on the device).
24DEFAULT_CATEGORIES = 'sched gfx view dalvik webview input disk am wm'.split()
25# The command to list trace categories.
26LIST_CATEGORIES_ARGS = ATRACE_BASE_ARGS + ['--list_categories']
27# Minimum number of seconds between displaying status updates.
28MIN_TIME_BETWEEN_STATUS_UPDATES = 0.2
29# ADB sends this text to indicate the beginning of the trace data.
30TRACE_START_REGEXP = r'TRACE\:'
31# Plain-text trace data should always start with this string.
32TRACE_TEXT_HEADER = '# tracer'
33
34# This list is based on the tags in frameworks/native/include/utils/Trace.h for
35# legacy platform.
36LEGACY_TRACE_TAG_BITS = (
37  ('gfx',       1<<1),
38  ('input',     1<<2),
39  ('view',      1<<3),
40  ('webview',   1<<4),
41  ('wm',        1<<5),
42  ('am',        1<<6),
43  ('sm',        1<<7),
44  ('audio',     1<<8),
45  ('video',     1<<9),
46  ('camera',    1<<10),
47)
48
49
50def try_create_agent(options, categories):
51  if options.from_file is not None:
52    return AtraceAgent(options, categories)
53
54  device_sdk_version = util.get_device_sdk_version()
55  if device_sdk_version >= 18:
56    return AtraceAgent(options, categories)
57  elif device_sdk_version >= 16:
58    return AtraceLegacyAgent(options, categories)
59
60
61class AtraceAgent(systrace_agent.SystraceAgent):
62  def __init__(self, options, categories):
63    super(AtraceAgent, self).__init__(options, categories)
64    self._expect_trace = False
65    self._adb = None
66    self._trace_data = None
67
68  def start(self):
69    tracer_args = self._construct_trace_command()
70    try:
71      self._adb = subprocess.Popen(tracer_args, stdout=subprocess.PIPE,
72                                   stderr=subprocess.PIPE)
73    except OSError as error:
74      print >> sys.stderr, (
75          'The command "%s" failed with the following error:' %
76          ' '.join(tracer_args))
77      print >> sys.stderr, '    ', error
78      sys.exit(1)
79
80  def collect_result(self):
81    trace_data = self._collect_trace_data();
82    if self._expect_trace:
83      self._trace_data = self._preprocess_trace_data(trace_data);
84
85  def expect_trace(self):
86    return self._expect_trace
87
88  def get_trace_data(self):
89    return self._trace_data
90
91  def get_class_name(self):
92    return 'trace-data'
93
94  def _construct_list_categories_command(self):
95    return util.construct_adb_shell_command(
96          LIST_CATEGORIES_ARGS, self._options.device_serial)
97
98  def _construct_extra_trace_command(self):
99    extra_args = []
100    if self._options.app_name is not None:
101      extra_args.extend(['-a', self._options.app_name])
102
103    if self._options.kfuncs is not None:
104      extra_args.extend(['-k', self._options.kfuncs])
105
106    if not self._categories:
107      self._categories = get_default_categories(self._options.device_serial)
108    extra_args.extend(self._categories)
109    return extra_args
110
111  def _construct_trace_command(self):
112    """Builds a command-line used to invoke a trace process.
113
114    Returns:
115      A tuple where the first element is an array of command-line arguments, and
116      the second element is a boolean which will be true if the commend will
117      stream trace data.
118    """
119    if self._options.list_categories:
120      tracer_args = self._construct_list_categories_command()
121      self._expect_trace = False
122    elif self._options.from_file is not None:
123      tracer_args = ['cat', self._options.from_file]
124      self._expect_trace = True
125    else:
126      atrace_args = ATRACE_BASE_ARGS[:]
127      self._expect_trace = True
128      if self._options.compress_trace_data:
129        atrace_args.extend(['-z'])
130
131      if ((self._options.trace_time is not None)
132          and (self._options.trace_time > 0)):
133        atrace_args.extend(['-t', str(self._options.trace_time)])
134
135      if ((self._options.trace_buf_size is not None)
136          and (self._options.trace_buf_size > 0)):
137        atrace_args.extend(['-b', str(self._options.trace_buf_size)])
138      extra_args = self._construct_extra_trace_command()
139      atrace_args.extend(extra_args)
140
141      if self._options.fix_threads:
142        atrace_args.extend([';', 'ps', '-t'])
143      tracer_args = util.construct_adb_shell_command(
144          atrace_args, self._options.device_serial)
145
146    return tracer_args
147
148  def _collect_trace_data(self):
149    # Read the output from ADB in a worker thread.  This allows us to monitor
150    # the progress of ADB and bail if ADB becomes unresponsive for any reason.
151
152    # Limit the stdout_queue to 128 entries because we will initially be reading
153    # one byte at a time.  When the queue fills up, the reader thread will
154    # block until there is room in the queue.  Once we start downloading the
155    # trace data, we will switch to reading data in larger chunks, and 128
156    # entries should be plenty for that purpose.
157    stdout_queue = Queue.Queue(maxsize=128)
158    stderr_queue = Queue.Queue()
159
160    if self._expect_trace:
161      # Use stdout.write() (here and for the rest of this function) instead
162      # of print() to avoid extra newlines.
163      sys.stdout.write('Capturing trace...')
164
165    # Use a chunk_size of 1 for stdout so we can display the output to
166    # the user without waiting for a full line to be sent.
167    stdout_thread = FileReaderThread(self._adb.stdout, stdout_queue,
168                                     text_file=False, chunk_size=1)
169    stderr_thread = FileReaderThread(self._adb.stderr, stderr_queue,
170                                     text_file=True)
171    stdout_thread.start()
172    stderr_thread.start()
173
174    # Holds the trace data returned by ADB.
175    trace_data = []
176    # Keep track of the current line so we can find the TRACE_START_REGEXP.
177    current_line = ''
178    # Set to True once we've received the TRACE_START_REGEXP.
179    reading_trace_data = False
180
181    last_status_update_time = time.time()
182
183    while (stdout_thread.isAlive() or stderr_thread.isAlive() or
184           not stdout_queue.empty() or not stderr_queue.empty()):
185      if self._expect_trace:
186        last_status_update_time = status_update(last_status_update_time)
187
188      while not stderr_queue.empty():
189        # Pass along errors from adb.
190        line = stderr_queue.get()
191        sys.stderr.write(line)
192
193      # Read stdout from adb.  The loop exits if we don't get any data for
194      # ADB_STDOUT_READ_TIMEOUT seconds.
195      while True:
196        try:
197          chunk = stdout_queue.get(True, ADB_STDOUT_READ_TIMEOUT)
198        except Queue.Empty:
199          # Didn't get any data, so exit the loop to check that ADB is still
200          # alive and print anything sent to stderr.
201          break
202
203        if reading_trace_data:
204          # Save, but don't print, the trace data.
205          trace_data.append(chunk)
206        else:
207          if not self._expect_trace:
208            sys.stdout.write(chunk)
209          else:
210            # Buffer the output from ADB so we can remove some strings that
211            # don't need to be shown to the user.
212            current_line += chunk
213            if re.match(TRACE_START_REGEXP, current_line):
214              # We are done capturing the trace.
215              sys.stdout.write('Done.\n')
216              # Now we start downloading the trace data.
217              sys.stdout.write('Downloading trace...')
218              current_line = ''
219              # Use a larger chunk size for efficiency since we no longer
220              # need to worry about parsing the stream.
221              stdout_thread.set_chunk_size(4096)
222              reading_trace_data = True
223            elif chunk == '\n' or chunk == '\r':
224              # Remove ADB output that we don't care about.
225              current_line = re.sub(ADB_IGNORE_REGEXP, '', current_line)
226              if len(current_line) > 1:
227                # ADB printed something that we didn't understand, so show it
228                # it to the user (might be helpful for debugging).
229                sys.stdout.write(current_line)
230              # Reset our current line.
231              current_line = ''
232
233    if self._expect_trace:
234      if reading_trace_data:
235        # Indicate to the user that the data download is complete.
236        sys.stdout.write('Done.\n')
237      else:
238        # We didn't receive the trace start tag, so something went wrong.
239        sys.stdout.write('ERROR.\n')
240        # Show any buffered ADB output to the user.
241        current_line = re.sub(ADB_IGNORE_REGEXP, '', current_line)
242        if current_line:
243          sys.stdout.write(current_line)
244          sys.stdout.write('\n')
245
246    # The threads should already have stopped, so this is just for cleanup.
247    stdout_thread.join()
248    stderr_thread.join()
249
250    self._adb.stdout.close()
251    self._adb.stderr.close()
252
253    # The adb process should be done since it's io pipes are closed.  Call
254    # poll() to set the returncode.
255    self._adb.poll()
256
257    if self._adb.returncode != 0:
258      print >> sys.stderr, ('The command "%s" returned error code %d.' %
259                            (' '.join(tracer_args), self._adb.returncode))
260      sys.exit(1)
261
262    return trace_data
263
264  def _preprocess_trace_data(self, trace_data):
265    """Performs various processing on atrace data.
266
267    Args:
268      trace_data: The raw trace data.
269    Returns:
270      The processed trace data.
271    """
272    trace_data = ''.join(trace_data)
273
274    if self._options.fix_threads:
275      # Extract the thread list dumped by ps.
276      trace_data, thread_names = extract_thread_list(trace_data)
277
278    if trace_data:
279      trace_data = strip_and_decompress_trace(trace_data)
280
281    if not trace_data:
282      print >> sys.stderr, ('No data was captured.  Output file was not '
283                            'written.')
284      sys.exit(1)
285
286    if self._options.fix_threads:
287      trace_data = fix_thread_names(trace_data, thread_names)
288
289    if self._options.fix_circular:
290      trace_data = fix_circular_traces(trace_data)
291
292    return trace_data
293
294
295class AtraceLegacyAgent(AtraceAgent):
296  def _construct_list_categories_command(self):
297    LEGACY_CATEGORIES = """       sched - CPU Scheduling
298        freq - CPU Frequency
299        idle - CPU Idle
300        load - CPU Load
301        disk - Disk I/O (requires root)
302         bus - Bus utilization (requires root)
303   workqueue - Kernel workqueues (requires root)"""
304    return ["echo", LEGACY_CATEGORIES]
305
306  def start(self):
307    super(AtraceLegacyAgent, self).start()
308    if self.expect_trace():
309      SHELL_ARGS = ['getprop', 'debug.atrace.tags.enableflags']
310      output, return_code = util.run_adb_shell(SHELL_ARGS, self._options.device_serial)
311      flags = 0
312      if return_code == 0:
313        try:
314          if output.startswith('0x'):
315            flags = int(output, 16)
316          elif output.startswith('0'):
317            flags = int(output, 8)
318          else:
319            flags = int(output)
320        except ValueError, e:
321          pass
322
323      if flags:
324        tags = []
325        for desc, bit in LEGACY_TRACE_TAG_BITS:
326          if bit & flags:
327            tags.append(desc)
328        categories = tags + self._categories
329        print 'Collecting data with following categories:', ' '.join(categories)
330
331  def _construct_extra_trace_command(self):
332    extra_args = []
333    if not self._categories:
334      self._categories = ['sched', ]
335    if 'sched' in self._categories:
336      extra_args.append('-s')
337    if 'freq' in self._categories:
338      extra_args.append('-f')
339    if 'idle' in self._categories:
340      extra_args.append('-i')
341    if 'load' in self._categories:
342      extra_args.append('-l')
343    if 'disk' in self._categories:
344      extra_args.append('-d')
345    if 'bus' in self._categories:
346      extra_args.append('-u')
347    if 'workqueue' in self._categories:
348      extra_args.append('-w')
349
350    return extra_args
351
352
353class FileReaderThread(threading.Thread):
354  """Reads data from a file/pipe on a worker thread.
355
356  Use the standard threading. Thread object API to start and interact with the
357  thread (start(), join(), etc.).
358  """
359
360  def __init__(self, file_object, output_queue, text_file, chunk_size=-1):
361    """Initializes a FileReaderThread.
362
363    Args:
364      file_object: The file or pipe to read from.
365      output_queue: A Queue.Queue object that will receive the data
366      text_file: If True, the file will be read one line at a time, and
367          chunk_size will be ignored.  If False, line breaks are ignored and
368          chunk_size must be set to a positive integer.
369      chunk_size: When processing a non-text file (text_file = False),
370          chunk_size is the amount of data to copy into the queue with each
371          read operation.  For text files, this parameter is ignored.
372    """
373    threading.Thread.__init__(self)
374    self._file_object = file_object
375    self._output_queue = output_queue
376    self._text_file = text_file
377    self._chunk_size = chunk_size
378    assert text_file or chunk_size > 0
379
380  def run(self):
381    """Overrides Thread's run() function.
382
383    Returns when an EOF is encountered.
384    """
385    if self._text_file:
386      # Read a text file one line at a time.
387      for line in self._file_object:
388        self._output_queue.put(line)
389    else:
390      # Read binary or text data until we get to EOF.
391      while True:
392        chunk = self._file_object.read(self._chunk_size)
393        if not chunk:
394          break
395        self._output_queue.put(chunk)
396
397  def set_chunk_size(self, chunk_size):
398    """Change the read chunk size.
399
400    This function can only be called if the FileReaderThread object was
401    created with an initial chunk_size > 0.
402    Args:
403      chunk_size: the new chunk size for this file.  Must be > 0.
404    """
405    # The chunk size can be changed asynchronously while a file is being read
406    # in a worker thread.  However, type of file can not be changed after the
407    # the FileReaderThread has been created.  These asserts verify that we are
408    # only changing the chunk size, and not the type of file.
409    assert not self._text_file
410    assert chunk_size > 0
411    self._chunk_size = chunk_size
412
413
414def get_default_categories(device_serial):
415  categories_output, return_code = util.run_adb_shell(LIST_CATEGORIES_ARGS,
416                                                    device_serial)
417
418  if return_code == 0 and categories_output:
419    categories = [c.split('-')[0].strip()
420                  for c in categories_output.splitlines()]
421    return [c for c in categories if c in DEFAULT_CATEGORIES]
422
423  return []
424
425
426def status_update(last_update_time):
427  current_time = time.time()
428  if (current_time - last_update_time) >= MIN_TIME_BETWEEN_STATUS_UPDATES:
429    # Gathering a trace may take a while.  Keep printing something so users
430    # don't think the script has hung.
431    sys.stdout.write('.')
432    sys.stdout.flush()
433    return current_time
434
435  return last_update_time
436
437
438def extract_thread_list(trace_data):
439  """Removes the thread list from the given trace data.
440
441  Args:
442    trace_data: The raw trace data (before decompression).
443  Returns:
444    A tuple containing the trace data and a map of thread ids to thread names.
445  """
446  threads = {}
447  parts = re.split('USER +PID +PPID +VSIZE +RSS +WCHAN +PC +NAME',
448                   trace_data, 1)
449  if len(parts) == 2:
450    trace_data = parts[0]
451    for line in parts[1].splitlines():
452      cols = line.split(None, 8)
453      if len(cols) == 9:
454        tid = int(cols[1])
455        name = cols[8]
456        threads[tid] = name
457
458  return (trace_data, threads)
459
460
461def strip_and_decompress_trace(trace_data):
462  """Fixes new-lines and decompresses trace data.
463
464  Args:
465    trace_data: The trace data returned by atrace.
466  Returns:
467    The decompressed trace data.
468  """
469  # Collapse CRLFs that are added by adb shell.
470  if trace_data.startswith('\r\n'):
471    trace_data = trace_data.replace('\r\n', '\n')
472  elif trace_data.startswith('\r\r\n'):
473    # On windows, adb adds an extra '\r' character for each line.
474    trace_data = trace_data.replace('\r\r\n', '\n')
475
476  # Skip the initial newline.
477  trace_data = trace_data[1:]
478
479  if not trace_data.startswith(TRACE_TEXT_HEADER):
480    # No header found, so assume the data is compressed.
481    trace_data = zlib.decompress(trace_data)
482
483  # Enforce Unix line-endings.
484  trace_data = trace_data.replace('\r', '')
485
486  # Skip any initial newlines.
487  while trace_data and trace_data[0] == '\n':
488    trace_data = trace_data[1:]
489
490  return trace_data
491
492
493def fix_thread_names(trace_data, thread_names):
494  """Replaces thread ids with their names.
495
496  Args:
497    trace_data: The atrace data.
498    thread_names: A mapping of thread ids to thread names.
499  Returns:
500    The updated trace data.
501  """
502  def repl(m):
503    tid = int(m.group(2))
504    if tid > 0:
505      name = thread_names.get(tid)
506      if name is None:
507        name = m.group(1)
508        if name == '<...>':
509          name = '<' + str(tid) + '>'
510        thread_names[tid] = name
511      return name + '-' + m.group(2)
512    else:
513      return m.group(0)
514  trace_data = re.sub(r'^\s*(\S+)-(\d+)', repl, trace_data,
515                      flags=re.MULTILINE)
516  return trace_data
517
518
519def fix_circular_traces(out):
520  """Fix inconsistentcies in traces due to circular buffering.
521
522  The circular buffers are kept per CPU, so it is not guaranteed that the
523  beginning of a slice is overwritten before the end. To work around this, we
524  throw away the prefix of the trace where not all CPUs have events yet.
525
526  Args:
527    out: The data to fix.
528  Returns:
529    The updated trace data.
530  """
531  # If any of the CPU's buffers have filled up and
532  # older events have been dropped, the kernel
533  # emits markers of the form '##### CPU 2 buffer started ####' on
534  # the line before the first event in the trace on that CPU.
535  #
536  # No such headers are emitted if there were no overflows or the trace
537  # was captured with non-circular buffers.
538  buffer_start_re = re.compile(r'^#+ CPU \d+ buffer started', re.MULTILINE)
539
540  start_of_full_trace = 0
541
542  while True:
543    result = buffer_start_re.search(out, start_of_full_trace + 1)
544    if result:
545      start_of_full_trace = result.start()
546    else:
547      break
548
549  if start_of_full_trace > 0:
550    # Need to keep the header intact to make the importer happy.
551    end_of_header = re.search(r'^[^#]', out, re.MULTILINE).start()
552    out = out[:end_of_header] + out[start_of_full_trace:]
553  return out
554