binary_search_state.py revision 702bf8793c082c144b57e445694ccc548a90ac21 
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
346
347def Main(argv):
348  """The main function."""
349  # Common initializations
350
351  parser = argparse.ArgumentParser()
352  parser.add_argument('-n',
353                      '--iterations',
354                      dest='iterations',
355                      help='Number of iterations to try in the search.',
356                      default=50)
357  parser.add_argument('-i',
358                      '--get_initial_items',
359                      dest='get_initial_items',
360                      help=('Script to run to get the initial objects. '
361                            'If your script requires user input '
362                            'the --verbose option must be used'))
363  parser.add_argument('-g',
364                      '--switch_to_good',
365                      dest='switch_to_good',
366                      help=('Script to run to switch to good. '
367                            'If your switch script requires user input '
368                            'the --verbose option must be used'))
369  parser.add_argument('-b',
370                      '--switch_to_bad',
371                      dest='switch_to_bad',
372                      help=('Script to run to switch to bad. '
373                            'If your switch script requires user input '
374                            'the --verbose option must be used'))
375  parser.add_argument('-I',
376                      '--install_script',
377                      dest='install_script',
378                      default=None,
379                      help=('Optional script to perform building, flashing, '
380                            'and other setup before the test script runs.'))
381  parser.add_argument('-t',
382                      '--test_script',
383                      dest='test_script',
384                      help=('Script to run to test the '
385                            'output after packages are built.'))
386  parser.add_argument('-p',
387                      '--prune',
388                      dest='prune',
389                      action='store_true',
390                      default=False,
391                      help=('Script to run to test the output after '
392                            'packages are built.'))
393  parser.add_argument('-c',
394                      '--noincremental',
395                      dest='noincremental',
396                      action='store_true',
397                      default=False,
398                      help='Do not propagate good/bad changes incrementally.')
399  parser.add_argument('-f',
400                      '--file_args',
401                      dest='file_args',
402                      action='store_true',
403                      default=False,
404                      help='Use a file to pass arguments to scripts.')
405  parser.add_argument('-v',
406                      '--verify_level',
407                      dest='verify_level',
408                      default=1,
409                      help=('Check binary search assumptions N times '
410                            'before starting.'))
411  parser.add_argument('-N',
412                      '--prune_iterations',
413                      dest='prune_iterations',
414                      help='Number of prune iterations to try in the search.',
415                      default=100)
416  parser.add_argument('-V',
417                      '--verbose',
418                      dest='verbose',
419                      action='store_true',
420                      help='Print full output to console.')
421  parser.add_argument('-r',
422                      '--resume',
423                      dest='resume',
424                      action='store_true',
425                      help=('Resume bisection tool execution from state file.'
426                            'Useful if the last bisection was terminated '
427                            'before it could properly finish.'))
428
429  logger.GetLogger().LogOutput(' '.join(argv))
430  options = parser.parse_args(argv)
431
432  if not (options.get_initial_items and options.switch_to_good and
433          options.switch_to_bad and options.test_script) and not options.resume:
434    parser.print_help()
435    return 1
436
437  if options.resume:
438    logger.GetLogger().LogOutput('Resuming from %s' % STATE_FILE)
439    if len(argv) > 1:
440      logger.GetLogger().LogOutput(('Note: resuming from previous state, '
441                                    'ignoring given options and loading saved '
442                                    'options instead.'))
443    bss = BinarySearchState.LoadState()
444    if not bss:
445      logger.GetLogger().LogOutput(
446          '%s is not a valid binary_search_tool state file, cannot resume!' %
447          STATE_FILE)
448      return 1
449  else:
450    iterations = int(options.iterations)
451    switch_to_good = _CanonicalizeScript(options.switch_to_good)
452    switch_to_bad = _CanonicalizeScript(options.switch_to_bad)
453    install_script = options.install_script
454    if install_script:
455      install_script = _CanonicalizeScript(options.install_script)
456    test_script = _CanonicalizeScript(options.test_script)
457    get_initial_items = _CanonicalizeScript(options.get_initial_items)
458    prune = options.prune
459    prune_iterations = options.prune_iterations
460    verify_level = options.verify_level
461    file_args = options.file_args
462    verbose = options.verbose
463    binary_search_perforce.verbose = verbose
464
465    if options.noincremental:
466      incremental = False
467    else:
468      incremental = True
469
470    bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad,
471                            install_script, test_script, incremental, prune,
472                            iterations, prune_iterations, verify_level,
473                            file_args, verbose)
474    bss.DoVerify()
475
476  try:
477    bss.DoSearch()
478    bss.RemoveState()
479  except Error as e:
480    logger.GetLogger().LogError(e)
481    return 1
482
483  return 0
484
485
486if __name__ == '__main__':
487  sys.exit(Main(sys.argv[1:]))
488