binary_search_state.py revision 8f7e4bc040879b7a2eb22033f2c83d99cedbe6f4 
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 Exception:
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