binary_search_state.py revision 19b6f5fc11dcf97144e9723c8f78534cce27423a 
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
12import time
13
14# Adds utils to PYTHONPATH
15import common
16
17# Now we do import from utils
18from utils import command_executer
19from utils import logger
20
21import binary_search_perforce
22
23STATE_FILE = '%s.state' % sys.argv[0]
24HIDDEN_STATE_FILE = os.path.join(
25    os.path.dirname(STATE_FILE), '.%s' % os.path.basename(STATE_FILE))
26
27class Error(Exception):
28  """The general binary search tool error class."""
29  pass
30
31
32class BinarySearchState(object):
33  """The binary search state class."""
34
35  def __init__(self, get_initial_items, switch_to_good, switch_to_bad,
36               install_script, test_script, incremental, prune, iterations,
37               prune_iterations, verify_level, file_args, verbose):
38    """BinarySearchState constructor, see Run for full args documentation."""
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.resumed = False
56    self.prune_cycles = 0
57    self.search_cycles = 0
58    self.binary_search = None
59    self.all_items = None
60    self.PopulateItemsUsingCommand(self.get_initial_items)
61    self.currently_good_items = set([])
62    self.currently_bad_items = set([])
63    self.found_items = set([])
64
65    self.start_time = time.time()
66
67  def SwitchToGood(self, item_list):
68    """Switch given items to "good" set."""
69    if self.incremental:
70      self.l.LogOutput('Incremental set. Wanted to switch %s to good' %
71                       str(item_list), print_to_console=self.verbose)
72      incremental_items = [
73          item for item in item_list if item not in self.currently_good_items
74      ]
75      item_list = incremental_items
76      self.l.LogOutput('Incremental set. Actually switching %s to good' %
77                       str(item_list), print_to_console=self.verbose)
78
79    if not item_list:
80      return
81
82    self.l.LogOutput('Switching %s to good' % str(item_list),
83                     print_to_console=self.verbose)
84    self.RunSwitchScript(self.switch_to_good, item_list)
85    self.currently_good_items = self.currently_good_items.union(set(item_list))
86    self.currently_bad_items.difference_update(set(item_list))
87
88  def SwitchToBad(self, item_list):
89    """Switch given items to "bad" set."""
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    """Pass given items to switch script.
111
112    Args:
113      switch_script: path to switch script
114      item_list: list of all items to be switched
115    """
116    if self.file_args:
117      with tempfile.NamedTemporaryFile() as f:
118        f.write('\n'.join(item_list))
119        f.flush()
120        command = '%s %s' % (switch_script, f.name)
121        ret, _, _ = self.ce.RunCommandWExceptionCleanup(
122            command, print_to_console=self.verbose)
123    else:
124      command = '%s %s' % (switch_script, ' '.join(item_list))
125      ret, _, _ = self.ce.RunCommandWExceptionCleanup(
126          command, print_to_console=self.verbose)
127    assert ret == 0, 'Switch script %s returned %d' % (switch_script, ret)
128
129  def TestScript(self):
130    """Run test script and return exit code from script."""
131    command = self.test_script
132    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
133    return ret
134
135  def InstallScript(self):
136    """Run install script and return exit code from script."""
137    if not self.install_script:
138      return 0
139
140    command = self.install_script
141    ret, _, _ = self.ce.RunCommandWExceptionCleanup(command)
142    return ret
143
144  def DoVerify(self):
145    """Verify correctness of test environment.
146
147    Verify that a "good" set of items produces a "good" result and that a "bad"
148    set of items produces a "bad" result. To be run directly before running
149    DoSearch. If verify_level is 0 this step is skipped.
150    """
151    if not self.verify_level:
152      return
153
154    self.l.LogOutput('VERIFICATION')
155    self.l.LogOutput('Beginning %d tests to verify good/bad sets\n' %
156                     self.verify_level)
157    for _ in range(int(self.verify_level)):
158      self.l.LogOutput('Resetting all items to good to verify.')
159      self.SwitchToGood(self.all_items)
160      status = self.InstallScript()
161      assert status == 0, 'When reset_to_good, install should succeed.'
162      status = self.TestScript()
163      assert status == 0, 'When reset_to_good, status should be 0.'
164
165      self.l.LogOutput('Resetting all items to bad to verify.')
166      self.SwitchToBad(self.all_items)
167      status = self.InstallScript()
168      assert status == 0, 'When reset_to_bad, install should succeed.'
169      status = self.TestScript()
170      assert status == 1, 'When reset_to_bad, status should be 1.'
171
172  def DoSearch(self):
173    """Perform full search for bad items.
174
175    Perform full search until prune_iterations number of bad items are found.
176    """
177    while (True and
178           len(self.all_items) > 1 and
179           self.prune_cycles < self.prune_iterations):
180      terminated = self.DoBinarySearch()
181      self.prune_cycles += 1
182      if not terminated:
183        break
184      # Prune is set.
185      prune_index = self.binary_search.current
186
187      # If already seen item we have no new bad items to find, finish up
188      if self.all_items[prune_index] in self.found_items:
189        self.l.LogOutput(('Found item already found before: %s. '
190                          'Done searching.' % self.all_items[prune_index]))
191        self.l.LogOutput('Bad items are: %s' %
192                         ' '.join(self.found_items))
193        break
194
195      new_all_items = list(self.all_items)
196      # Move prune item to the end of the list.
197      new_all_items.append(new_all_items.pop(prune_index))
198      self.found_items.add(new_all_items[-1])
199
200      if prune_index:
201        new_all_items = new_all_items[prune_index - 1:]
202
203      self.l.LogOutput('Old list: %s. New list: %s' % (str(self.all_items),
204                                                       str(new_all_items)),
205                       print_to_console=self.verbose)
206
207      if not self.prune:
208        self.l.LogOutput('Not continuning further, --prune is not set')
209        break
210      # FIXME: Do we need to Convert the currently good items to bad
211      self.PopulateItemsUsingList(new_all_items)
212
213  def DoBinarySearch(self):
214    """Perform single iteration of binary search."""
215    # If in resume mode don't reset search_cycles
216    if not self.resumed:
217      self.search_cycles = 0
218    else:
219      self.resumed = False
220
221    terminated = False
222    while self.search_cycles < self.iterations and not terminated:
223      self.SaveState()
224      self.OutputProgress()
225
226      self.search_cycles += 1
227      [bad_items, good_items] = self.GetNextItems()
228
229      # TODO: bad_items should come first.
230      self.SwitchToGood(good_items)
231      self.SwitchToBad(bad_items)
232      status = self.InstallScript()
233      if status == 0:
234        status = self.TestScript()
235      else:
236        # Install script failed, treat as skipped item
237        status = 2
238      terminated = self.binary_search.SetStatus(status)
239
240      if terminated:
241        self.l.LogOutput('Terminated!')
242    if not terminated:
243      self.l.LogOutput('Ran out of iterations searching...')
244    self.l.LogOutput(str(self), print_to_console=self.verbose)
245    return terminated
246
247  def PopulateItemsUsingCommand(self, command):
248    """Update all_items and binary search logic from executable.
249
250    This method is mainly required for enumerating the initial list of items
251    from the get_initial_items script.
252
253    Args:
254      command: path to executable that will enumerate items.
255    """
256    ce = command_executer.GetCommandExecuter()
257    _, out, _ = ce.RunCommandWExceptionCleanup(command,
258                                               return_output=True,
259                                               print_to_console=self.verbose)
260    all_items = out.split()
261    self.PopulateItemsUsingList(all_items)
262
263  def PopulateItemsUsingList(self, all_items):
264    """Update all_items and binary searching logic from list.
265
266    Args:
267      all_items: new list of all_items
268    """
269    self.all_items = all_items
270    self.binary_search = binary_search_perforce.BinarySearcher(
271        logger_to_set=self.l)
272    self.binary_search.SetSortedList(self.all_items)
273
274  def SaveState(self):
275    """Save state to STATE_FILE.
276
277    SaveState will create a new unique, hidden state file to hold data from
278    object. Then atomically overwrite the STATE_FILE symlink to point to the
279    new data.
280
281    Raises:
282      Error if STATE_FILE already exists but is not a symlink.
283    """
284    ce, l = self.ce, self.l
285    self.ce, self.l, self.binary_search.logger = None, None, None
286    old_state = None
287
288    _, path = tempfile.mkstemp(prefix=HIDDEN_STATE_FILE, dir='.')
289    with open(path, 'wb') as f:
290      pickle.dump(self, f)
291
292    if os.path.exists(STATE_FILE):
293      if os.path.islink(STATE_FILE):
294        old_state = os.readlink(STATE_FILE)
295      else:
296        raise Error(('%s already exists and is not a symlink!\n'
297                     'State file saved to %s' % (STATE_FILE, path)))
298
299    # Create new link and atomically overwrite old link
300    temp_link = '%s.link' % HIDDEN_STATE_FILE
301    os.symlink(path, temp_link)
302    os.rename(temp_link, STATE_FILE)
303
304    if old_state:
305      os.remove(old_state)
306
307    self.ce, self.l, self.binary_search.logger = ce, l, l
308
309  @classmethod
310  def LoadState(cls):
311    """Create BinarySearchState object from STATE_FILE."""
312    if not os.path.isfile(STATE_FILE):
313      return None
314    try:
315      bss = pickle.load(file(STATE_FILE))
316      bss.l = logger.GetLogger()
317      bss.ce = command_executer.GetCommandExecuter()
318      bss.binary_search.logger = bss.l
319      bss.start_time = time.time()
320      bss.resumed = True
321      binary_search_perforce.verbose = bss.verbose
322      return bss
323    except Exception:
324      return None
325
326  def RemoveState(self):
327    """Remove STATE_FILE and its symlinked data from file system."""
328    if os.path.exists(STATE_FILE):
329      if os.path.islink(STATE_FILE):
330        real_file = os.readlink(STATE_FILE)
331        os.remove(real_file)
332        os.remove(STATE_FILE)
333
334  def GetNextItems(self):
335    """Get next items for binary search based on result of the last test run."""
336    border_item = self.binary_search.GetNext()
337    index = self.all_items.index(border_item)
338
339    next_bad_items = self.all_items[:index + 1]
340    next_good_items = self.all_items[index + 1:]
341
342    return [next_bad_items, next_good_items]
343
344  def ElapsedTimeString(self):
345    """Return h m s format of elapsed time since execution has started."""
346    diff = int(time.time() - self.start_time)
347    seconds = diff % 60
348    minutes = (diff / 60) % 60
349    hours = diff / (60 * 60)
350
351    seconds = str(seconds).rjust(2)
352    minutes = str(minutes).rjust(2)
353    hours = str(hours).rjust(2)
354
355    return '%sh %sm %ss' % (hours, minutes, seconds)
356
357  def OutputProgress(self):
358    """Output current progress of binary search to console and logs."""
359    out = ('\n***** PROGRESS (elapsed time: %s) *****\n'
360           'Search %d of estimated %d.\n'
361           'Prune %d of max %d.\n'
362           'Current bad items found:\n'
363           '%s\n'
364           '************************************************')
365    out = out % (self.ElapsedTimeString(),
366                 self.search_cycles + 1,
367                 math.ceil(math.log(len(self.all_items), 2)),
368                 self.prune_cycles + 1,
369                 self.prune_iterations,
370                 str(self.found_items))
371
372    self.l.LogOutput(out)
373
374  def __str__(self):
375    ret = ''
376    ret += 'all: %s\n' % str(self.all_items)
377    ret += 'currently_good: %s\n' % str(self.currently_good_items)
378    ret += 'currently_bad: %s\n' % str(self.currently_bad_items)
379    ret += str(self.binary_search)
380    return ret
381
382
383class MockBinarySearchState(BinarySearchState):
384  """Mock class for BinarySearchState."""
385
386  def __init__(self, **kwargs):
387    # Initialize all arguments to None
388    default_kwargs = {
389        'get_initial_items': 'echo "1"',
390        'switch_to_good': None,
391        'switch_to_bad': None,
392        'install_script': None,
393        'test_script': None,
394        'incremental': True,
395        'prune': False,
396        'iterations': 50,
397        'prune_iterations': 100,
398        'verify_level': 1,
399        'file_args': False,
400        'verbose': False
401    }
402    default_kwargs.update(kwargs)
403    super(MockBinarySearchState, self).__init__(**default_kwargs)
404
405
406def _CanonicalizeScript(script_name):
407  """Return canonical path to script.
408
409  Args:
410    script_name: Relative or absolute path to script
411
412  Returns:
413    Canonicalized script path
414  """
415  script_name = os.path.expanduser(script_name)
416  if not script_name.startswith('/'):
417    return os.path.join('.', script_name)
418
419
420def Run(get_initial_items, switch_to_good, switch_to_bad, test_script,
421        install_script=None, iterations=50, prune=True, noincremental=False,
422        file_args=False, verify_level=1, prune_iterations=100, verbose=False,
423        resume=False):
424  """Run binary search tool. Equivalent to running through terminal.
425
426  Args:
427    get_initial_items: Script to enumerate all items being binary searched
428    switch_to_good: Script that will take items as input and switch them to good
429                    set
430    switch_to_bad: Script that will take items as input and switch them to bad
431                   set
432    test_script: Script that will determine if the current combination of good
433                 and bad items make a "good" or "bad" result.
434    install_script: Script to do necessary setup (building, compilation, etc.)
435                    for test_script
436    iterations: How many binary search iterations to run before exiting.
437    prune: If False the binary search tool will stop when the first bad item is
438           found. Otherwise then binary search tool will continue searching
439           until all bad items are found (or prune_iterations is reached).
440    noincremental: Whether to send "diffs" of good/bad items to switch scripts.
441    file_args: If True then arguments to switch scripts will be a file name
442               containing a newline separated list of the items to switch.
443    verify_level: How many verification tests to run to ensure initial good/bad
444                  sets actually produce a good/bad result.
445    prune_iterations: Max number of bad items to search for.
446    verbose: If True will print extra debug information to user.
447    resume: If True will resume using STATE_FILE.
448
449  Returns:
450    0 for success, error otherwise
451  """
452  if resume:
453    bss = BinarySearchState.LoadState()
454    if not bss:
455      logger.GetLogger().LogOutput(
456          '%s is not a valid binary_search_tool state file, cannot resume!' %
457          STATE_FILE)
458      return 1
459  else:
460    switch_to_good = _CanonicalizeScript(switch_to_good)
461    switch_to_bad = _CanonicalizeScript(switch_to_bad)
462    if install_script:
463      install_script = _CanonicalizeScript(install_script)
464    test_script = _CanonicalizeScript(test_script)
465    get_initial_items = _CanonicalizeScript(get_initial_items)
466    incremental = not noincremental
467
468    binary_search_perforce.verbose = verbose
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    logger.GetLogger().LogOutput('Total execution time: %s' %
480                                 bss.ElapsedTimeString())
481  except Error as e:
482    logger.GetLogger().LogError(e)
483    return 1
484
485  return 0
486
487
488def Main(argv):
489  """The main function."""
490  # Common initializations
491
492  parser = argparse.ArgumentParser()
493  common.BuildArgParser(parser)
494  logger.GetLogger().LogOutput(' '.join(argv))
495  options = parser.parse_args(argv)
496
497  if not (options.get_initial_items and options.switch_to_good and
498          options.switch_to_bad and options.test_script) and not options.resume:
499    parser.print_help()
500    return 1
501
502  if options.resume:
503    logger.GetLogger().LogOutput('Resuming from %s' % STATE_FILE)
504    if len(argv) > 1:
505      logger.GetLogger().LogOutput(('Note: resuming from previous state, '
506                                    'ignoring given options and loading saved '
507                                    'options instead.'))
508
509  # Get dictionary of all options
510  args = vars(options)
511  return Run(**args)
512
513
514if __name__ == '__main__':
515  sys.exit(Main(sys.argv[1:]))
516