binary_search_state.py revision a5e3929d7a155606bc541777fc84e4989389e13c 
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               test_setup_script, test_script, incremental, prune, iterations,
66               prune_iterations, verify, 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.test_setup_script = test_setup_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 = verify
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 TestSetupScript(self):
172    """Run test setup script and return exit code from script."""
173    if not self.test_setup_script:
174      return 0
175
176    command = self.test_setup_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 is False this step is skipped.
186    """
187    if not self.verify:
188      return
189
190    self.l.LogOutput('VERIFICATION')
191    self.l.LogOutput('Beginning tests to verify good/bad sets\n')
192
193    self._OutputProgress('Verifying items from GOOD set\n')
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.TestSetupScript()
198      assert status == 0, 'When reset_to_good, test setup should succeed.'
199      status = self.TestScript()
200      assert status == 0, 'When reset_to_good, status should be 0.'
201
202    self._OutputProgress('Verifying items from BAD set\n')
203    with SetFile(GOOD_SET_VAR, []), SetFile(BAD_SET_VAR, self.all_items):
204      self.l.LogOutput('Resetting all items to bad to verify.')
205      self.SwitchToBad(self.all_items)
206      status = self.TestSetupScript()
207      assert status == 0, 'When reset_to_bad, test setup should succeed.'
208      status = self.TestScript()
209      assert status == 1, 'When reset_to_bad, status should be 1.'
210
211  def DoSearch(self):
212    """Perform full search for bad items.
213
214    Perform full search until prune_iterations number of bad items are found.
215    """
216    while (True and
217           len(self.all_items) > 1 and
218           self.prune_cycles < self.prune_iterations):
219      terminated = self.DoBinarySearch()
220      self.prune_cycles += 1
221      if not terminated:
222        break
223      # Prune is set.
224      prune_index = self.binary_search.current
225
226      # If found item is last item, no new items can be found
227      if prune_index == len(self.all_items) - 1:
228        self.l.LogOutput('First bad item is the last item. Breaking.')
229        self.l.LogOutput('Bad items are: %s' % self.all_items[-1])
230        break
231
232      # If already seen item we have no new bad items to find, finish up
233      if self.all_items[prune_index] in self.found_items:
234        self.l.LogOutput('Found item already found before: %s.' %
235                         self.all_items[prune_index],
236                         print_to_console=self.verbose)
237        self.l.LogOutput('No more bad items remaining. Done searching.')
238        self.l.LogOutput('Bad items are: %s' %
239                         ' '.join(self.found_items))
240        break
241
242      new_all_items = list(self.all_items)
243      # Move prune item to the end of the list.
244      new_all_items.append(new_all_items.pop(prune_index))
245      self.found_items.add(new_all_items[-1])
246
247      # Everything below newly found bad item is now known to be a good item.
248      # Take these good items out of the equation to save time on the next
249      # search. We save these known good items so they are still sent to the
250      # switch_to_good script.
251      if prune_index:
252        self.known_good.update(new_all_items[:prune_index])
253        new_all_items = new_all_items[prune_index:]
254
255      self.l.LogOutput('Old list: %s. New list: %s' % (str(self.all_items),
256                                                       str(new_all_items)),
257                       print_to_console=self.verbose)
258
259      if not self.prune:
260        self.l.LogOutput('Not continuning further, --prune is not set')
261        break
262      # FIXME: Do we need to Convert the currently good items to bad
263      self.PopulateItemsUsingList(new_all_items)
264
265  def DoBinarySearch(self):
266    """Perform single iteration of binary search."""
267    # If in resume mode don't reset search_cycles
268    if not self.resumed:
269      self.search_cycles = 0
270    else:
271      self.resumed = False
272
273    terminated = False
274    while self.search_cycles < self.iterations and not terminated:
275      self.SaveState()
276      self.OutputIterationProgress()
277
278      self.search_cycles += 1
279      [bad_items, good_items] = self.GetNextItems()
280
281      with SetFile(GOOD_SET_VAR, good_items), SetFile(BAD_SET_VAR, bad_items):
282        # TODO: bad_items should come first.
283        self.SwitchToGood(good_items)
284        self.SwitchToBad(bad_items)
285        status = self.TestSetupScript()
286        if status == 0:
287          status = self.TestScript()
288        else:
289          # Test setup script failed, treat as skipped item
290          status = 125
291        terminated = self.binary_search.SetStatus(status)
292
293      if terminated:
294        self.l.LogOutput('Terminated!', print_to_console=self.verbose)
295    if not terminated:
296      self.l.LogOutput('Ran out of iterations searching...')
297    self.l.LogOutput(str(self), print_to_console=self.verbose)
298    return terminated
299
300  def PopulateItemsUsingCommand(self, command):
301    """Update all_items and binary search logic from executable.
302
303    This method is mainly required for enumerating the initial list of items
304    from the get_initial_items script.
305
306    Args:
307      command: path to executable that will enumerate items.
308    """
309    ce = command_executer.GetCommandExecuter()
310    _, out, _ = ce.RunCommandWExceptionCleanup(command,
311                                               return_output=True,
312                                               print_to_console=self.verbose)
313    all_items = out.split()
314    self.PopulateItemsUsingList(all_items)
315
316  def PopulateItemsUsingList(self, all_items):
317    """Update all_items and binary searching logic from list.
318
319    Args:
320      all_items: new list of all_items
321    """
322    self.all_items = all_items
323    self.binary_search = binary_search_perforce.BinarySearcher(
324        logger_to_set=self.l)
325    self.binary_search.SetSortedList(self.all_items)
326
327  def SaveState(self):
328    """Save state to STATE_FILE.
329
330    SaveState will create a new unique, hidden state file to hold data from
331    object. Then atomically overwrite the STATE_FILE symlink to point to the
332    new data.
333
334    Raises:
335      Error if STATE_FILE already exists but is not a symlink.
336    """
337    ce, l = self.ce, self.l
338    self.ce, self.l, self.binary_search.logger = None, None, None
339    old_state = None
340
341    _, path = tempfile.mkstemp(prefix=HIDDEN_STATE_FILE, dir='.')
342    with open(path, 'wb') as f:
343      pickle.dump(self, f)
344
345    if os.path.exists(STATE_FILE):
346      if os.path.islink(STATE_FILE):
347        old_state = os.readlink(STATE_FILE)
348      else:
349        raise Error(('%s already exists and is not a symlink!\n'
350                     'State file saved to %s' % (STATE_FILE, path)))
351
352    # Create new link and atomically overwrite old link
353    temp_link = '%s.link' % HIDDEN_STATE_FILE
354    os.symlink(path, temp_link)
355    os.rename(temp_link, STATE_FILE)
356
357    if old_state:
358      os.remove(old_state)
359
360    self.ce, self.l, self.binary_search.logger = ce, l, l
361
362  @classmethod
363  def LoadState(cls):
364    """Create BinarySearchState object from STATE_FILE."""
365    if not os.path.isfile(STATE_FILE):
366      return None
367    try:
368      bss = pickle.load(file(STATE_FILE))
369      bss.l = logger.GetLogger()
370      bss.ce = command_executer.GetCommandExecuter()
371      bss.binary_search.logger = bss.l
372      bss.start_time = time.time()
373      bss.resumed = True
374      binary_search_perforce.verbose = bss.verbose
375      return bss
376    except StandardError:
377      return None
378
379  def RemoveState(self):
380    """Remove STATE_FILE and its symlinked data from file system."""
381    if os.path.exists(STATE_FILE):
382      if os.path.islink(STATE_FILE):
383        real_file = os.readlink(STATE_FILE)
384        os.remove(real_file)
385        os.remove(STATE_FILE)
386
387  def GetNextItems(self):
388    """Get next items for binary search based on result of the last test run."""
389    border_item = self.binary_search.GetNext()
390    index = self.all_items.index(border_item)
391
392    next_bad_items = self.all_items[:index + 1]
393    next_good_items = self.all_items[index + 1:] + list(self.known_good)
394
395    return [next_bad_items, next_good_items]
396
397  def ElapsedTimeString(self):
398    """Return h m s format of elapsed time since execution has started."""
399    diff = int(time.time() - self.start_time)
400    seconds = diff % 60
401    minutes = (diff / 60) % 60
402    hours = diff / (60 * 60)
403
404    seconds = str(seconds).rjust(2)
405    minutes = str(minutes).rjust(2)
406    hours = str(hours).rjust(2)
407
408    return '%sh %sm %ss' % (hours, minutes, seconds)
409
410  def _OutputProgress(self, progress_text):
411    """Output current progress of binary search to console and logs.
412
413    Args:
414      progress_text: The progress to display to the user.
415    """
416    progress = ('\n***** PROGRESS (elapsed time: %s) *****\n'
417                '%s'
418                '************************************************')
419    progress = progress % (self.ElapsedTimeString(), progress_text)
420    self.l.LogOutput(progress)
421
422  def OutputIterationProgress(self):
423    out = ('Search %d of estimated %d.\n'
424           'Prune %d of max %d.\n'
425           'Current bad items found:\n'
426           '%s\n')
427    out = out % (self.search_cycles + 1,
428                 math.ceil(math.log(len(self.all_items), 2)),
429                 self.prune_cycles + 1,
430                 self.prune_iterations,
431                 ', '.join(self.found_items))
432    self._OutputProgress(out)
433
434  def __str__(self):
435    ret = ''
436    ret += 'all: %s\n' % str(self.all_items)
437    ret += 'currently_good: %s\n' % str(self.currently_good_items)
438    ret += 'currently_bad: %s\n' % str(self.currently_bad_items)
439    ret += str(self.binary_search)
440    return ret
441
442
443class MockBinarySearchState(BinarySearchState):
444  """Mock class for BinarySearchState."""
445
446  def __init__(self, **kwargs):
447    # Initialize all arguments to None
448    default_kwargs = {
449        'get_initial_items': 'echo "1"',
450        'switch_to_good': None,
451        'switch_to_bad': None,
452        'test_setup_script': None,
453        'test_script': None,
454        'incremental': True,
455        'prune': False,
456        'iterations': 50,
457        'prune_iterations': 100,
458        'verify': True,
459        'file_args': False,
460        'verbose': False
461    }
462    default_kwargs.update(kwargs)
463    super(MockBinarySearchState, self).__init__(**default_kwargs)
464
465
466def _CanonicalizeScript(script_name):
467  """Return canonical path to script.
468
469  Args:
470    script_name: Relative or absolute path to script
471
472  Returns:
473    Canonicalized script path
474  """
475  script_name = os.path.expanduser(script_name)
476  if not script_name.startswith('/'):
477    return os.path.join('.', script_name)
478
479
480def Run(get_initial_items, switch_to_good, switch_to_bad, test_script,
481        test_setup_script=None, iterations=50, prune=False, noincremental=False,
482        file_args=False, verify=True, prune_iterations=100, verbose=False,
483        resume=False):
484  """Run binary search tool. Equivalent to running through terminal.
485
486  Args:
487    get_initial_items: Script to enumerate all items being binary searched
488    switch_to_good: Script that will take items as input and switch them to good
489                    set
490    switch_to_bad: Script that will take items as input and switch them to bad
491                   set
492    test_script: Script that will determine if the current combination of good
493                 and bad items make a "good" or "bad" result.
494    test_setup_script: Script to do necessary setup (building, compilation,
495                       etc.) for test_script.
496    iterations: How many binary search iterations to run before exiting.
497    prune: If False the binary search tool will stop when the first bad item is
498           found. Otherwise then binary search tool will continue searching
499           until all bad items are found (or prune_iterations is reached).
500    noincremental: Whether to send "diffs" of good/bad items to switch scripts.
501    file_args: If True then arguments to switch scripts will be a file name
502               containing a newline separated list of the items to switch.
503    verify: If True, run tests to ensure initial good/bad sets actually
504            produce a good/bad result.
505    prune_iterations: Max number of bad items to search for.
506    verbose: If True will print extra debug information to user.
507    resume: If True will resume using STATE_FILE.
508
509  Returns:
510    0 for success, error otherwise
511  """
512  if resume:
513    bss = BinarySearchState.LoadState()
514    if not bss:
515      logger.GetLogger().LogOutput(
516          '%s is not a valid binary_search_tool state file, cannot resume!' %
517          STATE_FILE)
518      return 1
519  else:
520    switch_to_good = _CanonicalizeScript(switch_to_good)
521    switch_to_bad = _CanonicalizeScript(switch_to_bad)
522    if test_setup_script:
523      test_setup_script = _CanonicalizeScript(test_setup_script)
524    test_script = _CanonicalizeScript(test_script)
525    get_initial_items = _CanonicalizeScript(get_initial_items)
526    incremental = not noincremental
527
528    binary_search_perforce.verbose = verbose
529
530    bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad,
531                            test_setup_script, test_script, incremental, prune,
532                            iterations, prune_iterations, verify, file_args,
533                            verbose)
534    bss.DoVerify()
535
536  try:
537    bss.DoSearch()
538    bss.RemoveState()
539    logger.GetLogger().LogOutput('Total execution time: %s' %
540                                 bss.ElapsedTimeString())
541  except Error as e:
542    logger.GetLogger().LogError(e)
543    return 1
544
545  return 0
546
547
548def Main(argv):
549  """The main function."""
550  # Common initializations
551
552  parser = argparse.ArgumentParser()
553  common.BuildArgParser(parser)
554  logger.GetLogger().LogOutput(' '.join(argv))
555  options = parser.parse_args(argv)
556
557  if not (options.get_initial_items and options.switch_to_good and
558          options.switch_to_bad and options.test_script) and not options.resume:
559    parser.print_help()
560    return 1
561
562  if options.resume:
563    logger.GetLogger().LogOutput('Resuming from %s' % STATE_FILE)
564    if len(argv) > 1:
565      logger.GetLogger().LogOutput(('Note: resuming from previous state, '
566                                    'ignoring given options and loading saved '
567                                    'options instead.'))
568
569  # Get dictionary of all options
570  args = vars(options)
571  return Run(**args)
572
573
574if __name__ == '__main__':
575  sys.exit(Main(sys.argv[1:]))
576