1# Copyright 2013 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 json
6import logging
7import socket
8import time
9import traceback
10
11from telemetry import decorators
12from telemetry.internal.backends.chrome_inspector import inspector_websocket
13from telemetry.internal.backends.chrome_inspector import websocket
14from telemetry.timeline import trace_data as trace_data_module
15
16
17class TracingUnsupportedException(Exception):
18  pass
19
20
21class TracingTimeoutException(Exception):
22  pass
23
24
25class TracingUnrecoverableException(Exception):
26  pass
27
28
29class TracingHasNotRunException(Exception):
30  pass
31
32
33class TracingUnexpectedResponseException(Exception):
34  pass
35
36
37class _DevToolsStreamReader(object):
38  def __init__(self, inspector_socket, stream_handle):
39    self._inspector_websocket = inspector_socket
40    self._handle = stream_handle
41    self._callback = None
42    self._data = None
43
44  def Read(self, callback):
45    # Do not allow the instance of this class to be reused, as
46    # we only read data sequentially at the moment, so a stream
47    # can only be read once.
48    assert not self._callback
49    self._data = []
50    self._callback = callback
51    self._ReadChunkFromStream()
52    # The below is not a typo -- queue one extra read ahead to avoid latency.
53    self._ReadChunkFromStream()
54
55  def _ReadChunkFromStream(self):
56    # Limit max block size to avoid fragmenting memory in sock.recv(),
57    # (see https://github.com/liris/websocket-client/issues/163 for details)
58    req = {'method': 'IO.read', 'params': {
59        'handle': self._handle, 'size': 32768}}
60    self._inspector_websocket.AsyncRequest(req, self._GotChunkFromStream)
61
62  def _GotChunkFromStream(self, response):
63    # Quietly discard responses from reads queued ahead after EOF.
64    if self._data is None:
65      return
66    if 'error' in response:
67      raise TracingUnrecoverableException(
68          'Reading trace failed: %s' % response['error']['message'])
69    result = response['result']
70    self._data.append(result['data'])
71    if not result.get('eof', False):
72      self._ReadChunkFromStream()
73      return
74    req = {'method': 'IO.close', 'params': {'handle': self._handle}}
75    self._inspector_websocket.SendAndIgnoreResponse(req)
76    trace_string = ''.join(self._data)
77    self._data = None
78    self._callback(trace_string)
79
80
81class TracingBackend(object):
82
83  _TRACING_DOMAIN = 'Tracing'
84
85  def __init__(self, inspector_socket, is_tracing_running=False):
86    self._inspector_websocket = inspector_socket
87    self._inspector_websocket.RegisterDomain(
88        self._TRACING_DOMAIN, self._NotificationHandler)
89    self._trace_events = []
90    self._is_tracing_running = is_tracing_running
91    self._has_received_all_tracing_data = False
92
93  @property
94  def is_tracing_running(self):
95    return self._is_tracing_running
96
97  def StartTracing(self, trace_options, custom_categories=None, timeout=10):
98    """When first called, starts tracing, and returns True.
99
100    If called during tracing, tracing is unchanged, and it returns False.
101    """
102    if self.is_tracing_running:
103      return False
104    # Reset collected tracing data from previous tracing calls.
105    self._trace_events = []
106
107    if not self.IsTracingSupported():
108      raise TracingUnsupportedException(
109          'Chrome tracing not supported for this app.')
110
111    req = {
112      'method': 'Tracing.start',
113      'params': {
114        'options': trace_options.GetTraceOptionsStringForChromeDevtool(),
115        'transferMode': 'ReturnAsStream'
116      }
117    }
118    if custom_categories:
119      req['params']['categories'] = custom_categories
120    logging.info('Start Tracing Request: %s', repr(req))
121    response = self._inspector_websocket.SyncRequest(req, timeout)
122
123    if 'error' in response:
124      raise TracingUnexpectedResponseException(
125          'Inspector returned unexpected response for '
126          'Tracing.start:\n' + json.dumps(response, indent=2))
127
128    self._is_tracing_running = True
129    return True
130
131  def StopTracing(self, trace_data_builder, timeout=30):
132    """Stops tracing and pushes results to the supplied TraceDataBuilder.
133
134    If this is called after tracing has been stopped, trace data from the last
135    tracing run is pushed.
136    """
137    if not self.is_tracing_running:
138      if not self._trace_events:
139        raise TracingHasNotRunException()
140    else:
141      req = {'method': 'Tracing.end'}
142      self._inspector_websocket.SendAndIgnoreResponse(req)
143      # After Tracing.end, chrome browser will send asynchronous notifications
144      # containing trace data. This is until Tracing.tracingComplete is sent,
145      # which means there is no trace buffers pending flush.
146      self._CollectTracingData(timeout)
147    self._is_tracing_running = False
148    trace_data_builder.AddEventsTo(
149      trace_data_module.CHROME_TRACE_PART, self._trace_events)
150
151  def DumpMemory(self, timeout=30):
152    """Dumps memory.
153
154    Returns:
155      GUID of the generated dump if successful, None otherwise.
156
157    Raises:
158      TracingTimeoutException: If more than |timeout| seconds has passed
159      since the last time any data is received.
160      TracingUnrecoverableException: If there is a websocket error.
161      TracingUnexpectedResponseException: If the response contains an error
162      or does not contain the expected result.
163    """
164    request = {
165      'method': 'Tracing.requestMemoryDump'
166    }
167    try:
168      response = self._inspector_websocket.SyncRequest(request, timeout)
169    except websocket.WebSocketTimeoutException:
170      raise TracingTimeoutException(
171          'Exception raised while sending a Tracing.requestMemoryDump '
172          'request:\n' + traceback.format_exc())
173    except (socket.error, websocket.WebSocketException,
174            inspector_websocket.WebSocketDisconnected):
175      raise TracingUnrecoverableException(
176          'Exception raised while sending a Tracing.requestMemoryDump '
177          'request:\n' + traceback.format_exc())
178
179
180    if ('error' in response or
181        'result' not in response or
182        'success' not in response['result'] or
183        'dumpGuid' not in response['result']):
184      raise TracingUnexpectedResponseException(
185          'Inspector returned unexpected response for '
186          'Tracing.requestMemoryDump:\n' + json.dumps(response, indent=2))
187
188    result = response['result']
189    return result['dumpGuid'] if result['success'] else None
190
191  def _CollectTracingData(self, timeout):
192    """Collects tracing data. Assumes that Tracing.end has already been sent.
193
194    Args:
195      timeout: The timeout in seconds.
196
197    Raises:
198      TracingTimeoutException: If more than |timeout| seconds has passed
199      since the last time any data is received.
200      TracingUnrecoverableException: If there is a websocket error.
201    """
202    self._has_received_all_tracing_data = False
203    start_time = time.time()
204    while True:
205      try:
206        self._inspector_websocket.DispatchNotifications(timeout)
207        start_time = time.time()
208      except websocket.WebSocketTimeoutException:
209        pass
210      except (socket.error, websocket.WebSocketException):
211        raise TracingUnrecoverableException(
212            'Exception raised while collecting tracing data:\n' +
213                traceback.format_exc())
214
215      if self._has_received_all_tracing_data:
216        break
217
218      elapsed_time = time.time() - start_time
219      if elapsed_time > timeout:
220        raise TracingTimeoutException(
221            'Only received partial trace data due to timeout after %s seconds. '
222            'If the trace data is big, you may want to increase the timeout '
223            'amount.' % elapsed_time)
224
225  def _NotificationHandler(self, res):
226    if 'Tracing.dataCollected' == res.get('method'):
227      value = res.get('params', {}).get('value')
228      self._trace_events.extend(value)
229    elif 'Tracing.tracingComplete' == res.get('method'):
230      stream_handle = res.get('params', {}).get('stream')
231      if not stream_handle:
232        self._has_received_all_tracing_data = True
233        return
234
235      if self._trace_events:
236        raise TracingUnexpectedResponseException(
237            'Got both dataCollected events and a stream from server')
238      reader = _DevToolsStreamReader(self._inspector_websocket, stream_handle)
239      reader.Read(self._ReceivedAllTraceDataFromStream)
240
241  def _ReceivedAllTraceDataFromStream(self, data):
242    self._trace_events = json.loads(data)
243    self._has_received_all_tracing_data = True
244
245  def Close(self):
246    self._inspector_websocket.UnregisterDomain(self._TRACING_DOMAIN)
247    self._inspector_websocket = None
248
249  @decorators.Cache
250  def IsTracingSupported(self):
251    req = {'method': 'Tracing.hasCompleted'}
252    res = self._inspector_websocket.SyncRequest(req)
253    return not res.get('response')
254