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