Support traces containing all parts/agents returned via DevTools streams

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

Index: telemetry/telemetry/internal/backends/chrome_inspector/tracing_backend.py
diff --git a/telemetry/telemetry/internal/backends/chrome_inspector/tracing_backend.py b/telemetry/telemetry/internal/backends/chrome_inspector/tracing_backend.py
index ad068e162559309790a88b6003204bced76e3ae8..5cb8607c58292c2421f7c1ed7e4268dc6c66ecf4 100644
--- a/telemetry/telemetry/internal/backends/chrome_inspector/tracing_backend.py
+++ b/telemetry/telemetry/internal/backends/chrome_inspector/tracing_backend.py
@@ -91,13 +91,13 @@ class TracingBackend(object):
     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._start_issued = False
     self._can_collect_data = False
     self._has_received_all_tracing_data = 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):
@@ -112,7 +112,6 @@ class TracingBackend(object):
       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 = []
 
     if not self.IsTracingSupported():
       raise TracingUnsupportedException(
@@ -162,8 +161,7 @@ class TracingBackend(object):
     tracing run is pushed.
     """
     if not self.is_tracing_running:
-      if not self._trace_events:
-        raise TracingHasNotRunException()
+      raise TracingHasNotRunException()
     else:
       if not self._start_issued:
         # Tracing is running but start was not issued so, startup tracing must
@@ -226,15 +224,14 @@ class TracingBackend(object):
   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:
@@ -244,45 +241,61 @@ class TracingBackend(object):
     """
     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())
-
-      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)
+    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
+
+        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)
     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:
-        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):