binary_search_state.py revision 333fe383dcd0c35981a5473bc913d07b52c310a1 
1#!/usr/bin/python2
2"""The binary search wrapper."""
3
4from __future__ import print_function
5
6import argparse
7import os
8import pickle
9import sys
10import tempfile
11
12# Programtically adding utils python path to PYTHONPATH
13if os.path.isabs(sys.argv[0]):
14  utils_pythonpath = os.path.abspath('{0}/..'.format(os.path.dirname(sys.argv[
15      0])))
16else:
17  wdir = os.getcwd()
18  utils_pythonpath = os.path.abspath('{0}/{1}/..'.format(wdir, os.path.dirname(
19      sys.argv[0])))
20sys.path.append(utils_pythonpath)
21# Now we do import from utils
22from utils import command_executer
23from utils import logger
24
25import binary_search_perforce
26
27STATE_FILE = '%s.state' % sys.argv[0]
28
29
30class BinarySearchState(object):
31  """The binary search state class."""
32
33  def __init__(self, get_initial_items, switch_to_good, switch_to_bad,
34               install_script, test_script, incremental, prune, iterations,
35               prune_iterations, verify_level, file_args, verbose):
36    self.get_initial_items = get_initial_items
37    self.switch_to_good = switch_to_good
38    self.switch_to_bad = switch_to_bad
39    self.install_script = install_script
40    self.test_script = test_script
41    self.incremental = incremental
42    self.prune = prune
43    self.iterations = iterations
44    self.prune_iterations = prune_iterations
45    self.verify_level = verify_level
46    self.file_args = file_args
47    self.verbose = verbose
48
49    self.l = logger.GetLogger()
50    self.ce = command_executer.GetCommandExecuter()
51
52    self.bs = None
53    self.all_items = None
54    self.PopulateItemsUsingCommand(self.get_initial_items)
55    self.currently_good_items = set([])
56    self.currently_bad_items = set([])
57
58  def SwitchToGood(self, item_list):
59    if self.incremental:
60      self.l.LogOutput('Incremental set. Wanted to switch %s to good' %
61                       str(item_list), print_to_console=self.verbose)
62      incremental_items = [
63          item for item in item_list if item not in self.currently_good_items
64      ]
65      item_list = incremental_items
66      self.l.LogOutput('Incremental set. Actually switching %s to good' %
67                       str(item_list), print_to_console=self.verbose)
68
69    if not item_list:
70      return
71
72    self.l.LogOutput('Switching %s to good' % str(item_list),
73                     print_to_console=self.verbose)
74    self.RunSwitchScript(self.switch_to_good, item_list)
75    self.currently_good_items = self.currently_good_items.union(set(item_list))
76    self.currently_bad_items.difference_update(set(item_list))
77
78  def SwitchToBad(self, item_list):
79    if self.incremental:
80      self.l.LogOutput('Incremental set. Wanted to switch %s to bad' %
81                       str(item_list), print_to_console=self.verbose)
82      incremental_items = [
83          item for item in item_list if item not in self.currently_bad_items
84      ]
85      item_list = incremental_items
86      self.l.LogOutput('Incremental set. Actually switching %s to bad' %
87                       str(item_list), print_to_console=self.verbose)
88
89    if not item_list:
90      return
91
92    self.l.LogOutput('Switching %s to bad' % str(item_list),
93                     print_to_console=self.verbose)
94    self.RunSwitchScript(self.switch_to_bad, item_list)
95    self.currently_bad_items = self.currently_bad_items.union(set(item_list))
96    self.currently_good_items.difference_update(set(item_list))
97
98  def RunSwitchScript(self, switch_script, item_list):
99    if self.file_args:
100      temp_file = tempfile.mktemp()
101      f = open(temp_file, 'wb')
102      f.write('\n'.join(item_list))
103      f.close()
104      command = '%s %s' % (switch_script, temp_file)
105    else:
106      command = '%s %s' % (switch_script, ' '.join(item_list))
107    ret, _, _ = self.ce.RunCommandWExceptionCleanup(
108        command, print_to_console=self.verbose)
109    assert ret == 0, 'Switch script %s returned %d' % (switch_script, ret)
110
111  def TestScript(self):
112    command = self.test_script
113    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
114    return ret
115
116  def InstallScript(self):
117    if not self.install_script:
118      return 0
119
120    command = self.install_script
121    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
122    return ret
123
124  def DoVerify(self):
125    for _ in range(int(self.verify_level)):
126      self.l.LogOutput('Resetting all items to good to verify.')
127      self.SwitchToGood(self.all_items)
128      status = self.InstallScript()
129      assert status == 0, 'When reset_to_good, install should succeed.'
130      status = self.TestScript()
131      assert status == 0, 'When reset_to_good, status should be 0.'
132
133      self.l.LogOutput('Resetting all items to bad to verify.')
134      self.SwitchToBad(self.all_items)
135      status = self.InstallScript()
136      assert status == 0, 'When reset_to_bad, install should succeed.'
137      status = self.TestScript()
138      assert status == 1, 'When reset_to_bad, status should be 1.'
139
140  def DoSearch(self):
141    num_bad_items_history = []
142    i = 0
143    while True and len(self.all_items) > 1 and i < self.prune_iterations:
144      i += 1
145      terminated = self.DoBinarySearch()
146      if not terminated:
147        break
148      if not self.prune:
149        self.l.LogOutput('Not continuning further, --prune is not set')
150        break
151      # Prune is set.
152      prune_index = self.bs.current
153
154      if prune_index == len(self.all_items) - 1:
155        self.l.LogOutput('First bad item is the last item. Breaking.')
156        self.l.LogOutput('Only bad item is: %s' % self.all_items[-1])
157        break
158
159      num_bad_items = len(self.all_items) - prune_index
160      num_bad_items_history.append(num_bad_items)
161
162      if (num_bad_items_history[-num_bad_items:] ==
163          [num_bad_items for _ in range(num_bad_items)]):
164        self.l.LogOutput('num_bad_items_history: %s for past %d iterations. '
165                         'Breaking.' % (str(num_bad_items_history),
166                                        num_bad_items))
167        self.l.LogOutput('Bad items are: %s' %
168                         ' '.join(self.all_items[prune_index:]))
169        break
170
171      new_all_items = list(self.all_items)
172      # Move prune item to the end of the list.
173      new_all_items.append(new_all_items.pop(prune_index))
174
175      if prune_index:
176        new_all_items = new_all_items[prune_index - 1:]
177
178      self.l.LogOutput('Old list: %s. New list: %s' % (str(self.all_items),
179                                                       str(new_all_items)),
180                       print_to_console=self.verbose)
181
182      # FIXME: Do we need to Convert the currently good items to bad
183      self.PopulateItemsUsingList(new_all_items)
184
185  def DoBinarySearch(self):
186    i = 0
187    terminated = False
188    while i < self.iterations and not terminated:
189      i += 1
190      [bad_items, good_items] = self.GetNextItems()
191
192      # TODO: bad_items should come first.
193      self.SwitchToGood(good_items)
194      self.SwitchToBad(bad_items)
195      status = self.InstallScript()
196      if status == 0:
197        status = self.TestScript()
198      else:
199        # Install script failed, treat as skipped item
200        status = 2
201      terminated = self.bs.SetStatus(status)
202
203      if terminated:
204        self.l.LogOutput('Terminated!')
205    if not terminated:
206      self.l.LogOutput('Ran out of iterations searching...')
207    self.l.LogOutput(str(self), print_to_console=self.verbose)
208    return terminated
209
210  def PopulateItemsUsingCommand(self, command):
211    ce = command_executer.GetCommandExecuter()
212    _, out, _ = ce.RunCommandWExceptionCleanup(command,
213                                               return_output=True,
214                                               print_to_console=self.verbose)
215    all_items = out.split()
216    self.PopulateItemsUsingList(all_items)
217
218  def PopulateItemsUsingList(self, all_items):
219    self.all_items = all_items
220    self.bs = binary_search_perforce.BinarySearcher()
221    self.bs.SetSortedList(self.all_items)
222
223  def SaveState(self):
224    self.l = None
225    self.ce = None
226    # TODO Implement save/restore
227    ###    return
228    f = open(STATE_FILE, 'wb')
229    pickle.dump(self, f)
230    f.close()
231
232  @classmethod
233  def LoadState(cls):
234    if not os.path.isfile(STATE_FILE):
235      return None
236    return pickle.load(file(STATE_FILE))
237
238  def GetNextItems(self):
239    border_item = self.bs.GetNext()
240    index = self.all_items.index(border_item)
241
242    next_bad_items = self.all_items[:index + 1]
243    next_good_items = self.all_items[index + 1:]
244
245    return [next_bad_items, next_good_items]
246
247  def __str__(self):
248    ret = ''
249    ret += 'all: %s\n' % str(self.all_items)
250    ret += 'currently_good: %s\n' % str(self.currently_good_items)
251    ret += 'currently_bad: %s\n' % str(self.currently_bad_items)
252    ret += str(self.bs)
253    return ret
254
255
256def _CanonicalizeScript(script_name):
257  script_name = os.path.expanduser(script_name)
258  if not script_name.startswith('/'):
259    return os.path.join('.', script_name)
260
261
262def Main(argv):
263  """The main function."""
264  # Common initializations
265
266  parser = argparse.ArgumentParser()
267  parser.add_argument('-n',
268                      '--iterations',
269                      dest='iterations',
270                      help='Number of iterations to try in the search.',
271                      default=50)
272  parser.add_argument('-i',
273                      '--get_initial_items',
274                      dest='get_initial_items',
275                      help=('Script to run to get the initial objects. '
276                            'If your script requires user input '
277                            'the --verbose option must be used'))
278  parser.add_argument('-g',
279                      '--switch_to_good',
280                      dest='switch_to_good',
281                      help=('Script to run to switch to good. '
282                            'If your switch script requires user input '
283                            'the --verbose option must be used'))
284  parser.add_argument('-b',
285                      '--switch_to_bad',
286                      dest='switch_to_bad',
287                      help=('Script to run to switch to bad. '
288                            'If your switch script requires user input '
289                            'the --verbose option must be used'))
290  parser.add_argument('-I',
291                      '--install_script',
292                      dest='install_script',
293                      default=None,
294                      help=('Optional script to perform building, flashing, '
295                            'and other setup before the test script runs.'))
296  parser.add_argument('-t',
297                      '--test_script',
298                      dest='test_script',
299                      help=('Script to run to test the '
300                            'output after packages are built.'))
301  parser.add_argument('-p',
302                      '--prune',
303                      dest='prune',
304                      action='store_true',
305                      default=False,
306                      help=('Script to run to test the output after '
307                            'packages are built.'))
308  parser.add_argument('-c',
309                      '--noincremental',
310                      dest='noincremental',
311                      action='store_true',
312                      default=False,
313                      help='Do not propagate good/bad changes incrementally.')
314  parser.add_argument('-f',
315                      '--file_args',
316                      dest='file_args',
317                      action='store_true',
318                      default=False,
319                      help='Use a file to pass arguments to scripts.')
320  parser.add_argument('-v',
321                      '--verify_level',
322                      dest='verify_level',
323                      default=1,
324                      help=('Check binary search assumptions N times '
325                            'before starting.'))
326  parser.add_argument('-N',
327                      '--prune_iterations',
328                      dest='prune_iterations',
329                      help='Number of prune iterations to try in the search.',
330                      default=100)
331  parser.add_argument('-V',
332                      '--verbose',
333                      dest='verbose',
334                      action='store_true',
335                      help='Print full output to console.')
336
337  logger.GetLogger().LogOutput(' '.join(argv))
338  options = parser.parse_args(argv)
339
340  if not (options.get_initial_items and options.switch_to_good and
341          options.switch_to_bad and options.test_script):
342    parser.print_help()
343    return 1
344
345  iterations = int(options.iterations)
346  switch_to_good = _CanonicalizeScript(options.switch_to_good)
347  switch_to_bad = _CanonicalizeScript(options.switch_to_bad)
348  install_script = options.install_script
349  if install_script:
350    install_script = _CanonicalizeScript(options.install_script)
351  test_script = _CanonicalizeScript(options.test_script)
352  get_initial_items = _CanonicalizeScript(options.get_initial_items)
353  prune = options.prune
354  prune_iterations = options.prune_iterations
355  verify_level = options.verify_level
356  file_args = options.file_args
357  verbose = options.verbose
358  binary_search_perforce.verbose = verbose
359
360  if options.noincremental:
361    incremental = False
362  else:
363    incremental = True
364
365  try:
366    bss = BinarySearchState(get_initial_items, switch_to_good, switch_to_bad,
367                            install_script, test_script, incremental, prune,
368                            iterations, prune_iterations, verify_level,
369                            file_args, verbose)
370    bss.DoVerify()
371    bss.DoSearch()
372
373  except (KeyboardInterrupt, SystemExit):
374    print('C-c pressed')
375    bss.SaveState()
376  return 0
377
378
379if __name__ == '__main__':
380  sys.exit(Main(sys.argv[1:]))
381