tools/android/loading Switch from multiprocessing Pool to Process

tools/android/loading/cloud/backend/trace_task_handler.py

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import multiprocessing
 import os
 import re
-import resource
 import sys
 import traceback
 
-import psutil
-
 import common.clovis_paths
 from common.clovis_task import ClovisTask
 from common.loading_trace_database import LoadingTraceDatabase
 import controller
 from failure_database import FailureDatabase
 import loading_trace
+import multiprocessing_helper
 import options
 import xvfb_helper
 
 
-def LimitMemory(memory_share):
-  """Limits the memory available to this process, to avoid OOM issues.
-
-  Args:
-    memory_share: (float) Share coefficient of the total physical memory that
-                          the process can use.
-  """
-  total_memory = psutil.virtual_memory().total
-  memory_limit = memory_share * total_memory
-  resource.setrlimit(resource.RLIMIT_AS, (memory_limit, -1L))
-
-
 def GenerateTrace(url, emulate_device, emulate_network, filename, log_filename):
   """ Generates a trace.
 
   Args:
     url: URL as a string.
     emulate_device: Name of the device to emulate. Empty for no emulation.
     emulate_network: Type of network emulation. Empty for no emulation.
     filename: Name of the file where the trace is saved.
     log_filename: Name of the file where standard output and errors are
                   logged.
@@ skipping 11 matching lines @@
   old_stderr = sys.stderr
 
   trace_metadata = { 'succeeded' : False, 'url' : url }
   trace = None
   if not url.startswith('http') and not url.startswith('file'):
     url = 'http://' + url
   with open(log_filename, 'w') as sys.stdout:
     try:
       sys.stderr = sys.stdout
 
+      sys.stdout.write('Starting trace generation for: %s.\n' % url)
+
       # Set up the controller.
       chrome_ctl = controller.LocalChromeController()
       chrome_ctl.SetChromeEnvOverride(xvfb_helper.GetChromeEnvironment())
       if emulate_device:
         chrome_ctl.SetDeviceEmulation(emulate_device)
       if emulate_network:
         chrome_ctl.SetNetworkEmulation(emulate_network)
 
       # Record and write the trace.
       with chrome_ctl.Open() as connection:
         connection.ClearCache()
         trace = loading_trace.LoadingTrace.RecordUrlNavigation(
             url, connection, chrome_ctl.ChromeMetadata())
         trace_metadata['succeeded'] = True
         trace_metadata.update(trace.ToJsonDict()[trace._METADATA_KEY])
+        sys.stdout.write('Trace generation success.\n')
     except controller.ChromeControllerError as e:
       e.Dump(sys.stderr)
     except Exception as e:
       sys.stderr.write('Unknown exception:\n' + str(e))
       traceback.print_exc(file=sys.stderr)
 
     if trace:
+      sys.stdout.write('Dumping trace to file.\n')
       trace.ToJsonFile(filename)
+    else:
+      sys.stderr.write('No trace generated.\n')
+
+    sys.stdout.write('Trace generation finished.\n')
 
   sys.stdout = old_stdout
   sys.stderr = old_stderr
 
   return trace_metadata
 
 
 class TraceTaskHandler(object):
   """Handles 'trace' tasks."""
 
@@ skipping 59 matching lines @@
     """ Generates a trace in a separate process by calling GenerateTrace().
 
     The generation is done out of process to avoid issues where the system would
     run out of memory when the trace is very large. This ensures that the system
     can reclaim all the memory when the trace generation is done.
 
     See the GenerateTrace() documentation for a description of the parameters
     and return values.
     """
     self._logger.info('Starting external process for trace generation.')
-    failed_metadata = {'succeeded':False, 'url':url}
-    failed = False
-    pool = multiprocessing.Pool(1, initializer=LimitMemory, initargs=(0.9,))
+    result = multiprocessing_helper.RunInSeparateProcess(
+        GenerateTrace,
+        (url, emulate_device, emulate_network, filename, log_filename),
+        self._logger, timeout_seconds=180, memory_share=0.9)
 
-    apply_result = pool.apply_async(
-        GenerateTrace,
-        (url, emulate_device, emulate_network, filename, log_filename))
-    pool.close()
-    apply_result.wait(timeout=180)
+    self._logger.info('Cleaning up Chrome processes.')
+    controller.LocalChromeController.KillChromeProcesses()
 
-    if not apply_result.ready():
-      self._logger.error('Process timeout for trace generation of URL: ' + url)
+    if not result:
       self._failure_database.AddFailure('trace_process_timeout', url)
-      # Explicitly kill Chrome now, or pool.terminate() will hang.
-      controller.LocalChromeController.KillChromeProcesses()
-      pool.terminate()
-      failed = True
-    elif not apply_result.successful():
-      # Try to reraise the exception that killed the subprocess and add it to
-      # the error log.
-      try:
-        apply_result.get()
-      except Exception as e:
-        with file(log_filename, 'w+') as error_log:
-          error_log.write('Unhandled exception caught by apply_result: {}'
-                          .format(e))
-          traceback.print_exc(file=error_log)
-      else:
-        with file(log_filename, 'w+') as error_log:
-          error_log.write('No exception found for unsuccessful apply_result')
-      self._logger.error('Process failure for trace generation of URL: ' + url)
-      self._failure_database.AddFailure('trace_process_error', url)
-      failed = True
+      return {'succeeded':False, 'url':url}
+    return result
 
-    self._logger.info('Cleaning up external process.')
-    pool.join()
-
-    if failed:
-      return failed_metadata
-    return apply_result.get()
 
   def _HandleTraceGenerationResults(self, local_filename, log_filename,
                                     remote_filename, trace_metadata):
     """Updates the trace database and the failure database after a trace
     generation. Uploads the trace and the log.
     Results related to successful traces are uploaded in the 'traces' directory,
     and failures are uploaded in the 'failures' directory.
 
     Args:
       local_filename (str): Path to the local file containing the trace.
@@ skipping 76 matching lines @@
         trace_metadata = self._GenerateTraceOutOfProcess(
             url, emulate_device, emulate_network, local_filename, log_filename)
         if trace_metadata['succeeded']:
           success_happened = True
         remote_filename = os.path.join(local_filename, str(repeat))
         self._HandleTraceGenerationResults(
             local_filename, log_filename, remote_filename, trace_metadata)
 
     if success_happened:
       self._UploadTraceDatabase()