| OLD | NEW |
| 1 # Copyright 2016 The Chromium Authors. All rights reserved. | 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 | 2 # Use of this source code is governed by a BSD-style license that can be |
| 3 # found in the LICENSE file. | 3 # found in the LICENSE file. |
| 4 | 4 |
| 5 """This module contains util functions that local scripts can use.""" | 5 """This module contains util functions that local scripts can use.""" |
| 6 | 6 |
| 7 import atexit |
| 8 import cgi |
| 9 import json |
| 7 import logging | 10 import logging |
| 8 import os | 11 import os |
| 12 import Queue |
| 13 import re |
| 9 import subprocess | 14 import subprocess |
| 10 import sys | 15 import sys |
| 16 import threading |
| 17 import traceback |
| 18 import time |
| 11 | 19 |
| 12 from lib.cache_decorator import Cached | 20 MAX_THREAD_NUMBER = 10 |
| 13 from local_cache import LocalCacher # pylint: disable=W | 21 TASK_QUEUE = None |
| 14 | 22 |
| 15 | 23 |
| 16 def SetUpSystemPaths(): # pragma: no cover | 24 def SetUpSystemPaths(): # pragma: no cover |
| 17 """Sets system paths so as to import modules in findit, third_party and | 25 """Sets system paths so as to import modules in findit, third_party and |
| 18 appengine.""" | 26 appengine.""" |
| 19 findit_root_dir = os.path.join(os.path.dirname(__file__), os.path.pardir) | 27 findit_root_dir = os.path.join(os.path.dirname(__file__), os.path.pardir) |
| 20 third_party_dir = os.path.join(findit_root_dir, 'third_party') | 28 third_party_dir = os.path.join(findit_root_dir, 'third_party') |
| 21 appengine_sdk_dir = os.path.join(findit_root_dir, os.path.pardir, | 29 appengine_sdk_dir = os.path.join(findit_root_dir, os.path.pardir, |
| 22 os.path.pardir, os.path.pardir, | 30 os.path.pardir, os.path.pardir, |
| 23 'google_appengine') | 31 'google_appengine') |
| 24 | 32 |
| 25 # Add App Engine SDK dir to sys.path. | 33 # Add App Engine SDK dir to sys.path. |
| 26 sys.path.insert(1, appengine_sdk_dir) | 34 sys.path.insert(1, appengine_sdk_dir) |
| 27 sys.path.insert(1, third_party_dir) | 35 sys.path.insert(1, third_party_dir) |
| 28 import dev_appserver | 36 import dev_appserver |
| 29 dev_appserver.fix_sys_path() | 37 dev_appserver.fix_sys_path() |
| 30 | 38 |
| 31 # Add Findit root dir to sys.path so that modules in Findit is available. | 39 # Add Findit root dir to sys.path so that modules in Findit is available. |
| 32 sys.path.insert(1, findit_root_dir) | 40 sys.path.insert(1, findit_root_dir) |
| 33 | 41 |
| 34 | 42 |
| 43 SetUpSystemPaths() |
| 44 |
| 45 # The lib is in predator/ root dir, and can be imported only when sys.path get |
| 46 # set up. |
| 47 from lib.cache_decorator import Cached |
| 48 from local_cache import LocalCacher # pylint: disable=W |
| 49 |
| 50 |
| 51 def SignalWorkerThreads(): # pragma: no cover |
| 52 """Puts signal worker threads into task queue.""" |
| 53 global TASK_QUEUE # pylint: disable=W0602 |
| 54 if not TASK_QUEUE: |
| 55 return |
| 56 |
| 57 for _ in range(MAX_THREAD_NUMBER): |
| 58 TASK_QUEUE.put(None) |
| 59 |
| 60 # Give worker threads a chance to exit. |
| 61 # Workaround the harmless bug in python 2.7 below. |
| 62 time.sleep(1) |
| 63 |
| 64 |
| 65 atexit.register(SignalWorkerThreads) |
| 66 |
| 67 |
| 68 def Worker(): # pragma: no cover |
| 69 global TASK_QUEUE # pylint: disable=W0602 |
| 70 while True: |
| 71 try: |
| 72 task = TASK_QUEUE.get() |
| 73 if not task: |
| 74 return |
| 75 except TypeError: |
| 76 # According to http://bugs.python.org/issue14623, this is a harmless bug |
| 77 # in python 2.7 which won't be fixed. |
| 78 # The exception is raised on daemon threads when python interpreter is |
| 79 # shutting down. |
| 80 return |
| 81 |
| 82 function, args, kwargs, result_semaphore = task |
| 83 try: |
| 84 function(*args, **kwargs) |
| 85 except Exception as e: |
| 86 print e |
| 87 # Continue to process tasks in queue, in case every thread fails, the |
| 88 # main thread will be waiting forever. |
| 89 continue |
| 90 finally: |
| 91 # Signal one task is done in case of exception. |
| 92 result_semaphore.release() |
| 93 |
| 94 |
| 95 def RunTasks(tasks): # pragma: no cover |
| 96 """Run given tasks. Not thread-safe: no concurrent calls of this function. |
| 97 |
| 98 Return after all tasks were completed. A task is a dict as below: |
| 99 { |
| 100 'function': the function to call, |
| 101 'args': the positional argument to pass to the function, |
| 102 'kwargs': the key-value arguments to pass to the function, |
| 103 } |
| 104 """ |
| 105 if not tasks: |
| 106 return |
| 107 |
| 108 global TASK_QUEUE |
| 109 if not TASK_QUEUE: |
| 110 TASK_QUEUE = Queue.Queue() |
| 111 for index in range(MAX_THREAD_NUMBER): |
| 112 thread = threading.Thread(target=Worker, name='worker_%s' % index) |
| 113 # Set as daemon, so no join is needed. |
| 114 thread.daemon = True |
| 115 thread.start() |
| 116 |
| 117 result_semaphore = threading.Semaphore(0) |
| 118 # Push task to task queue for execution. |
| 119 for task in tasks: |
| 120 TASK_QUEUE.put((task['function'], task.get('args', []), |
| 121 task.get('kwargs', {}), result_semaphore)) |
| 122 |
| 123 # Wait until all tasks to be executed. |
| 124 for _ in tasks: |
| 125 result_semaphore.acquire() |
| 126 |
| 127 |
| 35 @Cached(namespace='Command-output', cacher=LocalCacher()) | 128 @Cached(namespace='Command-output', cacher=LocalCacher()) |
| 36 def GetCommandOutput(command): # pragma: no cover | 129 def GetCommandOutput(command): # pragma: no cover |
| 37 """Gets the output stream of executable command. | 130 """Gets the output stream of executable command. |
| 38 | 131 |
| 39 Args: | 132 Args: |
| 40 command (str): Command to execute to get output. | 133 command (str): Command to execute to get output. |
| 41 | 134 |
| 42 Return: | 135 Return: |
| 43 Output steam of the command. | 136 Output steam of the command. |
| 44 """ | 137 """ |
| 45 p = subprocess.Popen( | 138 p = subprocess.Popen( |
| 46 command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) | 139 command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True) |
| 47 stdoutdata, stderrdata = p.communicate() | 140 stdoutdata, stderrdata = p.communicate() |
| 48 | 141 |
| 49 if p.returncode != 0: | 142 if p.returncode != 0: |
| 50 logging.error('Error running command %s: %s', command, stderrdata) | 143 logging.error('Error running command %s: %s', command, stderrdata) |
| 51 return None | 144 return None |
| 52 | 145 |
| 53 return stdoutdata | 146 return stdoutdata |
| OLD | NEW |