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