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