binary_search_state.py revision 068c5fb61394771043ed39b11cb996cfd9874c83 
1#!/usr/bin/python2
2"""The binary search wrapper."""
3
4from __future__ import print_function
5
6import argparse
7import math
8import os
9import pickle
10import sys
11import tempfile
12
13# Programtically adding utils python path to PYTHONPATH
14if os.path.isabs(sys.argv[0]):
15  utils_pythonpath = os.path.abspath('{0}/..'.format(os.path.dirname(sys.argv[
16      0])))
17else:
18  wdir = os.getcwd()
19  utils_pythonpath = os.path.abspath('{0}/{1}/..'.format(wdir, os.path.dirname(
20      sys.argv[0])))
21sys.path.append(utils_pythonpath)
22# Now we do import from utils
23from utils import command_executer
24from utils import logger
25
26import binary_search_perforce
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
37class BinarySearchState(object):
38  """The binary search state class."""
39
40  def __init__(self, get_initial_items, switch_to_good, switch_to_bad,
41               install_script, test_script, incremental, prune, iterations,
42               prune_iterations, verify_level, file_args, verbose):
43    self.get_initial_items = get_initial_items
44    self.switch_to_good = switch_to_good
45    self.switch_to_bad = switch_to_bad
46    self.install_script = install_script
47    self.test_script = test_script
48    self.incremental = incremental
49    self.prune = prune
50    self.iterations = iterations
51    self.prune_iterations = prune_iterations
52    self.verify_level = verify_level
53    self.file_args = file_args
54    self.verbose = verbose
55
56    self.l = logger.GetLogger()
57    self.ce = command_executer.GetCommandExecuter()
58
59    self.resumed = False
60    self.prune_cycles = 0
61    self.search_cycles = 0
62    self.binary_search = None
63    self.all_items = None
64    self.PopulateItemsUsingCommand(self.get_initial_items)
65    self.currently_good_items = set([])
66    self.currently_bad_items = set([])
67    self.found_items = set([])
68
69  def SwitchToGood(self, item_list):
70    if self.incremental:
71      self.l.LogOutput('Incremental set. Wanted to switch %s to good' %
72                       str(item_list), print_to_console=self.verbose)
73      incremental_items = [
74          item for item in item_list if item not in self.currently_good_items
75      ]
76      item_list = incremental_items
77      self.l.LogOutput('Incremental set. Actually switching %s to good' %
78                       str(item_list), print_to_console=self.verbose)
79
80    if not item_list:
81      return
82
83    self.l.LogOutput('Switching %s to good' % str(item_list),
84                     print_to_console=self.verbose)
85    self.RunSwitchScript(self.switch_to_good, item_list)
86    self.currently_good_items = self.currently_good_items.union(set(item_list))
87    self.currently_bad_items.difference_update(set(item_list))
88
89  def SwitchToBad(self, item_list):
90    if self.incremental:
91      self.l.LogOutput('Incremental set. Wanted to switch %s to bad' %
92                       str(item_list), print_to_console=self.verbose)
93      incremental_items = [
94          item for item in item_list if item not in self.currently_bad_items
95      ]
96      item_list = incremental_items
97      self.l.LogOutput('Incremental set. Actually switching %s to bad' %
98                       str(item_list), print_to_console=self.verbose)
99
100    if not item_list:
101      return
102
103    self.l.LogOutput('Switching %s to bad' % str(item_list),
104                     print_to_console=self.verbose)
105    self.RunSwitchScript(self.switch_to_bad, item_list)
106    self.currently_bad_items = self.currently_bad_items.union(set(item_list))
107    self.currently_good_items.difference_update(set(item_list))
108
109  def RunSwitchScript(self, switch_script, item_list):
110    if self.file_args:
111      with tempfile.NamedTemporaryFile() as f:
112        f.write('\n'.join(item_list))
113        f.flush()
114        command = '%s %s' % (switch_script, f.name)
115        ret, _, _ = self.ce.RunCommandWExceptionCleanup(
116            command, print_to_console=self.verbose)
117    else:
118      command = '%s %s' % (switch_script, ' '.join(item_list))
119      ret, _, _ = self.ce.RunCommandWExceptionCleanup(
120          command, print_to_console=self.verbose)
121    assert ret == 0, 'Switch script %s returned %d' % (switch_script, ret)
122
123  def TestScript(self):
124    command = self.test_script
125    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
126    return ret
127
128  def InstallScript(self):
129    if not self.install_script:
130      return 0
131
132    command = self.install_script
133    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
134    return ret
135
136  def DoVerify(self):
137    self.l.LogOutput('VERIFICATION')
138    self.l.LogOutput('Beginning %d tests to verify good/bad sets\n')
139    for _ in range(int(self.verify_level)):
140      self.l.LogOutput('Resetting all items to good to verify.')
141      self.SwitchToGood(self.all_items)
142      status = self.InstallScript()
143      assert status == 0, 'When reset_to_good, install should succeed.'
144      status = self.TestScript()
145      assert status == 0, 'When reset_to_good, status should be 0.'
146
147      self.l.LogOutput('Resetting all items to bad to verify.')
148      self.SwitchToBad(self.all_items)
149      status = self.InstallScript()
150      assert status == 0, 'When reset_to_bad, install should succeed.'
151      status = self.TestScript()
152      assert status == 1, 'When reset_to_bad, status should be 1.'
153
154  def DoSearch(self):
155    while (True and
156           len(self.all_items) > 1 and
157           self.prune_cycles < self.prune_iterations):
158      terminated = self.DoBinarySearch()
159      self.prune_cycles += 1
160      if not terminated:
161        break
162      if not self.prune:
163        self.l.LogOutput('Not continuning further, --prune is not set')
164        break
165      # Prune is set.
166      prune_index = self.binary_search.current
167
168      # If already seen item we have no new bad items to find, finish up
169      if self.all_items[prune_index] in self.found_items:
170        self.l.LogOutput(('Found item already found before: %s. '
171                          'Done searching.' % self.all_items[prune_index]))
172        self.l.LogOutput('Bad items are: %s' %
173                         ' '.join(self.found_items))
174        break
175
176      new_all_items = list(self.all_items)
177      # Move prune item to the end of the list.
178      new_all_items.append(new_all_items.pop(prune_index))
179      self.found_items.add(new_all_items[-1])
180
181      if prune_index:
182        new_all_items = new_all_items[prune_index - 1:]
183
184      self.l.LogOutput('Old list: %s. New list: %s' % (str(self.all_items),
185                                                       str(new_all_items)),
186                       print_to_console=self.verbose)
187
188      # FIXME: Do we need to Convert the currently good items to bad
189      self.PopulateItemsUsingList(new_all_items)
190
191  def DoBinarySearch(self):
192    # If in resume mode don't reset search_cycles
193    if not self.resumed:
194      self.search_cycles = 0
195    else:
196      self.resumed = False
197
198    terminated = False
199    while self.search_cycles < self.iterations and not terminated:
200      self.SaveState()
201      self.OutputProgress()
202      self.search_cycles += 1
203      [bad_items, good_items] = self.GetNextItems()
204
205      # TODO: bad_items should come first.
206      self.SwitchToGood(good_items)
207      self.SwitchToBad(bad_items)
208      status = self.InstallScript()
209      if status == 0:
210        status = self.TestScript()
211      else:
212        # Install script failed, treat as skipped item
213        status = 2
214      terminated = self.binary_search.SetStatus(status)
215
216      if terminated:
217        self.l.LogOutput('Terminated!')
218    if not terminated:
219      self.l.LogOutput('Ran out of iterations searching...')
220    self.l.LogOutput(str(self), print_to_console=self.verbose)
221    return terminated
222
223  def PopulateItemsUsingCommand(self, command):
224    ce = command_executer.GetCommandExecuter()
225    _, out, _ = ce.RunCommandWExceptionCleanup(command,
226                                               return_output=True,
227                                               print_to_console=self.verbose)
228    all_items = out.split()
229    self.PopulateItemsUsingList(all_items)
230
231  def PopulateItemsUsingList(self, all_items):
232    self.all_items = all_items
233    self.binary_search = binary_search_perforce.BinarySearcher(
234        logger_to_set=self.l)
235    self.binary_search.SetSortedList(self.all_items)
236
237  def SaveState(self):
238    ce, l = self.ce, self.l
239    self.ce, self.l, self.binary_search.logger = None, None, None
240    old_state = None
241
242    _, path = tempfile.mkstemp(prefix=HIDDEN_STATE_FILE, dir='.')
243    with open(path, 'wb') as f:
244      pickle.dump(self, f)
245
246    if os.path.exists(STATE_FILE):
247      if os.path.islink(STATE_FILE):
248        old_state = os.readlink(STATE_FILE)
249      else:
250        raise Error(('%s already exists and is not a symlink!\n'
251                     'State file saved to %s' % (STATE_FILE, path)))
252
253    # Create new link and atomically overwrite old link
254    temp_link = '%s.link' % HIDDEN_STATE_FILE
255    os.symlink(path, temp_link)
256    os.rename(temp_link, STATE_FILE)
257
258    if old_state:
259      os.remove(old_state)
260
261    self.ce, self.l, self.binary_search.logger = ce, l, l
262
263  @classmethod
264  def LoadState(cls):
265    if not os.path.isfile(STATE_FILE):
266      return None
267    try:
268      bss = pickle.load(file(STATE_FILE))
269      bss.l = logger.GetLogger()
270      bss.ce = command_executer.GetCommandExecuter()
271      bss.binary_search.logger = bss.l
272      bss.resumed = True
273      binary_search_perforce.verbose = bss.verbose
274      return bss
275    except Exception:
276      return None
277
278  def RemoveState(self):
279    if os.path.exists(STATE_FILE):
280      if os.path.islink(STATE_FILE):
281        real_file = os.readlink(STATE_FILE)
282        os.remove(real_file)
283        os.remove(STATE_FILE)
284
285  def GetNextItems(self):
286    border_item = self.binary_search.GetNext()
287    index = self.all_items.index(border_item)
288
289    next_bad_items = self.all_items[:index + 1]
290    next_good_items = self.all_items[index + 1:]
291
292    return [next_bad_items, next_good_items]
293
294  def OutputProgress(self):
295    out = ('\n********** PROGRESS **********\n'
296           'Search %d of estimated %d.\n'
297           'Prune %d of max %d.\n'
298           'Current bad items found:\n'
299           '%s\n'
300           '******************************')
301    out = out % (self.search_cycles + 1,
302                 math.ceil(math.log(len(self.all_items), 2)),
303                 self.prune_cycles + 1,
304                 self.prune_iterations,
305                 str(self.found_items))
306
307    self.l.LogOutput(out)
308
309  def __str__(self):
310    ret = ''
311    ret += 'all: %s\n' % str(self.all_items)
312    ret += 'currently_good: %s\n' % str(self.currently_good_items)
313    ret += 'currently_bad: %s\n' % str(self.currently_bad_items)
314    ret += str(self.binary_search)
315    return ret
316
317
318class MockBinarySearchState(BinarySearchState):
319  """Mock class for BinarySearchState."""
320
321  def __init__(self, **kwargs):
322    # Initialize all arguments to None
323    default_kwargs = {
324        'get_initial_items': 'echo "1"',
325        'switch_to_good': None,
326        'switch_to_bad': None,
327        'install_script': None,
328        'test_script': None,
329        'incremental': None,
330        'prune': None,
331        'iterations': None,
332        'prune_iterations': None,
333        'verify_level': None,
334        'file_args': None,
335        'verbose': None
336    }
337    default_kwargs.update(kwargs)
338    super(MockBinarySearchState, self).__init__(**default_kwargs)
339
340
341def _CanonicalizeScript(script_name):
342  script_name = os.path.expanduser(script_name)
343  if not script_name.startswith('/'):
344    return os.path.join('.', script_name)
345
346def Run(get_initial_items, switch_to_good, switch_to_bad, test_script,
347        install_script=None, iterations=50, prune=True, noincremental=False,
348        file_args=False, verify_level=1, prune_iterations=100, verbose=False,
349        resume=False):
350  """Run binary search tool. Equivalent to running through terminal."""
351  if resume:
352    bss = BinarySearchState.LoadState()
353    if not bss:
354      logger.GetLogger().LogOutput(
355          '%s is not a valid binary_search_tool state file, cannot resume!' %
356          STATE_FILE)
357      return 1
358  else:
359    switch_to_good = _CanonicalizeScript(switch_to_good)
360    switch_to_bad = _CanonicalizeScript(switch_to_bad)
361    if install_script:
362      install_script = _CanonicalizeScript(install_script)
363    test_script = _CanonicalizeScript(test_script)
364    get_initial_items = _CanonicalizeScript(get_initial_items)
365    incremental = not noincremental
366
367    binary_search_perforce.verbose = verbose
368
369    bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad,
370                            install_script, test_script, incremental, prune,
371                            iterations, prune_iterations, verify_level,
372                            file_args, verbose)
373    bss.DoVerify()
374
375  try:
376    bss.DoSearch()
377    bss.RemoveState()
378  except Error as e:
379    logger.GetLogger().LogError(e)
380    return 1
381
382  return 0
383
384
385def Main(argv):
386  """The main function."""
387  # Common initializations
388
389  parser = argparse.ArgumentParser()
390  parser.add_argument('-n',
391                      '--iterations',
392                      dest='iterations',
393                      type=int,
394                      help='Number of iterations to try in the search.',
395                      default=50)
396  parser.add_argument('-i',
397                      '--get_initial_items',
398                      dest='get_initial_items',
399                      help=('Script to run to get the initial objects. '
400                            'If your script requires user input '
401                            'the --verbose option must be used'))
402  parser.add_argument('-g',
403                      '--switch_to_good',
404                      dest='switch_to_good',
405                      help=('Script to run to switch to good. '
406                            'If your switch script requires user input '
407                            'the --verbose option must be used'))
408  parser.add_argument('-b',
409                      '--switch_to_bad',
410                      dest='switch_to_bad',
411                      help=('Script to run to switch to bad. '
412                            'If your switch script requires user input '
413                            'the --verbose option must be used'))
414  parser.add_argument('-I',
415                      '--install_script',
416                      dest='install_script',
417                      default=None,
418                      help=('Optional script to perform building, flashing, '
419                            'and other setup before the test script runs.'))
420  parser.add_argument('-t',
421                      '--test_script',
422                      dest='test_script',
423                      help=('Script to run to test the '
424                            'output after packages are built.'))
425  parser.add_argument('-p',
426                      '--prune',
427                      dest='prune',
428                      action='store_true',
429                      default=False,
430                      help=('Script to run to test the output after '
431                            'packages are built.'))
432  parser.add_argument('-c',
433                      '--noincremental',
434                      dest='noincremental',
435                      action='store_true',
436                      default=False,
437                      help='Do not propagate good/bad changes incrementally.')
438  parser.add_argument('-f',
439                      '--file_args',
440                      dest='file_args',
441                      action='store_true',
442                      default=False,
443                      help='Use a file to pass arguments to scripts.')
444  parser.add_argument('-v',
445                      '--verify_level',
446                      dest='verify_level',
447                      type=int,
448                      default=1,
449                      help=('Check binary search assumptions N times '
450                            'before starting.'))
451  parser.add_argument('-N',
452                      '--prune_iterations',
453                      dest='prune_iterations',
454                      type=int,
455                      help='Number of prune iterations to try in the search.',
456                      default=100)
457  parser.add_argument('-V',
458                      '--verbose',
459                      dest='verbose',
460                      action='store_true',
461                      help='Print full output to console.')
462  parser.add_argument('-r',
463                      '--resume',
464                      dest='resume',
465                      action='store_true',
466                      help=('Resume bisection tool execution from state file.'
467                            'Useful if the last bisection was terminated '
468                            'before it could properly finish.'))
469
470  logger.GetLogger().LogOutput(' '.join(argv))
471  options = parser.parse_args(argv)
472
473  if not (options.get_initial_items and options.switch_to_good and
474          options.switch_to_bad and options.test_script) and not options.resume:
475    parser.print_help()
476    return 1
477
478  if options.resume:
479    logger.GetLogger().LogOutput('Resuming from %s' % STATE_FILE)
480    if len(argv) > 1:
481      logger.GetLogger().LogOutput(('Note: resuming from previous state, '
482                                    'ignoring given options and loading saved '
483                                    'options instead.'))
484
485  # Get dictionary of all options
486  args = vars(options)
487  return Run(**args)
488
489
490if __name__ == '__main__':
491  sys.exit(Main(sys.argv[1:]))
492