Home | History | Annotate | Download | only in fdo_scripts 
      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