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