binary_search_state.py revision 9099a788cd7124024559c064e425ed9caef6e0ac
1#!/usr/bin/python2 2"""The binary search wrapper.""" 3 4from __future__ import print_function 5 6import argparse 7import contextlib 8import math 9import os 10import pickle 11import sys 12import tempfile 13import time 14 15# Adds cros_utils to PYTHONPATH 16import common 17 18# Now we do import from cros_utils 19from cros_utils import command_executer 20from cros_utils import logger 21 22import binary_search_perforce 23 24GOOD_SET_VAR = 'BISECT_GOOD_SET' 25BAD_SET_VAR = 'BISECT_BAD_SET' 26 27STATE_FILE = '%s.state' % sys.argv[0] 28HIDDEN_STATE_FILE = os.path.join( 29 os.path.dirname(STATE_FILE), '.%s' % os.path.basename(STATE_FILE)) 30 31class Error(Exception): 32 """The general binary search tool error class.""" 33 pass 34 35 36@contextlib.contextmanager 37def SetFile(env_var, items): 38 """Generate set files that can be used by switch/test scripts. 39 40 Generate temporary set file (good/bad) holding contents of good/bad items for 41 the current binary search iteration. Store the name of each file as an 42 environment variable so all child processes can access it. 43 44 This function is a contextmanager, meaning it's meant to be used with the 45 "with" statement in Python. This is so cleanup and setup happens automatically 46 and cleanly. Execution of the outer "with" statement happens at the "yield" 47 statement. 48 49 Args: 50 env_var: What environment variable to store the file name in. 51 items: What items are in this set. 52 """ 53 with tempfile.NamedTemporaryFile() as f: 54 os.environ[env_var] = f.name 55 f.write('\n'.join(items)) 56 f.flush() 57 yield 58 59 60class BinarySearchState(object): 61 """The binary search state class.""" 62 63 def __init__(self, get_initial_items, switch_to_good, switch_to_bad, 64 install_script, test_script, incremental, prune, iterations, 65 prune_iterations, verify_level, file_args, verbose): 66 """BinarySearchState constructor, see Run for full args documentation.""" 67 self.get_initial_items = get_initial_items 68 self.switch_to_good = switch_to_good 69 self.switch_to_bad = switch_to_bad 70 self.install_script = install_script 71 self.test_script = test_script 72 self.incremental = incremental 73 self.prune = prune 74 self.iterations = iterations 75 self.prune_iterations = prune_iterations 76 self.verify_level = verify_level 77 self.file_args = file_args 78 self.verbose = verbose 79 80 self.l = logger.GetLogger() 81 self.ce = command_executer.GetCommandExecuter() 82 83 self.resumed = False 84 self.prune_cycles = 0 85 self.search_cycles = 0 86 self.binary_search = None 87 self.all_items = None 88 self.PopulateItemsUsingCommand(self.get_initial_items) 89 self.currently_good_items = set([]) 90 self.currently_bad_items = set([]) 91 self.found_items = set([]) 92 self.known_good = set([]) 93 94 self.start_time = time.time() 95 96 def SwitchToGood(self, item_list): 97 """Switch given items to "good" set.""" 98 if self.incremental: 99 self.l.LogOutput('Incremental set. Wanted to switch %s to good' % 100 str(item_list), print_to_console=self.verbose) 101 incremental_items = [ 102 item for item in item_list if item not in self.currently_good_items 103 ] 104 item_list = incremental_items 105 self.l.LogOutput('Incremental set. Actually switching %s to good' % 106 str(item_list), print_to_console=self.verbose) 107 108 if not item_list: 109 return 110 111 self.l.LogOutput('Switching %s to good' % str(item_list), 112 print_to_console=self.verbose) 113 self.RunSwitchScript(self.switch_to_good, item_list) 114 self.currently_good_items = self.currently_good_items.union(set(item_list)) 115 self.currently_bad_items.difference_update(set(item_list)) 116 117 def SwitchToBad(self, item_list): 118 """Switch given items to "bad" set.""" 119 if self.incremental: 120 self.l.LogOutput('Incremental set. Wanted to switch %s to bad' % 121 str(item_list), print_to_console=self.verbose) 122 incremental_items = [ 123 item for item in item_list if item not in self.currently_bad_items 124 ] 125 item_list = incremental_items 126 self.l.LogOutput('Incremental set. Actually switching %s to bad' % 127 str(item_list), print_to_console=self.verbose) 128 129 if not item_list: 130 return 131 132 self.l.LogOutput('Switching %s to bad' % str(item_list), 133 print_to_console=self.verbose) 134 self.RunSwitchScript(self.switch_to_bad, item_list) 135 self.currently_bad_items = self.currently_bad_items.union(set(item_list)) 136 self.currently_good_items.difference_update(set(item_list)) 137 138 def RunSwitchScript(self, switch_script, item_list): 139 """Pass given items to switch script. 140 141 Args: 142 switch_script: path to switch script 143 item_list: list of all items to be switched 144 """ 145 if self.file_args: 146 with tempfile.NamedTemporaryFile() as f: 147 f.write('\n'.join(item_list)) 148 f.flush() 149 command = '%s %s' % (switch_script, f.name) 150 ret, _, _ = self.ce.RunCommandWExceptionCleanup( 151 command, print_to_console=self.verbose) 152 else: 153 command = '%s %s' % (switch_script, ' '.join(item_list)) 154 ret, _, _ = self.ce.RunCommandWExceptionCleanup( 155 command, print_to_console=self.verbose) 156 assert ret == 0, 'Switch script %s returned %d' % (switch_script, ret) 157 158 def TestScript(self): 159 """Run test script and return exit code from script.""" 160 command = self.test_script 161 ret, _, _ = self.ce.RunCommandWExceptionCleanup(command) 162 return ret 163 164 def InstallScript(self): 165 """Run install script and return exit code from script.""" 166 if not self.install_script: 167 return 0 168 169 command = self.install_script 170 ret, _, _ = self.ce.RunCommandWExceptionCleanup(command) 171 return ret 172 173 def DoVerify(self): 174 """Verify correctness of test environment. 175 176 Verify that a "good" set of items produces a "good" result and that a "bad" 177 set of items produces a "bad" result. To be run directly before running 178 DoSearch. If verify_level is 0 this step is skipped. 179 """ 180 if not self.verify_level: 181 return 182 183 self.l.LogOutput('VERIFICATION') 184 self.l.LogOutput('Beginning %d tests to verify good/bad sets\n' % 185 self.verify_level) 186 for _ in range(int(self.verify_level)): 187 with SetFile(GOOD_SET_VAR, self.all_items), SetFile(BAD_SET_VAR, []): 188 self.l.LogOutput('Resetting all items to good to verify.') 189 self.SwitchToGood(self.all_items) 190 status = self.InstallScript() 191 assert status == 0, 'When reset_to_good, install should succeed.' 192 status = self.TestScript() 193 assert status == 0, 'When reset_to_good, status should be 0.' 194 195 with SetFile(GOOD_SET_VAR, []), SetFile(BAD_SET_VAR, self.all_items): 196 self.l.LogOutput('Resetting all items to bad to verify.') 197 self.SwitchToBad(self.all_items) 198 status = self.InstallScript() 199 assert status == 0, 'When reset_to_bad, install should succeed.' 200 status = self.TestScript() 201 assert status == 1, 'When reset_to_bad, status should be 1.' 202 203 def DoSearch(self): 204 """Perform full search for bad items. 205 206 Perform full search until prune_iterations number of bad items are found. 207 """ 208 while (True and 209 len(self.all_items) > 1 and 210 self.prune_cycles < self.prune_iterations): 211 terminated = self.DoBinarySearch() 212 self.prune_cycles += 1 213 if not terminated: 214 break 215 # Prune is set. 216 prune_index = self.binary_search.current 217 218 # If already seen item we have no new bad items to find, finish up 219 if self.all_items[prune_index] in self.found_items: 220 self.l.LogOutput(('Found item already found before: %s. ' 221 'Done searching.' % self.all_items[prune_index])) 222 self.l.LogOutput('Bad items are: %s' % 223 ' '.join(self.found_items)) 224 break 225 226 new_all_items = list(self.all_items) 227 # Move prune item to the end of the list. 228 new_all_items.append(new_all_items.pop(prune_index)) 229 self.found_items.add(new_all_items[-1]) 230 231 # Everything below newly found bad item is now known to be a good item. 232 # Take these good items out of the equation to save time on the next 233 # search. We save these known good items so they are still sent to the 234 # switch_to_good script. 235 if prune_index: 236 self.known_good.update(new_all_items[:prune_index]) 237 new_all_items = new_all_items[prune_index:] 238 239 self.l.LogOutput('Old list: %s. New list: %s' % (str(self.all_items), 240 str(new_all_items)), 241 print_to_console=self.verbose) 242 243 if not self.prune: 244 self.l.LogOutput('Not continuning further, --prune is not set') 245 break 246 # FIXME: Do we need to Convert the currently good items to bad 247 self.PopulateItemsUsingList(new_all_items) 248 249 def DoBinarySearch(self): 250 """Perform single iteration of binary search.""" 251 # If in resume mode don't reset search_cycles 252 if not self.resumed: 253 self.search_cycles = 0 254 else: 255 self.resumed = False 256 257 terminated = False 258 while self.search_cycles < self.iterations and not terminated: 259 self.SaveState() 260 self.OutputProgress() 261 262 self.search_cycles += 1 263 [bad_items, good_items] = self.GetNextItems() 264 265 with SetFile(GOOD_SET_VAR, good_items), SetFile(BAD_SET_VAR, bad_items): 266 # TODO: bad_items should come first. 267 self.SwitchToGood(good_items) 268 self.SwitchToBad(bad_items) 269 status = self.InstallScript() 270 if status == 0: 271 status = self.TestScript() 272 else: 273 # Install script failed, treat as skipped item 274 status = 2 275 terminated = self.binary_search.SetStatus(status) 276 277 if terminated: 278 self.l.LogOutput('Terminated!') 279 if not terminated: 280 self.l.LogOutput('Ran out of iterations searching...') 281 self.l.LogOutput(str(self), print_to_console=self.verbose) 282 return terminated 283 284 def PopulateItemsUsingCommand(self, command): 285 """Update all_items and binary search logic from executable. 286 287 This method is mainly required for enumerating the initial list of items 288 from the get_initial_items script. 289 290 Args: 291 command: path to executable that will enumerate items. 292 """ 293 ce = command_executer.GetCommandExecuter() 294 _, out, _ = ce.RunCommandWExceptionCleanup(command, 295 return_output=True, 296 print_to_console=self.verbose) 297 all_items = out.split() 298 self.PopulateItemsUsingList(all_items) 299 300 def PopulateItemsUsingList(self, all_items): 301 """Update all_items and binary searching logic from list. 302 303 Args: 304 all_items: new list of all_items 305 """ 306 self.all_items = all_items 307 self.binary_search = binary_search_perforce.BinarySearcher( 308 logger_to_set=self.l) 309 self.binary_search.SetSortedList(self.all_items) 310 311 def SaveState(self): 312 """Save state to STATE_FILE. 313 314 SaveState will create a new unique, hidden state file to hold data from 315 object. Then atomically overwrite the STATE_FILE symlink to point to the 316 new data. 317 318 Raises: 319 Error if STATE_FILE already exists but is not a symlink. 320 """ 321 ce, l = self.ce, self.l 322 self.ce, self.l, self.binary_search.logger = None, None, None 323 old_state = None 324 325 _, path = tempfile.mkstemp(prefix=HIDDEN_STATE_FILE, dir='.') 326 with open(path, 'wb') as f: 327 pickle.dump(self, f) 328 329 if os.path.exists(STATE_FILE): 330 if os.path.islink(STATE_FILE): 331 old_state = os.readlink(STATE_FILE) 332 else: 333 raise Error(('%s already exists and is not a symlink!\n' 334 'State file saved to %s' % (STATE_FILE, path))) 335 336 # Create new link and atomically overwrite old link 337 temp_link = '%s.link' % HIDDEN_STATE_FILE 338 os.symlink(path, temp_link) 339 os.rename(temp_link, STATE_FILE) 340 341 if old_state: 342 os.remove(old_state) 343 344 self.ce, self.l, self.binary_search.logger = ce, l, l 345 346 @classmethod 347 def LoadState(cls): 348 """Create BinarySearchState object from STATE_FILE.""" 349 if not os.path.isfile(STATE_FILE): 350 return None 351 try: 352 bss = pickle.load(file(STATE_FILE)) 353 bss.l = logger.GetLogger() 354 bss.ce = command_executer.GetCommandExecuter() 355 bss.binary_search.logger = bss.l 356 bss.start_time = time.time() 357 bss.resumed = True 358 binary_search_perforce.verbose = bss.verbose 359 return bss 360 except StandardError: 361 return None 362 363 def RemoveState(self): 364 """Remove STATE_FILE and its symlinked data from file system.""" 365 if os.path.exists(STATE_FILE): 366 if os.path.islink(STATE_FILE): 367 real_file = os.readlink(STATE_FILE) 368 os.remove(real_file) 369 os.remove(STATE_FILE) 370 371 def GetNextItems(self): 372 """Get next items for binary search based on result of the last test run.""" 373 border_item = self.binary_search.GetNext() 374 index = self.all_items.index(border_item) 375 376 next_bad_items = self.all_items[:index + 1] 377 next_good_items = self.all_items[index + 1:] + list(self.known_good) 378 379 return [next_bad_items, next_good_items] 380 381 def ElapsedTimeString(self): 382 """Return h m s format of elapsed time since execution has started.""" 383 diff = int(time.time() - self.start_time) 384 seconds = diff % 60 385 minutes = (diff / 60) % 60 386 hours = diff / (60 * 60) 387 388 seconds = str(seconds).rjust(2) 389 minutes = str(minutes).rjust(2) 390 hours = str(hours).rjust(2) 391 392 return '%sh %sm %ss' % (hours, minutes, seconds) 393 394 def OutputProgress(self): 395 """Output current progress of binary search to console and logs.""" 396 out = ('\n***** PROGRESS (elapsed time: %s) *****\n' 397 'Search %d of estimated %d.\n' 398 'Prune %d of max %d.\n' 399 'Current bad items found:\n' 400 '%s\n' 401 '************************************************') 402 out = out % (self.ElapsedTimeString(), 403 self.search_cycles + 1, 404 math.ceil(math.log(len(self.all_items), 2)), 405 self.prune_cycles + 1, 406 self.prune_iterations, 407 str(self.found_items)) 408 409 self.l.LogOutput(out) 410 411 def __str__(self): 412 ret = '' 413 ret += 'all: %s\n' % str(self.all_items) 414 ret += 'currently_good: %s\n' % str(self.currently_good_items) 415 ret += 'currently_bad: %s\n' % str(self.currently_bad_items) 416 ret += str(self.binary_search) 417 return ret 418 419 420class MockBinarySearchState(BinarySearchState): 421 """Mock class for BinarySearchState.""" 422 423 def __init__(self, **kwargs): 424 # Initialize all arguments to None 425 default_kwargs = { 426 'get_initial_items': 'echo "1"', 427 'switch_to_good': None, 428 'switch_to_bad': None, 429 'install_script': None, 430 'test_script': None, 431 'incremental': True, 432 'prune': False, 433 'iterations': 50, 434 'prune_iterations': 100, 435 'verify_level': 1, 436 'file_args': False, 437 'verbose': False 438 } 439 default_kwargs.update(kwargs) 440 super(MockBinarySearchState, self).__init__(**default_kwargs) 441 442 443def _CanonicalizeScript(script_name): 444 """Return canonical path to script. 445 446 Args: 447 script_name: Relative or absolute path to script 448 449 Returns: 450 Canonicalized script path 451 """ 452 script_name = os.path.expanduser(script_name) 453 if not script_name.startswith('/'): 454 return os.path.join('.', script_name) 455 456 457def Run(get_initial_items, switch_to_good, switch_to_bad, test_script, 458 install_script=None, iterations=50, prune=True, noincremental=False, 459 file_args=False, verify_level=1, prune_iterations=100, verbose=False, 460 resume=False): 461 """Run binary search tool. Equivalent to running through terminal. 462 463 Args: 464 get_initial_items: Script to enumerate all items being binary searched 465 switch_to_good: Script that will take items as input and switch them to good 466 set 467 switch_to_bad: Script that will take items as input and switch them to bad 468 set 469 test_script: Script that will determine if the current combination of good 470 and bad items make a "good" or "bad" result. 471 install_script: Script to do necessary setup (building, compilation, etc.) 472 for test_script 473 iterations: How many binary search iterations to run before exiting. 474 prune: If False the binary search tool will stop when the first bad item is 475 found. Otherwise then binary search tool will continue searching 476 until all bad items are found (or prune_iterations is reached). 477 noincremental: Whether to send "diffs" of good/bad items to switch scripts. 478 file_args: If True then arguments to switch scripts will be a file name 479 containing a newline separated list of the items to switch. 480 verify_level: How many verification tests to run to ensure initial good/bad 481 sets actually produce a good/bad result. 482 prune_iterations: Max number of bad items to search for. 483 verbose: If True will print extra debug information to user. 484 resume: If True will resume using STATE_FILE. 485 486 Returns: 487 0 for success, error otherwise 488 """ 489 if resume: 490 bss = BinarySearchState.LoadState() 491 if not bss: 492 logger.GetLogger().LogOutput( 493 '%s is not a valid binary_search_tool state file, cannot resume!' % 494 STATE_FILE) 495 return 1 496 else: 497 switch_to_good = _CanonicalizeScript(switch_to_good) 498 switch_to_bad = _CanonicalizeScript(switch_to_bad) 499 if install_script: 500 install_script = _CanonicalizeScript(install_script) 501 test_script = _CanonicalizeScript(test_script) 502 get_initial_items = _CanonicalizeScript(get_initial_items) 503 incremental = not noincremental 504 505 binary_search_perforce.verbose = verbose 506 507 bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad, 508 install_script, test_script, incremental, prune, 509 iterations, prune_iterations, verify_level, 510 file_args, verbose) 511 bss.DoVerify() 512 513 try: 514 bss.DoSearch() 515 bss.RemoveState() 516 logger.GetLogger().LogOutput('Total execution time: %s' % 517 bss.ElapsedTimeString()) 518 except Error as e: 519 logger.GetLogger().LogError(e) 520 return 1 521 522 return 0 523 524 525def Main(argv): 526 """The main function.""" 527 # Common initializations 528 529 parser = argparse.ArgumentParser() 530 common.BuildArgParser(parser) 531 logger.GetLogger().LogOutput(' '.join(argv)) 532 options = parser.parse_args(argv) 533 534 if not (options.get_initial_items and options.switch_to_good and 535 options.switch_to_bad and options.test_script) and not options.resume: 536 parser.print_help() 537 return 1 538 539 if options.resume: 540 logger.GetLogger().LogOutput('Resuming from %s' % STATE_FILE) 541 if len(argv) > 1: 542 logger.GetLogger().LogOutput(('Note: resuming from previous state, ' 543 'ignoring given options and loading saved ' 544 'options instead.')) 545 546 # Get dictionary of all options 547 args = vars(options) 548 return Run(**args) 549 550 551if __name__ == '__main__': 552 sys.exit(Main(sys.argv[1:])) 553