parallel.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350
  1. from __future__ import absolute_import
  2. from __future__ import unicode_literals
  3. import logging
  4. import operator
  5. import sys
  6. from threading import Lock
  7. from threading import Semaphore
  8. from threading import Thread
  9. from docker.errors import APIError
  10. from docker.errors import ImageNotFound
  11. from six.moves import _thread as thread
  12. from six.moves.queue import Empty
  13. from six.moves.queue import Queue
  14. from compose.cli.colors import green
  15. from compose.cli.colors import red
  16. from compose.cli.signals import ShutdownException
  17. from compose.const import PARALLEL_LIMIT
  18. from compose.errors import HealthCheckFailed
  19. from compose.errors import NoHealthCheckConfigured
  20. from compose.errors import OperationFailedError
  21. from compose.utils import get_output_stream
  22. log = logging.getLogger(__name__)
  23. STOP = object()
  24. class GlobalLimit(object):
  25. """Simple class to hold a global semaphore limiter for a project. This class
  26. should be treated as a singleton that is instantiated when the project is.
  27. """
  28. global_limiter = Semaphore(PARALLEL_LIMIT)
  29. @classmethod
  30. def set_global_limit(cls, value):
  31. if value is None:
  32. value = PARALLEL_LIMIT
  33. cls.global_limiter = Semaphore(value)
  34. def parallel_execute_watch(events, writer, errors, results, msg, get_name):
  35. """ Watch events from a parallel execution, update status and fill errors and results.
  36. Returns exception to re-raise.
  37. """
  38. error_to_reraise = None
  39. for obj, result, exception in events:
  40. if exception is None:
  41. if callable(getattr(obj, 'containers', None)) and not obj.containers():
  42. # If service has no containers started
  43. writer.write(msg, get_name(obj), 'failed', red)
  44. else:
  45. writer.write(msg, get_name(obj), 'done', green)
  46. results.append(result)
  47. elif isinstance(exception, ImageNotFound):
  48. # This is to bubble up ImageNotFound exceptions to the client so we
  49. # can prompt the user if they want to rebuild.
  50. errors[get_name(obj)] = exception.explanation
  51. writer.write(msg, get_name(obj), 'error', red)
  52. error_to_reraise = exception
  53. elif isinstance(exception, APIError):
  54. errors[get_name(obj)] = exception.explanation
  55. writer.write(msg, get_name(obj), 'error', red)
  56. elif isinstance(exception, (OperationFailedError, HealthCheckFailed, NoHealthCheckConfigured)):
  57. errors[get_name(obj)] = exception.msg
  58. writer.write(msg, get_name(obj), 'error', red)
  59. elif isinstance(exception, UpstreamError):
  60. writer.write(msg, get_name(obj), 'error', red)
  61. else:
  62. errors[get_name(obj)] = exception
  63. error_to_reraise = exception
  64. return error_to_reraise
  65. def parallel_execute(objects, func, get_name, msg, get_deps=None, limit=None):
  66. """Runs func on objects in parallel while ensuring that func is
  67. ran on object only after it is ran on all its dependencies.
  68. get_deps called on object must return a collection with its dependencies.
  69. get_name called on object must return its name.
  70. """
  71. objects = list(objects)
  72. stream = get_output_stream(sys.stderr)
  73. if ParallelStreamWriter.instance:
  74. writer = ParallelStreamWriter.instance
  75. else:
  76. writer = ParallelStreamWriter(stream)
  77. for obj in objects:
  78. writer.add_object(msg, get_name(obj))
  79. for obj in objects:
  80. writer.write_initial(msg, get_name(obj))
  81. events = parallel_execute_iter(objects, func, get_deps, limit)
  82. errors = {}
  83. results = []
  84. error_to_reraise = parallel_execute_watch(events, writer, errors, results, msg, get_name)
  85. for obj_name, error in errors.items():
  86. stream.write("\nERROR: for {} {}\n".format(obj_name, error))
  87. if error_to_reraise:
  88. raise error_to_reraise
  89. return results, errors
  90. def _no_deps(x):
  91. return []
  92. class State(object):
  93. """
  94. Holds the state of a partially-complete parallel operation.
  95. state.started: objects being processed
  96. state.finished: objects which have been processed
  97. state.failed: objects which either failed or whose dependencies failed
  98. """
  99. def __init__(self, objects):
  100. self.objects = objects
  101. self.started = set()
  102. self.finished = set()
  103. self.failed = set()
  104. def is_done(self):
  105. return len(self.finished) + len(self.failed) >= len(self.objects)
  106. def pending(self):
  107. return set(self.objects) - self.started - self.finished - self.failed
  108. class NoLimit(object):
  109. def __enter__(self):
  110. pass
  111. def __exit__(self, *ex):
  112. pass
  113. def parallel_execute_iter(objects, func, get_deps, limit):
  114. """
  115. Runs func on objects in parallel while ensuring that func is
  116. ran on object only after it is ran on all its dependencies.
  117. Returns an iterator of tuples which look like:
  118. # if func returned normally when run on object
  119. (object, result, None)
  120. # if func raised an exception when run on object
  121. (object, None, exception)
  122. # if func raised an exception when run on one of object's dependencies
  123. (object, None, UpstreamError())
  124. """
  125. if get_deps is None:
  126. get_deps = _no_deps
  127. if limit is None:
  128. limiter = NoLimit()
  129. else:
  130. limiter = Semaphore(limit)
  131. results = Queue()
  132. state = State(objects)
  133. while True:
  134. feed_queue(objects, func, get_deps, results, state, limiter)
  135. try:
  136. event = results.get(timeout=0.1)
  137. except Empty:
  138. continue
  139. # See https://github.com/docker/compose/issues/189
  140. except thread.error:
  141. raise ShutdownException()
  142. if event is STOP:
  143. break
  144. obj, _, exception = event
  145. if exception is None:
  146. log.debug('Finished processing: {}'.format(obj))
  147. state.finished.add(obj)
  148. else:
  149. log.debug('Failed: {}'.format(obj))
  150. state.failed.add(obj)
  151. yield event
  152. def producer(obj, func, results, limiter):
  153. """
  154. The entry point for a producer thread which runs func on a single object.
  155. Places a tuple on the results queue once func has either returned or raised.
  156. """
  157. with limiter, GlobalLimit.global_limiter:
  158. try:
  159. result = func(obj)
  160. results.put((obj, result, None))
  161. except Exception as e:
  162. results.put((obj, None, e))
  163. def feed_queue(objects, func, get_deps, results, state, limiter):
  164. """
  165. Starts producer threads for any objects which are ready to be processed
  166. (i.e. they have no dependencies which haven't been successfully processed).
  167. Shortcuts any objects whose dependencies have failed and places an
  168. (object, None, UpstreamError()) tuple on the results queue.
  169. """
  170. pending = state.pending()
  171. log.debug('Pending: {}'.format(pending))
  172. for obj in pending:
  173. deps = get_deps(obj)
  174. try:
  175. if any(dep[0] in state.failed for dep in deps):
  176. log.debug('{} has upstream errors - not processing'.format(obj))
  177. results.put((obj, None, UpstreamError()))
  178. state.failed.add(obj)
  179. elif all(
  180. dep not in objects or (
  181. dep in state.finished and (not ready_check or ready_check(dep))
  182. ) for dep, ready_check in deps
  183. ):
  184. log.debug('Starting producer thread for {}'.format(obj))
  185. t = Thread(target=producer, args=(obj, func, results, limiter))
  186. t.daemon = True
  187. t.start()
  188. state.started.add(obj)
  189. except (HealthCheckFailed, NoHealthCheckConfigured) as e:
  190. log.debug(
  191. 'Healthcheck for service(s) upstream of {} failed - '
  192. 'not processing'.format(obj)
  193. )
  194. results.put((obj, None, e))
  195. if state.is_done():
  196. results.put(STOP)
  197. class UpstreamError(Exception):
  198. pass
  199. class ParallelStreamWriter(object):
  200. """Write out messages for operations happening in parallel.
  201. Each operation has its own line, and ANSI code characters are used
  202. to jump to the correct line, and write over the line.
  203. """
  204. noansi = False
  205. lock = Lock()
  206. instance = None
  207. @classmethod
  208. def set_noansi(cls, value=True):
  209. cls.noansi = value
  210. def __init__(self, stream):
  211. self.stream = stream
  212. self.lines = []
  213. self.width = 0
  214. ParallelStreamWriter.instance = self
  215. def add_object(self, msg, obj_index):
  216. if msg is None:
  217. return
  218. self.lines.append(msg + obj_index)
  219. self.width = max(self.width, len(msg + ' ' + obj_index))
  220. def write_initial(self, msg, obj_index):
  221. if msg is None:
  222. return
  223. return self._write_noansi(msg, obj_index, '')
  224. def _write_ansi(self, msg, obj_index, status):
  225. self.lock.acquire()
  226. position = self.lines.index(msg + obj_index)
  227. diff = len(self.lines) - position
  228. # move up
  229. self.stream.write("%c[%dA" % (27, diff))
  230. # erase
  231. self.stream.write("%c[2K\r" % 27)
  232. self.stream.write("{:<{width}} ... {}\r".format(msg + ' ' + obj_index,
  233. status, width=self.width))
  234. # move back down
  235. self.stream.write("%c[%dB" % (27, diff))
  236. self.stream.flush()
  237. self.lock.release()
  238. def _write_noansi(self, msg, obj_index, status):
  239. self.stream.write(
  240. "{:<{width}} ... {}\r\n".format(
  241. msg + ' ' + obj_index, status, width=self.width
  242. )
  243. )
  244. self.stream.flush()
  245. def write(self, msg, obj_index, status, color_func):
  246. if msg is None:
  247. return
  248. if self.noansi:
  249. self._write_noansi(msg, obj_index, status)
  250. else:
  251. self._write_ansi(msg, obj_index, color_func(status))
  252. def get_stream_writer():
  253. instance = ParallelStreamWriter.instance
  254. if instance is None:
  255. raise RuntimeError('ParallelStreamWriter has not yet been instantiated')
  256. return instance
  257. def parallel_operation(containers, operation, options, message):
  258. parallel_execute(
  259. containers,
  260. operator.methodcaller(operation, **options),
  261. operator.attrgetter('name'),
  262. message,
  263. )
  264. def parallel_remove(containers, options):
  265. stopped_containers = [c for c in containers if not c.is_running]
  266. parallel_operation(stopped_containers, 'remove', options, 'Removing')
  267. def parallel_pause(containers, options):
  268. parallel_operation(containers, 'pause', options, 'Pausing')
  269. def parallel_unpause(containers, options):
  270. parallel_operation(containers, 'unpause', options, 'Unpausing')
  271. def parallel_kill(containers, options):
  272. parallel_operation(containers, 'kill', options, 'Killing')