146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)# Copyright 2014 The Chromium Authors. All rights reserved. 246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)# Use of this source code is governed by a BSD-style license that can be 346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)# found in the LICENSE file. 446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)import gzip 646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)import json 746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)import os 846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)import shutil 946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)import sys 1046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)import zipfile 1146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 125f1c94371a64b3196d4be9466099bb892df9b88eTorne (Richard Coles)from profile_chrome import util 1346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 1446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)from pylib import constants 1546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 1646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)sys.path.append(os.path.join(constants.DIR_SOURCE_ROOT, 1746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 'third_party', 1846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 'trace-viewer')) 1946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)# pylint: disable=F0401 2046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)from trace_viewer.build import trace2html 2146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 2246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 2346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)def _PackageTracesAsHtml(trace_files, html_file): 2446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) with open(html_file, 'w') as f: 2546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) trace2html.WriteHTMLForTracesToFile(trace_files, f) 2646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) for trace_file in trace_files: 2746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) os.unlink(trace_file) 2846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 2946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 3046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)def _CompressFile(host_file, output): 3146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) with gzip.open(output, 'wb') as out, \ 3246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) open(host_file, 'rb') as input_file: 3346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) out.write(input_file.read()) 3446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) os.unlink(host_file) 3546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 3646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 3746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)def _ArchiveFiles(host_files, output): 3846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) with zipfile.ZipFile(output, 'w', zipfile.ZIP_DEFLATED) as z: 3946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) for host_file in host_files: 4046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) z.write(host_file) 4146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) os.unlink(host_file) 4246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 4346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 4446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)def _MergeTracesIfNeeded(trace_files): 4546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) if len(trace_files) <= 1: 4646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) return trace_files 4746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) merge_candidates = [] 4846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) for trace_file in trace_files: 4946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) with open(trace_file) as f: 5046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) # Try to detect a JSON file cheaply since that's all we can merge. 5146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) if f.read(1) != '{': 5246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) continue 5346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) f.seek(0) 5446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) try: 5546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) json_data = json.load(f) 5646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) except ValueError: 5746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) continue 5846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) merge_candidates.append((trace_file, json_data)) 5946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) if len(merge_candidates) <= 1: 6046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) return trace_files 6146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 6246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) other_files = [f for f in trace_files 6346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) if not f in [c[0] for c in merge_candidates]] 6446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) merged_file, merged_data = merge_candidates[0] 6546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) for trace_file, json_data in merge_candidates[1:]: 6646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) for key, value in json_data.items(): 6746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) if not merged_data.get(key) or json_data[key]: 6846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) merged_data[key] = value 6946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) os.unlink(trace_file) 7046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 7146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) with open(merged_file, 'w') as f: 7246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) json.dump(merged_data, f) 7346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) return [merged_file] + other_files 7446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 7546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 7646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles)def PackageTraces(trace_files, output=None, compress=False, write_json=False): 7746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) trace_files = _MergeTracesIfNeeded(trace_files) 7846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) if not write_json: 7946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) html_file = os.path.splitext(trace_files[0])[0] + '.html' 8046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) _PackageTracesAsHtml(trace_files, html_file) 8146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) trace_files = [html_file] 8246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) 8346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) if compress and len(trace_files) == 1: 8446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) result = output or trace_files[0] + '.gz' 8546d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) _CompressFile(trace_files[0], result) 8646d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) elif len(trace_files) > 1: 8746d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) result = output or 'chrome-combined-trace-%s.zip' % util.GetTraceTimestamp() 8846d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) _ArchiveFiles(trace_files, result) 8946d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) elif output: 9046d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) result = output 9146d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) shutil.move(trace_files[0], result) 9246d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) else: 9346d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) result = trace_files[0] 9446d4c2bc3267f3f028f39e7e311b0f89aba2e4fdTorne (Richard Coles) return result 95