divide_and_merge_profiles.py revision f81680c018729fd4499e1e200d04b48c4b90127c 
1#!/usr/bin/python2.6
2#
3# Copyright 2011 Google Inc. All Rights Reserved.
4
5"""Script to divide and merge profiles."""
6
7import copy
8import optparse
9import os
10import pickle
11import re
12import sys
13import tempfile
14
15import build_chrome_browser
16import lock_machine
17import run_tests
18from utils import command_executer
19from utils import logger
20
21
22class ProfileMerger:
23  def __init__(self, inputs, output, chunk_size, merge_program, multipliers):
24    self._inputs = inputs
25    self._output = output
26    self._chunk_size = chunk_size
27    self._merge_program = merge_program
28    self._multipliers = multipliers
29    self._ce = command_executer.GetCommandExecuter()
30    self._l = logger.GetLogger()
31
32  def _GetFilesSetForInputDir(self, input_dir):
33    output_file = tempfile.mktemp()
34    command = "find %s -name '*.gcda' -o -name '*.imports' > %s" % (input_dir, output_file)
35    self._ce.RunCommand(command)
36    files = open(output_file, "r").read()
37    files_set = set([])
38    for f in files.splitlines():
39      stripped_file = f.replace(input_dir, "", 1)
40      stripped_file = stripped_file.lstrip("/")
41      files_set.add(stripped_file)
42    return files_set
43
44  def _PopulateFilesSet(self):
45    self._files_set = set([])
46    for i in self._inputs:
47      current_files_set = self._GetFilesSetForInputDir(i)
48      self._files_set.update(current_files_set)
49
50  def _GetSubset(self):
51    ret = []
52    for i in range(self._chunk_size):
53      if not self._files_set:
54        break
55      ret.append(self._files_set.pop())
56    return ret
57
58  def _CopyFilesTree(self, input_dir, files, output_dir):
59    for f in files:
60      src_file = os.path.join(input_dir, f)
61      dst_file = os.path.join(output_dir, f)
62      if not os.path.isdir(os.path.dirname(dst_file)):
63        command = "mkdir -p %s" % os.path.dirname(dst_file)
64        self._ce.RunCommand(command)
65      command = "cp %s %s" % (src_file, dst_file)
66      self._ce.RunCommand(command)
67
68  def _DoChunkMerge(self, current_files):
69    temp_dirs = []
70    for i in self._inputs:
71      temp_dir = tempfile.mkdtemp()
72      temp_dirs.append(temp_dir)
73      self._CopyFilesTree(i, current_files, temp_dir)
74    # Now do the merge.
75    command = ("%s --inputs=%s --output=%s" %
76               (self._merge_program,
77                ",".join(temp_dirs),
78                self._output))
79    if self._multipliers:
80      command = ("%s --multipliers=%s" %
81                 (command, self._multipliers))
82    ret = self._ce.RunCommand(command)
83    assert ret == 0, "%s command failed!" % command
84    for temp_dir in temp_dirs:
85      command = "rm -rf %s" % temp_dir
86      self._ce.RunCommand(command)
87
88  def DoMerge(self):
89    self._PopulateFilesSet()
90    while True:
91      current_files = self._GetSubset()
92      if not current_files:
93        break
94      self._DoChunkMerge(current_files)
95
96
97def Main(argv):
98  """The main function."""
99  # Common initializations
100###  command_executer.InitCommandExecuter(True)
101  command_executer.InitCommandExecuter()
102  l = logger.GetLogger()
103  ce = command_executer.GetCommandExecuter()
104  parser = optparse.OptionParser()
105  parser.add_option("--inputs",
106                    dest="inputs",
107                    help="Comma-separated input profile directories to merge.")
108  parser.add_option("--output",
109                    dest="output",
110                    help="Output profile directory.")
111  parser.add_option("--chunk_size",
112                    dest="chunk_size",
113                    default="50",
114                    help="Chunk size to divide up the profiles into.")
115  parser.add_option("--merge_program",
116                    dest="merge_program",
117                    default="/home/xur/bin/profile_merge_v15.par",
118                    help="Merge program to use to do the actual merge.")
119  parser.add_option("--multipliers",
120                    dest="multipliers",
121                    help="multipliers to use when merging. (optional)")
122
123  options, _ = parser.parse_args(argv)
124
125  if not all([options.inputs,
126              options.output,]):
127    l.LogError("Must supply --inputs and --output")
128    return 1
129
130  try:
131    pm = ProfileMerger(options.inputs.split(","), options.output,
132                       int(options.chunk_size), options.merge_program,
133                       options.multipliers)
134    pm.DoMerge()
135    retval = 0
136  except:
137    retval = 1
138  finally:
139    print "My work is done..."
140  return retval
141
142
143if __name__ == "__main__":
144  retval = Main(sys.argv)
145  sys.exit(retval)
146