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 json 6import os 7import tempfile 8import zipfile 9 10class NonSerializableTraceData(Exception): 11 """Raised when raw trace data cannot be serialized to TraceData.""" 12 pass 13 14 15def _ValidateRawData(raw): 16 try: 17 json.dumps(raw) 18 except TypeError as e: 19 raise NonSerializableTraceData('TraceData is not serilizable: %s' % e) 20 except ValueError as e: 21 raise NonSerializableTraceData('TraceData is not serilizable: %s' % e) 22 23 24class TraceDataPart(object): 25 """TraceData can have a variety of events. 26 27 These are called "parts" and are accessed by the following fixed field names. 28 """ 29 def __init__(self, raw_field_name): 30 self._raw_field_name = raw_field_name 31 32 def __repr__(self): 33 return 'TraceDataPart("%s")' % self._raw_field_name 34 35 @property 36 def raw_field_name(self): 37 return self._raw_field_name 38 39 40ATRACE_PART = TraceDataPart('systemTraceEvents') 41BATTOR_TRACE_PART = TraceDataPart('powerTraceAsString') 42CHROME_TRACE_PART = TraceDataPart('traceEvents') 43CPU_TRACE_DATA = TraceDataPart('cpuSnapshots') 44INSPECTOR_TRACE_PART = TraceDataPart('inspectorTimelineEvents') 45SURFACE_FLINGER_PART = TraceDataPart('surfaceFlinger') 46TAB_ID_PART = TraceDataPart('tabIds') 47TELEMETRY_PART = TraceDataPart('telemetry') 48 49ALL_TRACE_PARTS = {ATRACE_PART, 50 BATTOR_TRACE_PART, 51 CHROME_TRACE_PART, 52 CPU_TRACE_DATA, 53 INSPECTOR_TRACE_PART, 54 SURFACE_FLINGER_PART, 55 TAB_ID_PART, 56 TELEMETRY_PART} 57 58 59def _HasTraceFor(part, raw): 60 assert isinstance(part, TraceDataPart) 61 if part.raw_field_name not in raw: 62 return False 63 return len(raw[part.raw_field_name]) > 0 64 65 66class TraceData(object): 67 """Validates, parses, and serializes raw data. 68 69 NOTE: raw data must only include primitive objects! 70 By design, TraceData must contain only data that is BOTH json-serializable 71 to a file, AND restorable once again from that file into TraceData without 72 assistance from other classes. 73 74 Raw data can be one of three standard trace_event formats: 75 1. Trace container format: a json-parseable dict. 76 2. A json-parseable array: assumed to be chrome trace data. 77 3. A json-parseable array missing the final ']': assumed to be chrome trace 78 data. 79 """ 80 def __init__(self, raw_data=None): 81 """Creates TraceData from the given data.""" 82 self._raw_data = {} 83 self._events_are_safely_mutable = False 84 if not raw_data: 85 return 86 _ValidateRawData(raw_data) 87 88 if isinstance(raw_data, basestring): 89 if raw_data.startswith('[') and not raw_data.endswith(']'): 90 if raw_data.endswith(','): 91 raw_data = raw_data[:-1] 92 raw_data += ']' 93 json_data = json.loads(raw_data) 94 # The parsed data isn't shared with anyone else, so we mark this value 95 # as safely mutable. 96 self._events_are_safely_mutable = True 97 else: 98 json_data = raw_data 99 100 if isinstance(json_data, dict): 101 self._raw_data = json_data 102 elif isinstance(json_data, list): 103 if len(json_data) == 0: 104 self._raw_data = {} 105 self._raw_data = {CHROME_TRACE_PART.raw_field_name: { 106 'traceEvents': json_data 107 }} 108 else: 109 raise Exception('Unrecognized data format.') 110 111 def _SetFromBuilder(self, d): 112 self._raw_data = d 113 self._events_are_safely_mutable = True 114 115 @property 116 def events_are_safely_mutable(self): 117 """Returns true if the events in this value are completely sealed. 118 119 Some importers want to take complex fields out of the TraceData and add 120 them to the model, changing them subtly as they do so. If the TraceData 121 was constructed with data that is shared with something outside the trace 122 data, for instance a test harness, then this mutation is unexpected. But, 123 if the values are sealed, then mutating the events is a lot faster. 124 125 We know if events are sealed if the value came from a string, or if the 126 value came from a TraceDataBuilder. 127 """ 128 return self._events_are_safely_mutable 129 130 @property 131 def active_parts(self): 132 return {p for p in ALL_TRACE_PARTS if p.raw_field_name in self._raw_data} 133 134 @property 135 def metadata_records(self): 136 part_field_names = {p.raw_field_name for p in ALL_TRACE_PARTS} 137 for k, v in self._raw_data.iteritems(): 138 if k in part_field_names: 139 continue 140 yield { 141 'name': k, 142 'value': v 143 } 144 145 def HasTraceFor(self, part): 146 return _HasTraceFor(part, self._raw_data) 147 148 def GetTraceFor(self, part): 149 if not self.HasTraceFor(part): 150 return [] 151 assert isinstance(part, TraceDataPart) 152 return self._raw_data[part.raw_field_name] 153 154 def Serialize(self, f, gzip_result=False): 155 """Serializes the trace result to a file-like object. 156 157 Write in trace container format if gzip_result=False. 158 Writes to a .zip file if gzip_result=True. 159 """ 160 if gzip_result: 161 zip_file = zipfile.ZipFile(f, mode='w') 162 try: 163 for part in self.active_parts: 164 tmp_file_name = None 165 with tempfile.NamedTemporaryFile(delete=False) as tmp_file: 166 tmp_file_name = tmp_file.name 167 tmp_file.write(str(self._raw_data[part.raw_field_name])) 168 zip_file.write(tmp_file_name, arcname=part.raw_field_name) 169 os.remove(tmp_file_name) 170 finally: 171 zip_file.close() 172 else: 173 json.dump(self._raw_data, f) 174 175 176class TraceDataBuilder(object): 177 """TraceDataBuilder helps build up a trace from multiple trace agents. 178 179 TraceData is supposed to be immutable, but it is useful during recording to 180 have a mutable version. That is TraceDataBuilder. 181 """ 182 def __init__(self): 183 self._raw_data = {} 184 185 def AsData(self): 186 if self._raw_data == None: 187 raise Exception('Can only AsData once') 188 189 data = TraceData() 190 data._SetFromBuilder(self._raw_data) 191 self._raw_data = None 192 return data 193 194 def AddEventsTo(self, part, events): 195 """Note: this won't work when called from multiple browsers. 196 197 Each browser's trace_event_impl zeros its timestamps when it writes them 198 out and doesn't write a timebase that can be used to re-sync them. 199 """ 200 assert isinstance(part, TraceDataPart) 201 assert isinstance(events, list) 202 if self._raw_data == None: 203 raise Exception('Already called AsData() on this builder.') 204 if part == CHROME_TRACE_PART: 205 target_events = self._raw_data.setdefault( 206 part.raw_field_name, {}).setdefault('traceEvents', []) 207 else: 208 target_events = self._raw_data.setdefault(part.raw_field_name, []) 209 target_events.extend(events) 210 211 def SetTraceFor(self, part, trace): 212 assert isinstance(part, TraceDataPart) 213 assert (isinstance(trace, basestring) or 214 isinstance(trace, dict) or 215 isinstance(trace, list)) 216 217 if self._raw_data == None: 218 raise Exception('Already called AsData() on this builder.') 219 220 if part.raw_field_name in self._raw_data: 221 raise Exception('Trace part %s is already set.' % part.raw_field_name) 222 223 self._raw_data[part.raw_field_name] = trace 224 225 def SetMetadataFor(self, part, metadata): 226 if part != CHROME_TRACE_PART: 227 raise Exception('Metadata are only supported for %s' 228 % CHROME_TRACE_PART.raw_field_name) 229 self._raw_data.setdefault(part.raw_field_name, {})['metadata'] = metadata 230 231 def HasTraceFor(self, part): 232 return _HasTraceFor(part, self._raw_data) 233