Fix bug where multiprocessing.Queue hangs on large outputs.

au_test_harness/cros_au_test_harness.py

 #!/usr/bin/python
 
 # Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """This module runs a suite of Auto Update tests.
 
   The tests can be run on either a virtual machine or actual device depending
   on parameters given.  Specific tests can be run by invoking --test_prefix.
@@ skipping 11 matching lines @@
 import constants
 sys.path.append(constants.CROSUTILS_LIB_DIR)
 import cros_build_lib as cros_lib
 
 import au_test
 import au_worker
 import dummy_au_worker
 import dev_server_wrapper
 import parallel_test_job
 import public_key_manager
+import tempfile
 import update_exception
 
 def _PrepareTestSuite(options, use_dummy_worker=False):
   """Returns a prepared test suite given by the options and test class."""
   au_test.AUTest.ProcessOptions(options, use_dummy_worker)
   test_loader = unittest.TestLoader()
   test_loader.testMethodPrefix = options.test_prefix
   return test_loader.loadTestsFromTestCase(au_test.AUTest)
 
 
 def _PregenerateUpdates(options):
   """Determines all deltas that will be generated and generates them.
 
   This method effectively pre-generates the dev server cache for all tests.
 
   Args:
     options: options from parsed parser.
   Returns:
     Dictionary of Update Identifiers->Relative cache locations.
   Raises:
     update_exception.UpdateException if we fail to generate an update.
   """
-  def _GenerateVMUpdate(target, src, private_key_path):
-    """Generates an update using the devserver."""
+  def _GenerateVMUpdate(target, src, private_key_path, log_file):
+    """Returns the error code from generating an update using the devserver."""
     command = ['sudo',
                'start_devserver',
                '--pregenerate_update',
                '--exit',
               ]
     # Add actual args to command.
     command.append('--image=%s' % cros_lib.ReinterpretPathForChroot(target))
     if src: command.append('--src_image=%s' %
                            cros_lib.ReinterpretPathForChroot(src))
     if options.type == 'vm': command.append('--for_vm')
     if private_key_path:
       command.append('--private_key=%s' %
                      cros_lib.ReinterpretPathForChroot(private_key_path))
 
-    return cros_lib.RunCommandCaptureOutput(command, combine_stdout_stderr=True,
-                                            enter_chroot=True,
-                                            print_cmd=True,
-                                            cwd=cros_lib.GetCrosUtilsPath())
+    return cros_lib.RunCommand(command, enter_chroot=True, print_cmd=True,
+                               cwd=cros_lib.GetCrosUtilsPath(),
+                               log_to_file=log_file, error_ok=True,
+                               exit_code=True)
+
+  def _ProcessGeneratorOutputs(log_files, return_codes):
+    """Processes results from the log files of GenerateVMUpdate calls.
+
+    Returns an array of cache entries from the log files.
+    """
+    return_array = []
+    for log_file, return_code in map(lambda x, y: (x, y), log_files,
+                                      return_codes):

[dgarrett, 2011/04/12 02:01:52]

off by one indention.

[sosa, 2011/04/12 02:09:50]

Done.

+      log_file_handle = open(log_file)
+      output = log_file_handle.read()
+      log_file_handle.close()
+
+      # Looking for this line in the output.
+      key_line_re = re.compile('^PREGENERATED_UPDATE=([\w/.]+)')

[dgarrett, 2011/04/12 02:01:52]

Can do this outside the outer for loop and not regenerate for each file.

[sosa, 2011/04/12 02:09:50]

Done.

+      if return_code != 0:

[dgarrett, 2011/04/12 02:01:52]

Seems like this line should be the first line after the for loop on line 82, and
that it doesn't need an else clause at all. Just unindent the normal case code.

[sosa, 2011/04/12 02:09:50]

I actually want to print out the output from a failed update (Warning(output))
so i need to get output first.

On 2011/04/12 02:01:52, dgarrett wrote:

+        cros_lib.Warning(output)
+        raise update_exception.UpdateException(return_code,
+                                               'Failed to generate update.')
+      else:
+        for line in output.splitlines():
+          match = key_line_re.search(line)
+          if match:
+            # Convert blah/blah/update.gz -> update/blah/blah.
+            path_to_update_gz = match.group(1).rstrip()
+            (path_to_update_dir, _, _) = path_to_update_gz.rpartition(
+                '/update.gz')
+            return_array.append('/'.join(['update', path_to_update_dir]))
+
+    assert len(return_array) == len(log_files), 'Return result size mismatch.'
+    return return_array
 
   # Use dummy class to mock out updates that would be run as part of a test.
   test_suite = _PrepareTestSuite(options, use_dummy_worker=True)
   test_result = unittest.TextTestRunner(verbosity=0).run(test_suite)
   if not test_result.wasSuccessful():
     raise update_exception.UpdateException(1,
                                            'Error finding updates to generate.')
 
   cros_lib.Info('The following delta updates are required.')
   update_ids = []
   jobs = []
   args = []
