Support traces containing all parts/agents returned via DevTools streams

telemetry/telemetry/internal/backends/chrome_inspector/tracing_backend.py

 # Copyright 2013 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import json
 import logging
 import socket
 import time
 import traceback
 
@@ skipping 73 matching lines @@
 
 class TracingBackend(object):
 
   _TRACING_DOMAIN = 'Tracing'
 
   def __init__(self, inspector_socket, is_tracing_running=False,
                support_modern_devtools_tracing_start_api=False):
     self._inspector_websocket = inspector_socket
     self._inspector_websocket.RegisterDomain(
         self._TRACING_DOMAIN, self._NotificationHandler)
-    self._trace_events = []
     self._is_tracing_running = is_tracing_running
     self._can_collect_data = False
     self._has_received_all_tracing_data = False
+    self._had_data_collected = False
     self._support_modern_devtools_tracing_start_api = (
         support_modern_devtools_tracing_start_api)
+    self._trace_data_builder = None
 
   @property
   def is_tracing_running(self):
     return self._is_tracing_running
 
   def StartTracing(self, chrome_trace_config, timeout=10):
     """When first called, starts tracing, and returns True.
 
     If called during tracing, tracing is unchanged, and it returns False.
     """
     if self.is_tracing_running:
       return False
     assert not self._can_collect_data, 'Data not collected from last trace.'
     # Reset collected tracing data from previous tracing calls.
-    self._trace_events = []
+    self._had_data_collected = False
 
     if not self.IsTracingSupported():
       raise TracingUnsupportedException(
           'Chrome tracing not supported for this app.')
 
     params = {'transferMode': 'ReturnAsStream'}
     if self._support_modern_devtools_tracing_start_api:
       params['traceConfig'] = (
           chrome_trace_config.GetChromeTraceConfigForDevTools())
     else:
@@ skipping 28 matching lines @@
     if 'error' in rc:
       raise ClockSyncResponseException(rc['error']['message'])
 
   def StopTracing(self):
     """Stops tracing and pushes results to the supplied TraceDataBuilder.
 
     If this is called after tracing has been stopped, trace data from the last
     tracing run is pushed.
     """
     if not self.is_tracing_running:
-      if not self._trace_events:
-        raise TracingHasNotRunException()
-    else:
-      req = {'method': 'Tracing.end'}
-      self._inspector_websocket.SendAndIgnoreResponse(req)
+      raise TracingHasNotRunException()
 
+    req = {'method': 'Tracing.end'}
+    self._inspector_websocket.SendAndIgnoreResponse(req)
     self._is_tracing_running = False
     self._can_collect_data = True
 
   def DumpMemory(self, timeout=30):
     """Dumps memory.
 
     Returns:
       GUID of the generated dump if successful, None otherwise.
 
     Raises:
@@ skipping 26 matching lines @@
       raise TracingUnexpectedResponseException(
           'Inspector returned unexpected response for '
           'Tracing.requestMemoryDump:\n' + json.dumps(response, indent=2))
 
     result = response['result']
     return result['dumpGuid'] if result['success'] else None
 
   def CollectTraceData(self, trace_data_builder, timeout=30):
     if not self._can_collect_data:
       raise Exception('Cannot collect before tracing is finished.')
-    self._CollectTracingData(timeout)
+    self._CollectTracingData(trace_data_builder, timeout)
     self._can_collect_data = False
-    trace_data_builder.AddEventsTo(
-      trace_data_module.CHROME_TRACE_PART, self._trace_events)
 
-  def _CollectTracingData(self, timeout):
+  def _CollectTracingData(self, trace_data_builder, timeout):
     """Collects tracing data. Assumes that Tracing.end has already been sent.
 
     Args:
