You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
598 lines
20 KiB
598 lines
20 KiB
#!/usr/bin/python2 |
|
"""The binary search wrapper.""" |
|
|
|
from __future__ import print_function |
|
|
|
import argparse |
|
import contextlib |
|
import errno |
|
import math |
|
import os |
|
import pickle |
|
import sys |
|
import tempfile |
|
import time |
|
|
|
# Adds cros_utils to PYTHONPATH |
|
import common |
|
|
|
# Now we do import from cros_utils |
|
from cros_utils import command_executer |
|
from cros_utils import logger |
|
|
|
import binary_search_perforce |
|
|
|
GOOD_SET_VAR = 'BISECT_GOOD_SET' |
|
BAD_SET_VAR = 'BISECT_BAD_SET' |
|
|
|
STATE_FILE = '%s.state' % sys.argv[0] |
|
HIDDEN_STATE_FILE = os.path.join( |
|
os.path.dirname(STATE_FILE), '.%s' % os.path.basename(STATE_FILE)) |
|
|
|
|
|
class Error(Exception): |
|
"""The general binary search tool error class.""" |
|
pass |
|
|
|
|
|
@contextlib.contextmanager |
|
def SetFile(env_var, items): |
|
"""Generate set files that can be used by switch/test scripts. |
|
|
|
Generate temporary set file (good/bad) holding contents of good/bad items for |
|
the current binary search iteration. Store the name of each file as an |
|
environment variable so all child processes can access it. |
|
|
|
This function is a contextmanager, meaning it's meant to be used with the |
|
"with" statement in Python. This is so cleanup and setup happens automatically |
|
and cleanly. Execution of the outer "with" statement happens at the "yield" |
|
statement. |
|
|
|
Args: |
|
env_var: What environment variable to store the file name in. |
|
items: What items are in this set. |
|
""" |
|
with tempfile.NamedTemporaryFile() as f: |
|
os.environ[env_var] = f.name |
|
f.write('\n'.join(items)) |
|
f.flush() |
|
yield |
|
|
|
|
|
class BinarySearchState(object): |
|
"""The binary search state class.""" |
|
|
|
def __init__(self, get_initial_items, switch_to_good, switch_to_bad, |
|
test_setup_script, test_script, incremental, prune, iterations, |
|
prune_iterations, verify, file_args, verbose): |
|
"""BinarySearchState constructor, see Run for full args documentation.""" |
|
self.get_initial_items = get_initial_items |
|
self.switch_to_good = switch_to_good |
|
self.switch_to_bad = switch_to_bad |
|
self.test_setup_script = test_setup_script |
|
self.test_script = test_script |
|
self.incremental = incremental |
|
self.prune = prune |
|
self.iterations = iterations |
|
self.prune_iterations = prune_iterations |
|
self.verify = verify |
|
self.file_args = file_args |
|
self.verbose = verbose |
|
|
|
self.l = logger.GetLogger() |
|
self.ce = command_executer.GetCommandExecuter() |
|
|
|
self.resumed = False |
|
self.prune_cycles = 0 |
|
self.search_cycles = 0 |
|
self.binary_search = None |
|
self.all_items = None |
|
self.PopulateItemsUsingCommand(self.get_initial_items) |
|
self.currently_good_items = set([]) |
|
self.currently_bad_items = set([]) |
|
self.found_items = set([]) |
|
self.known_good = set([]) |
|
|
|
self.start_time = time.time() |
|
|
|
def SwitchToGood(self, item_list): |
|
"""Switch given items to "good" set.""" |
|
if self.incremental: |
|
self.l.LogOutput( |
|
'Incremental set. Wanted to switch %s to good' % str(item_list), |
|
print_to_console=self.verbose) |
|
incremental_items = [ |
|
item for item in item_list if item not in self.currently_good_items |
|
] |
|
item_list = incremental_items |
|
self.l.LogOutput( |
|
'Incremental set. Actually switching %s to good' % str(item_list), |
|
print_to_console=self.verbose) |
|
|
|
if not item_list: |
|
return |
|
|
|
self.l.LogOutput( |
|
'Switching %s to good' % str(item_list), print_to_console=self.verbose) |
|
self.RunSwitchScript(self.switch_to_good, item_list) |
|
self.currently_good_items = self.currently_good_items.union(set(item_list)) |
|
self.currently_bad_items.difference_update(set(item_list)) |
|
|
|
def SwitchToBad(self, item_list): |
|
"""Switch given items to "bad" set.""" |
|
if self.incremental: |
|
self.l.LogOutput( |
|
'Incremental set. Wanted to switch %s to bad' % str(item_list), |
|
print_to_console=self.verbose) |
|
incremental_items = [ |
|
item for item in item_list if item not in self.currently_bad_items |
|
] |
|
item_list = incremental_items |
|
self.l.LogOutput( |
|
'Incremental set. Actually switching %s to bad' % str(item_list), |
|
print_to_console=self.verbose) |
|
|
|
if not item_list: |
|
return |
|
|
|
self.l.LogOutput( |
|
'Switching %s to bad' % str(item_list), print_to_console=self.verbose) |
|
self.RunSwitchScript(self.switch_to_bad, item_list) |
|
self.currently_bad_items = self.currently_bad_items.union(set(item_list)) |
|
self.currently_good_items.difference_update(set(item_list)) |
|
|
|
def RunSwitchScript(self, switch_script, item_list): |
|
"""Pass given items to switch script. |
|
|
|
Args: |
|
switch_script: path to switch script |
|
item_list: list of all items to be switched |
|
""" |
|
if self.file_args: |
|
with tempfile.NamedTemporaryFile() as f: |
|
f.write('\n'.join(item_list)) |
|
f.flush() |
|
command = '%s %s' % (switch_script, f.name) |
|
ret, _, _ = self.ce.RunCommandWExceptionCleanup( |
|
command, print_to_console=self.verbose) |
|
else: |
|
command = '%s %s' % (switch_script, ' '.join(item_list)) |
|
try: |
|
ret, _, _ = self.ce.RunCommandWExceptionCleanup( |
|
command, print_to_console=self.verbose) |
|
except OSError as e: |
|
if e.errno == errno.E2BIG: |
|
raise Error('Too many arguments for switch script! Use --file_args') |
|
else: |
|
raise |
|
assert ret == 0, 'Switch script %s returned %d' % (switch_script, ret) |
|
|
|
def TestScript(self): |
|
"""Run test script and return exit code from script.""" |
|
command = self.test_script |
|
ret, _, _ = self.ce.RunCommandWExceptionCleanup(command) |
|
return ret |
|
|
|
def TestSetupScript(self): |
|
"""Run test setup script and return exit code from script.""" |
|
if not self.test_setup_script: |
|
return 0 |
|
|
|
command = self.test_setup_script |
|
ret, _, _ = self.ce.RunCommandWExceptionCleanup(command) |
|
return ret |
|
|
|
def DoVerify(self): |
|
"""Verify correctness of test environment. |
|
|
|
Verify that a "good" set of items produces a "good" result and that a "bad" |
|
set of items produces a "bad" result. To be run directly before running |
|
DoSearch. If verify is False this step is skipped. |
|
""" |
|
if not self.verify: |
|
return |
|
|
|
self.l.LogOutput('VERIFICATION') |
|
self.l.LogOutput('Beginning tests to verify good/bad sets\n') |
|
|
|
self._OutputProgress('Verifying items from GOOD set\n') |
|
with SetFile(GOOD_SET_VAR, self.all_items), SetFile(BAD_SET_VAR, []): |
|
self.l.LogOutput('Resetting all items to good to verify.') |
|
self.SwitchToGood(self.all_items) |
|
status = self.TestSetupScript() |
|
assert status == 0, 'When reset_to_good, test setup should succeed.' |
|
status = self.TestScript() |
|
assert status == 0, 'When reset_to_good, status should be 0.' |
|
|
|
self._OutputProgress('Verifying items from BAD set\n') |
|
with SetFile(GOOD_SET_VAR, []), SetFile(BAD_SET_VAR, self.all_items): |
|
self.l.LogOutput('Resetting all items to bad to verify.') |
|
self.SwitchToBad(self.all_items) |
|
status = self.TestSetupScript() |
|
# The following assumption is not true; a bad image might not |
|
# successfully push onto a device. |
|
# assert status == 0, 'When reset_to_bad, test setup should succeed.' |
|
if status == 0: |
|
status = self.TestScript() |
|
assert status == 1, 'When reset_to_bad, status should be 1.' |
|
|
|
def DoSearch(self): |
|
"""Perform full search for bad items. |
|
|
|
Perform full search until prune_iterations number of bad items are found. |
|
""" |
|
while (True and len(self.all_items) > 1 and |
|
self.prune_cycles < self.prune_iterations): |
|
terminated = self.DoBinarySearch() |
|
self.prune_cycles += 1 |
|
if not terminated: |
|
break |
|
# Prune is set. |
|
prune_index = self.binary_search.current |
|
|
|
# If found item is last item, no new items can be found |
|
if prune_index == len(self.all_items) - 1: |
|
self.l.LogOutput('First bad item is the last item. Breaking.') |
|
self.l.LogOutput('Bad items are: %s' % self.all_items[-1]) |
|
break |
|
|
|
# If already seen item we have no new bad items to find, finish up |
|
if self.all_items[prune_index] in self.found_items: |
|
self.l.LogOutput( |
|
'Found item already found before: %s.' % |
|
self.all_items[prune_index], |
|
print_to_console=self.verbose) |
|
self.l.LogOutput('No more bad items remaining. Done searching.') |
|
self.l.LogOutput('Bad items are: %s' % ' '.join(self.found_items)) |
|
break |
|
|
|
new_all_items = list(self.all_items) |
|
# Move prune item to the end of the list. |
|
new_all_items.append(new_all_items.pop(prune_index)) |
|
self.found_items.add(new_all_items[-1]) |
|
|
|
# Everything below newly found bad item is now known to be a good item. |
|
# Take these good items out of the equation to save time on the next |
|
# search. We save these known good items so they are still sent to the |
|
# switch_to_good script. |
|
if prune_index: |
|
self.known_good.update(new_all_items[:prune_index]) |
|
new_all_items = new_all_items[prune_index:] |
|
|
|
self.l.LogOutput( |
|
'Old list: %s. New list: %s' % (str(self.all_items), |
|
str(new_all_items)), |
|
print_to_console=self.verbose) |
|
|
|
if not self.prune: |
|
self.l.LogOutput('Not continuning further, --prune is not set') |
|
break |
|
# FIXME: Do we need to Convert the currently good items to bad |
|
self.PopulateItemsUsingList(new_all_items) |
|
|
|
def DoBinarySearch(self): |
|
"""Perform single iteration of binary search.""" |
|
# If in resume mode don't reset search_cycles |
|
if not self.resumed: |
|
self.search_cycles = 0 |
|
else: |
|
self.resumed = False |
|
|
|
terminated = False |
|
while self.search_cycles < self.iterations and not terminated: |
|
self.SaveState() |
|
self.OutputIterationProgress() |
|
|
|
self.search_cycles += 1 |
|
[bad_items, good_items] = self.GetNextItems() |
|
|
|
with SetFile(GOOD_SET_VAR, good_items), SetFile(BAD_SET_VAR, bad_items): |
|
# TODO: bad_items should come first. |
|
self.SwitchToGood(good_items) |
|
self.SwitchToBad(bad_items) |
|
status = self.TestSetupScript() |
|
if status == 0: |
|
status = self.TestScript() |
|
terminated = self.binary_search.SetStatus(status) |
|
|
|
if terminated: |
|
self.l.LogOutput('Terminated!', print_to_console=self.verbose) |
|
if not terminated: |
|
self.l.LogOutput('Ran out of iterations searching...') |
|
self.l.LogOutput(str(self), print_to_console=self.verbose) |
|
return terminated |
|
|
|
def PopulateItemsUsingCommand(self, command): |
|
"""Update all_items and binary search logic from executable. |
|
|
|
This method is mainly required for enumerating the initial list of items |
|
from the get_initial_items script. |
|
|
|
Args: |
|
command: path to executable that will enumerate items. |
|
""" |
|
ce = command_executer.GetCommandExecuter() |
|
_, out, _ = ce.RunCommandWExceptionCleanup( |
|
command, return_output=True, print_to_console=self.verbose) |
|
all_items = out.split() |
|
self.PopulateItemsUsingList(all_items) |
|
|
|
def PopulateItemsUsingList(self, all_items): |
|
"""Update all_items and binary searching logic from list. |
|
|
|
Args: |
|
all_items: new list of all_items |
|
""" |
|
self.all_items = all_items |
|
self.binary_search = binary_search_perforce.BinarySearcher( |
|
logger_to_set=self.l) |
|
self.binary_search.SetSortedList(self.all_items) |
|
|
|
def SaveState(self): |
|
"""Save state to STATE_FILE. |
|
|
|
SaveState will create a new unique, hidden state file to hold data from |
|
object. Then atomically overwrite the STATE_FILE symlink to point to the |
|
new data. |
|
|
|
Raises: |
|
Error if STATE_FILE already exists but is not a symlink. |
|
""" |
|
ce, l = self.ce, self.l |
|
self.ce, self.l, self.binary_search.logger = None, None, None |
|
old_state = None |
|
|
|
_, path = tempfile.mkstemp(prefix=HIDDEN_STATE_FILE, dir='.') |
|
with open(path, 'wb') as f: |
|
pickle.dump(self, f) |
|
|
|
if os.path.exists(STATE_FILE): |
|
if os.path.islink(STATE_FILE): |
|
old_state = os.readlink(STATE_FILE) |
|
else: |
|
raise Error(('%s already exists and is not a symlink!\n' |
|
'State file saved to %s' % (STATE_FILE, path))) |
|
|
|
# Create new link and atomically overwrite old link |
|
temp_link = '%s.link' % HIDDEN_STATE_FILE |
|
os.symlink(path, temp_link) |
|
os.rename(temp_link, STATE_FILE) |
|
|
|
if old_state: |
|
os.remove(old_state) |
|
|
|
self.ce, self.l, self.binary_search.logger = ce, l, l |
|
|
|
@classmethod |
|
def LoadState(cls): |
|
"""Create BinarySearchState object from STATE_FILE.""" |
|
if not os.path.isfile(STATE_FILE): |
|
return None |
|
try: |
|
bss = pickle.load(file(STATE_FILE)) |
|
bss.l = logger.GetLogger() |
|
bss.ce = command_executer.GetCommandExecuter() |
|
bss.binary_search.logger = bss.l |
|
bss.start_time = time.time() |
|
|
|
# Set resumed to be True so we can enter DoBinarySearch without the method |
|
# resetting our current search_cycles to 0. |
|
bss.resumed = True |
|
|
|
# Set currently_good_items and currently_bad_items to empty so that the |
|
# first iteration after resuming will always be non-incremental. This is |
|
# just in case the environment changes, the user makes manual changes, or |
|
# a previous switch_script corrupted the environment. |
|
bss.currently_good_items = set([]) |
|
bss.currently_bad_items = set([]) |
|
|
|
binary_search_perforce.verbose = bss.verbose |
|
return bss |
|
except StandardError: |
|
return None |
|
|
|
def RemoveState(self): |
|
"""Remove STATE_FILE and its symlinked data from file system.""" |
|
if os.path.exists(STATE_FILE): |
|
if os.path.islink(STATE_FILE): |
|
real_file = os.readlink(STATE_FILE) |
|
os.remove(real_file) |
|
os.remove(STATE_FILE) |
|
|
|
def GetNextItems(self): |
|
"""Get next items for binary search based on result of the last test run.""" |
|
border_item = self.binary_search.GetNext() |
|
index = self.all_items.index(border_item) |
|
|
|
next_bad_items = self.all_items[:index + 1] |
|
next_good_items = self.all_items[index + 1:] + list(self.known_good) |
|
|
|
return [next_bad_items, next_good_items] |
|
|
|
def ElapsedTimeString(self): |
|
"""Return h m s format of elapsed time since execution has started.""" |
|
diff = int(time.time() - self.start_time) |
|
seconds = diff % 60 |
|
minutes = (diff / 60) % 60 |
|
hours = diff / (60 * 60) |
|
|
|
seconds = str(seconds).rjust(2) |
|
minutes = str(minutes).rjust(2) |
|
hours = str(hours).rjust(2) |
|
|
|
return '%sh %sm %ss' % (hours, minutes, seconds) |
|
|
|
def _OutputProgress(self, progress_text): |
|
"""Output current progress of binary search to console and logs. |
|
|
|
Args: |
|
progress_text: The progress to display to the user. |
|
""" |
|
progress = ('\n***** PROGRESS (elapsed time: %s) *****\n' |
|
'%s' |
|
'************************************************') |
|
progress = progress % (self.ElapsedTimeString(), progress_text) |
|
self.l.LogOutput(progress) |
|
|
|
def OutputIterationProgress(self): |
|
out = ('Search %d of estimated %d.\n' |
|
'Prune %d of max %d.\n' |
|
'Current bad items found:\n' |
|
'%s\n') |
|
out = out % (self.search_cycles + 1, |
|
math.ceil(math.log(len(self.all_items), 2)), |
|
self.prune_cycles + 1, self.prune_iterations, |
|
', '.join(self.found_items)) |
|
self._OutputProgress(out) |
|
|
|
def __str__(self): |
|
ret = '' |
|
ret += 'all: %s\n' % str(self.all_items) |
|
ret += 'currently_good: %s\n' % str(self.currently_good_items) |
|
ret += 'currently_bad: %s\n' % str(self.currently_bad_items) |
|
ret += str(self.binary_search) |
|
return ret |
|
|
|
|
|
class MockBinarySearchState(BinarySearchState): |
|
"""Mock class for BinarySearchState.""" |
|
|
|
def __init__(self, **kwargs): |
|
# Initialize all arguments to None |
|
default_kwargs = { |
|
'get_initial_items': 'echo "1"', |
|
'switch_to_good': None, |
|
'switch_to_bad': None, |
|
'test_setup_script': None, |
|
'test_script': None, |
|
'incremental': True, |
|
'prune': False, |
|
'iterations': 50, |
|
'prune_iterations': 100, |
|
'verify': True, |
|
'file_args': False, |
|
'verbose': False |
|
} |
|
default_kwargs.update(kwargs) |
|
super(MockBinarySearchState, self).__init__(**default_kwargs) |
|
|
|
|
|
def _CanonicalizeScript(script_name): |
|
"""Return canonical path to script. |
|
|
|
Args: |
|
script_name: Relative or absolute path to script |
|
|
|
Returns: |
|
Canonicalized script path |
|
""" |
|
script_name = os.path.expanduser(script_name) |
|
if not script_name.startswith('/'): |
|
return os.path.join('.', script_name) |
|
|
|
|
|
def Run(get_initial_items, |
|
switch_to_good, |
|
switch_to_bad, |
|
test_script, |
|
test_setup_script=None, |
|
iterations=50, |
|
prune=False, |
|
noincremental=False, |
|
file_args=False, |
|
verify=True, |
|
prune_iterations=100, |
|
verbose=False, |
|
resume=False): |
|
"""Run binary search tool. Equivalent to running through terminal. |
|
|
|
Args: |
|
get_initial_items: Script to enumerate all items being binary searched |
|
switch_to_good: Script that will take items as input and switch them to good |
|
set |
|
switch_to_bad: Script that will take items as input and switch them to bad |
|
set |
|
test_script: Script that will determine if the current combination of good |
|
and bad items make a "good" or "bad" result. |
|
test_setup_script: Script to do necessary setup (building, compilation, |
|
etc.) for test_script. |
|
iterations: How many binary search iterations to run before exiting. |
|
prune: If False the binary search tool will stop when the first bad item is |
|
found. Otherwise then binary search tool will continue searching |
|
until all bad items are found (or prune_iterations is reached). |
|
noincremental: Whether to send "diffs" of good/bad items to switch scripts. |
|
file_args: If True then arguments to switch scripts will be a file name |
|
containing a newline separated list of the items to switch. |
|
verify: If True, run tests to ensure initial good/bad sets actually |
|
produce a good/bad result. |
|
prune_iterations: Max number of bad items to search for. |
|
verbose: If True will print extra debug information to user. |
|
resume: If True will resume using STATE_FILE. |
|
|
|
Returns: |
|
0 for success, error otherwise |
|
""" |
|
if resume: |
|
bss = BinarySearchState.LoadState() |
|
if not bss: |
|
logger.GetLogger().LogOutput( |
|
'%s is not a valid binary_search_tool state file, cannot resume!' % |
|
STATE_FILE) |
|
return 1 |
|
else: |
|
switch_to_good = _CanonicalizeScript(switch_to_good) |
|
switch_to_bad = _CanonicalizeScript(switch_to_bad) |
|
if test_setup_script: |
|
test_setup_script = _CanonicalizeScript(test_setup_script) |
|
test_script = _CanonicalizeScript(test_script) |
|
get_initial_items = _CanonicalizeScript(get_initial_items) |
|
incremental = not noincremental |
|
|
|
binary_search_perforce.verbose = verbose |
|
|
|
bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad, |
|
test_setup_script, test_script, incremental, prune, |
|
iterations, prune_iterations, verify, file_args, |
|
verbose) |
|
bss.DoVerify() |
|
|
|
try: |
|
bss.DoSearch() |
|
bss.RemoveState() |
|
logger.GetLogger().LogOutput('Total execution time: %s' % |
|
bss.ElapsedTimeString()) |
|
except Error as e: |
|
logger.GetLogger().LogError(e) |
|
return 1 |
|
|
|
return 0 |
|
|
|
|
|
def Main(argv): |
|
"""The main function.""" |
|
# Common initializations |
|
|
|
parser = argparse.ArgumentParser() |
|
common.BuildArgParser(parser) |
|
logger.GetLogger().LogOutput(' '.join(argv)) |
|
options = parser.parse_args(argv) |
|
|
|
if not (options.get_initial_items and options.switch_to_good and |
|
options.switch_to_bad and options.test_script) and not options.resume: |
|
parser.print_help() |
|
return 1 |
|
|
|
if options.resume: |
|
logger.GetLogger().LogOutput('Resuming from %s' % STATE_FILE) |
|
if len(argv) > 1: |
|
logger.GetLogger().LogOutput(('Note: resuming from previous state, ' |
|
'ignoring given options and loading saved ' |
|
'options instead.')) |
|
|
|
# Get dictionary of all options |
|
args = vars(options) |
|
return Run(**args) |
|
|
|
|
|
if __name__ == '__main__': |
|
sys.exit(Main(sys.argv[1:]))
|
|
|