1 #!/usr/bin/python 2 # 3 # Copyright 2011 Google Inc. All Rights Reserved. 4 """Script to divide and merge profiles.""" 5 6 import copy 7 import optparse 8 import os 9 import pickle 10 import re 11 import sys 12 import tempfile 13 14 import build_chrome_browser 15 import lock_machine 16 import run_tests 17 from cros_utils import command_executer 18 from cros_utils import logger 19 20 21 class 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 95 def 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 138 if __name__ == '__main__': 139 retval = Main(sys.argv) 140 sys.exit(retval) 141