Chromium Code Reviews| OLD | NEW |
|---|---|
| (Empty) | |
| 1 # Copyright 2016 The Chromium Authors. All rights reserved. | |
| 2 # Use of this source code is governed by a BSD-style license that can be | |
| 3 # found in the LICENSE file. | |
| 4 | |
| 5 import multiprocessing | |
| 6 import os | |
| 7 import Queue | |
| 8 import resource | |
| 9 import signal | |
| 10 | |
| 11 import psutil | |
| 12 | |
| 13 | |
| 14 def _LimitMemory(memory_share): | |
| 15 """Limits the memory available to this process, to avoid OOM issues. | |
| 16 | |
| 17 Args: | |
| 18 memory_share: (float) Share coefficient of the total physical memory that | |
| 19 the process can use. | |
| 20 """ | |
| 21 total_memory = psutil.virtual_memory().total | |
| 22 memory_limit = memory_share * total_memory | |
| 23 resource.setrlimit(resource.RLIMIT_AS, (memory_limit, -1L)) | |
| 24 | |
| 25 | |
| 26 def _MultiprocessingWrapper(queue, memory_share, function, args): | |
| 27 """Helper function that sets a memory limit on the current process, then | |
| 28 calls |function| on |args| and write the results to |queue|. | |
| 29 | |
| 30 Args: | |
| 31 queue: (multiprocessing.Queue) Queue where the results of the wrapped | |
| 32 function are written. | |
| 33 memory_share: (float) Share coefficient of the total physical memory that | |
| 34 the process can use. | |
| 35 function: The wrapped function. | |
| 36 args: (list) Arguments for the wrapped function. | |
| 37 """ | |
| 38 try: | |
| 39 if memory_share: | |
| 40 _LimitMemory(memory_share) | |
| 41 | |
| 42 queue.put(function(*args)) | |
| 43 except Exception: | |
| 44 queue.put(None) | |
| 45 | |
| 46 | |
| 47 def RunOutOfProcess(function, args, logger, timeout_seconds, memory_share=None): | |
|
blundell
2016/06/06 14:12:00
nit: maybe RunInSeparateProcess()?
| |
| 48 """Runs a function in a separate process, and kills it after the timeout is | |
| 49 reached. | |
| 50 | |
| 51 Args: | |
| 52 function: The function to run. | |
| 53 args: (list) Arguments for the wrapped function. | |
| 54 timeout_seconds: (float) Timeout in seconds after which the subprocess is | |
| 55 terminated. | |
| 56 memory_share: (float) Set this parameter to limit the memory available to | |
| 57 the spawned subprocess. This is a ratio of the total system | |
| 58 memory (between 0 and 1). | |
| 59 Returns: | |
| 60 The result of the wrapped function, or None if the call failed. | |
| 61 """ | |
| 62 queue = multiprocessing.Queue() | |
| 63 process = multiprocessing.Process(target=_MultiprocessingWrapper, | |
| 64 args=(queue, memory_share, function, args)) | |
| 65 process.daemon = True | |
| 66 process.start() | |
| 67 | |
| 68 result = None | |
| 69 | |
| 70 try: | |
| 71 logger.info('Wait for result.') | |
| 72 # Note: If the subprocess somehow crashes (e.g. Python crashing), this | |
| 73 # process will wait the full timeout. Could be avoided but probably not | |
| 74 # worth the extra complexity. | |
| 75 result = queue.get(block=True, timeout=timeout_seconds) | |
| 76 except Queue.Empty: | |
| 77 logger.warning('Subprocess timeout.') | |
| 78 process.terminate() | |
| 79 | |
| 80 logger.info('Wait for process to terminate.') | |
| 81 process.join(timeout=5) | |
| 82 | |
| 83 if process.is_alive(): | |
| 84 logger.warning('Process still alive, hard killing now.') | |
| 85 os.kill(process.pid, signal.SIGKILL) | |
| 86 | |
| 87 return result | |
| OLD | NEW |