+  log_files = []
   modified_images = set()
   for target, srcs in dummy_au_worker.DummyAUWorker.delta_list.items():
     modified_images.add(target)
     for src_key in srcs:
+      log_file = tempfile.mktemp('GenerateVMUpdate')
       (src, _ , key) = src_key.partition('+')
       if src: modified_images.add(src)
       # TODO(sosa): Add private key as part of caching name once devserver can
       # handle it its own cache.
       update_id = dev_server_wrapper.GenerateUpdateId(target, src, key)
       print >> sys.stderr, 'AU: %s' % update_id
       update_ids.append(update_id)
       jobs.append(_GenerateVMUpdate)
-      args.append((target, src, key))
+      args.append((target, src, key, log_file))
+      log_files.append(log_file)
 
   # Always add the base image path.  This is only useful for non-delta updates.
   modified_images.add(options.base_image)
 
   # Add public key to all images we are using.
   if options.public_key:
     cros_lib.Info('Adding public keys to images for testing.')
     for image in modified_images:
       manager = public_key_manager.PublicKeyManager(image, options.public_key)
       manager.AddKeyToImage()
       au_test.AUTest.public_key_managers.append(manager)
 
-  raw_results = parallel_test_job.RunParallelJobs(options.jobs, jobs, args,
-                                                  print_status=True)
-  results = []
-
-  # Looking for this line in the output.
-  key_line_re = re.compile('^PREGENERATED_UPDATE=([\w/.]+)')
-  for result in raw_results:
-    (return_code, output, _) = result
-    if return_code != 0:
-      cros_lib.Warning(output)
-      raise update_exception.UpdateException(return_code,
-                                             'Failed to generate all updates.')
-    else:
-      for line in output.splitlines():
-        match = key_line_re.search(line)
-        if match:
-          # Convert blah/blah/update.gz -> update/blah/blah.
-          path_to_update_gz = match.group(1).rstrip()
-          (path_to_update_dir, _, _) = path_to_update_gz.rpartition(
-              '/update.gz')
-          results.append('/'.join(['update', path_to_update_dir]))
-          break
-
-  # Make sure all generation of updates returned cached locations.
-  if len(raw_results) != len(results):
-    raise update_exception.UpdateException(
-        1, 'Insufficient number cache directories returned.')
+  error_codes = parallel_test_job.RunParallelJobs(options.jobs, jobs, args)
+  results = _ProcessGeneratorOutputs(log_files, error_codes)
 
   # Build the dictionary from our id's and returned cache paths.
   cache_dictionary = {}
   for index, id in enumerate(update_ids):
     cache_dictionary[id] = results[index]
 
   return cache_dictionary
 
 
 def _RunTestsInParallel(options):
   """Runs the tests given by the options in parallel."""
   threads = []
   args = []
   test_suite = _PrepareTestSuite(options)
   for test in test_suite:
     test_name = test.id()
     test_case = unittest.TestLoader().loadTestsFromName(test_name)
     threads.append(unittest.TextTestRunner(verbosity=2).run)
     args.append(test_case)
 
-  results = parallel_test_job.RunParallelJobs(options.jobs, threads, args,
-                                              print_status=False)
+  results = parallel_test_job.RunParallelJobs(options.jobs, threads, args)
   for test_result in results:
     if not test_result.wasSuccessful():
       cros_lib.Die('Test harness was not successful')
 
 
 def _CleanPreviousWork(options):
   """Cleans up previous work from the devserver cache and local image cache."""
   cros_lib.Info('Cleaning up previous work.')
   # Wipe devserver cache.
   cros_lib.RunCommandCaptureOutput(
@@ skipping 102 matching lines @@
     if options.public_key:
       cros_lib.Info('Cleaning up.  Removing keys added as part of testing.')
       target_directory = os.path.dirname(options.target_image)
       for key_manager in au_test.AUTest.public_key_managers:
         if key_manager.image_path.startswith(target_directory):
           key_manager.RemoveKeyFromImage()
 
 
 if __name__ == '__main__':
   main()