+      trace_data_builder: An instance of TraceDataBuilder to put results into.
       timeout: The timeout in seconds.
 
     Raises:
       TracingTimeoutException: If more than |timeout| seconds has passed
       since the last time any data is received.
       TracingUnrecoverableException: If there is a websocket error.
     """
     self._has_received_all_tracing_data = False
     start_time = time.time()
-    while True:
-      try:
-        self._inspector_websocket.DispatchNotifications(timeout)
-        start_time = time.time()
-      except websocket.WebSocketTimeoutException:
-        pass
-      except (socket.error, websocket.WebSocketException):
-        raise TracingUnrecoverableException(
-            'Exception raised while collecting tracing data:\n' +
-                traceback.format_exc())
+    self._trace_data_builder = trace_data_builder
+    try:
+      while True:
+        try:
+          self._inspector_websocket.DispatchNotifications(timeout)
+          start_time = time.time()
+        except websocket.WebSocketTimeoutException:
+          pass
+        except (socket.error, websocket.WebSocketException):
+          raise TracingUnrecoverableException(
+              'Exception raised while collecting tracing data:\n' +
+                  traceback.format_exc())
 
-      if self._has_received_all_tracing_data:
-        break
+        if self._has_received_all_tracing_data:
+          break
 
-      elapsed_time = time.time() - start_time
-      if elapsed_time > timeout:
-        raise TracingTimeoutException(
-            'Only received partial trace data due to timeout after %s seconds. '
-            'If the trace data is big, you may want to increase the timeout '
-            'amount.' % elapsed_time)
+        elapsed_time = time.time() - start_time
+        if elapsed_time > timeout:
+          raise TracingTimeoutException(
+              'Only received partial trace data due to timeout after %s '
+              'seconds. If the trace data is big, you may want to increase '
+              'the timeout amount.' % elapsed_time)
+    finally:
+      self._trace_data_builder = None
 
   def _NotificationHandler(self, res):
     if 'Tracing.dataCollected' == res.get('method'):
       value = res.get('params', {}).get('value')
-      self._trace_events.extend(value)
+      self._trace_data_builder.AddEventsTo(
+        trace_data_module.CHROME_TRACE_PART, value)
+      self._had_data_collected = True
     elif 'Tracing.tracingComplete' == res.get('method'):
       stream_handle = res.get('params', {}).get('stream')
       if not stream_handle:
         self._has_received_all_tracing_data = True
         return
-
-      if self._trace_events:
+      if self._had_data_collected:

[Zhen Wang, 2016/07/29 14:45:42]

So this will never happen? If so, I would suggest to remove _had_data_collected
to reduce the complexity here.

         raise TracingUnexpectedResponseException(
             'Got both dataCollected events and a stream from server')
       reader = _DevToolsStreamReader(self._inspector_websocket, stream_handle)
       reader.Read(self._ReceivedAllTraceDataFromStream)
 
   def _ReceivedAllTraceDataFromStream(self, data):
-    self._trace_events = json.loads(data)
+    trace = json.loads(data)
+    if type(trace) == dict:
+      for part in trace_data_module.ALL_TRACE_PARTS:
+        field_name = part.raw_field_name
+        if field_name in trace:
+          self._trace_data_builder.AddEventsTo(part, trace[field_name])
+
+      if 'metadata' in trace:
+        self._trace_data_builder.SetMetadataFor(
+            trace_data_module.CHROME_TRACE_PART, trace['metadata'])
+
+    elif type(trace) == list:
+      self._trace_data_builder.AddEventsTo(
+        trace_data_module.CHROME_TRACE_PART, trace)
+    else:
+      raise TracingUnexpectedResponseException('Unexpected trace type')
     self._has_received_all_tracing_data = True
 
   def Close(self):
     self._inspector_websocket.UnregisterDomain(self._TRACING_DOMAIN)
     self._inspector_websocket = None
 
   @decorators.Cache
   def IsTracingSupported(self):
     req = {'method': 'Tracing.hasCompleted'}
     res = self._inspector_websocket.SyncRequest(req, timeout=10)
     return not res.get('response')