thread.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217
  1. # Copyright 2009 Brian Quinlan. All Rights Reserved.
  2. # Licensed to PSF under a Contributor Agreement.
  3. """Implements ThreadPoolExecutor."""
  4. __author__ = 'Brian Quinlan (brian@sweetapp.com)'
  5. import atexit
  6. from concurrent.futures import _base
  7. import itertools
  8. import queue
  9. import threading
  10. import weakref
  11. import os
  12. # Workers are created as daemon threads. This is done to allow the interpreter
  13. # to exit when there are still idle threads in a ThreadPoolExecutor's thread
  14. # pool (i.e. shutdown() was not called). However, allowing workers to die with
  15. # the interpreter has two undesirable properties:
  16. # - The workers would still be running during interpreter shutdown,
  17. # meaning that they would fail in unpredictable ways.
  18. # - The workers could be killed while evaluating a work item, which could
  19. # be bad if the callable being evaluated has external side-effects e.g.
  20. # writing to a file.
  21. #
  22. # To work around this problem, an exit handler is installed which tells the
  23. # workers to exit when their work queues are empty and then waits until the
  24. # threads finish.
  25. _threads_queues = weakref.WeakKeyDictionary()
  26. _shutdown = False
  27. def _python_exit():
  28. global _shutdown
  29. _shutdown = True
  30. items = list(_threads_queues.items())
  31. for t, q in items:
  32. q.put(None)
  33. for t, q in items:
  34. t.join()
  35. atexit.register(_python_exit)
  36. class _WorkItem(object):
  37. def __init__(self, future, fn, args, kwargs):
  38. self.future = future
  39. self.fn = fn
  40. self.args = args
  41. self.kwargs = kwargs
  42. def run(self):
  43. if not self.future.set_running_or_notify_cancel():
  44. return
  45. try:
  46. result = self.fn(*self.args, **self.kwargs)
  47. except BaseException as exc:
  48. self.future.set_exception(exc)
  49. # Break a reference cycle with the exception 'exc'
  50. self = None
  51. else:
  52. self.future.set_result(result)
  53. def _worker(executor_reference, work_queue, initializer, initargs):
  54. if initializer is not None:
  55. try:
  56. initializer(*initargs)
  57. except BaseException:
  58. _base.LOGGER.critical('Exception in initializer:', exc_info=True)
  59. executor = executor_reference()
  60. if executor is not None:
  61. executor._initializer_failed()
  62. return
  63. try:
  64. while True:
  65. work_item = work_queue.get(block=True)
  66. if work_item is not None:
  67. work_item.run()
  68. # Delete references to object. See issue16284
  69. del work_item
  70. continue
  71. executor = executor_reference()
  72. # Exit if:
  73. # - The interpreter is shutting down OR
  74. # - The executor that owns the worker has been collected OR
  75. # - The executor that owns the worker has been shutdown.
  76. if _shutdown or executor is None or executor._shutdown:
  77. # Flag the executor as shutting down as early as possible if it
  78. # is not gc-ed yet.
  79. if executor is not None:
  80. executor._shutdown = True
  81. # Notice other workers
  82. work_queue.put(None)
  83. return
  84. del executor
  85. except BaseException:
  86. _base.LOGGER.critical('Exception in worker', exc_info=True)
  87. class BrokenThreadPool(_base.BrokenExecutor):
  88. """
  89. Raised when a worker thread in a ThreadPoolExecutor failed initializing.
  90. """
  91. class ThreadPoolExecutor(_base.Executor):
  92. # Used to assign unique thread names when thread_name_prefix is not supplied.
  93. _counter = itertools.count().__next__
  94. def __init__(self, max_workers=None, thread_name_prefix='',
  95. initializer=None, initargs=()):
  96. """Initializes a new ThreadPoolExecutor instance.
  97. Args:
  98. max_workers: The maximum number of threads that can be used to
  99. execute the given calls.
  100. thread_name_prefix: An optional name prefix to give our threads.
  101. initializer: A callable used to initialize worker threads.
  102. initargs: A tuple of arguments to pass to the initializer.
  103. """
  104. if max_workers is None:
  105. # Use this number because ThreadPoolExecutor is often
  106. # used to overlap I/O instead of CPU work.
  107. max_workers = (os.cpu_count() or 1) * 5
  108. if max_workers <= 0:
  109. raise ValueError("max_workers must be greater than 0")
  110. if initializer is not None and not callable(initializer):
  111. raise TypeError("initializer must be a callable")
  112. self._max_workers = max_workers
  113. self._work_queue = queue.SimpleQueue()
  114. self._threads = set()
  115. self._broken = False
  116. self._shutdown = False
  117. self._shutdown_lock = threading.Lock()
  118. self._thread_name_prefix = (thread_name_prefix or
  119. ("ThreadPoolExecutor-%d" % self._counter()))
  120. self._initializer = initializer
  121. self._initargs = initargs
  122. def submit(*args, **kwargs):
  123. if len(args) >= 2:
  124. self, fn, *args = args
  125. elif not args:
  126. raise TypeError("descriptor 'submit' of 'ThreadPoolExecutor' object "
  127. "needs an argument")
  128. elif 'fn' in kwargs:
  129. fn = kwargs.pop('fn')
  130. self, *args = args
  131. else:
  132. raise TypeError('submit expected at least 1 positional argument, '
  133. 'got %d' % (len(args)-1))
  134. with self._shutdown_lock:
  135. if self._broken:
  136. raise BrokenThreadPool(self._broken)
  137. if self._shutdown:
  138. raise RuntimeError('cannot schedule new futures after shutdown')
  139. if _shutdown:
  140. raise RuntimeError('cannot schedule new futures after '
  141. 'interpreter shutdown')
  142. f = _base.Future()
  143. w = _WorkItem(f, fn, args, kwargs)
  144. self._work_queue.put(w)
  145. self._adjust_thread_count()
  146. return f
  147. submit.__doc__ = _base.Executor.submit.__doc__
  148. def _adjust_thread_count(self):
  149. # When the executor gets lost, the weakref callback will wake up
  150. # the worker threads.
  151. def weakref_cb(_, q=self._work_queue):
  152. q.put(None)
  153. # TODO(bquinlan): Should avoid creating new threads if there are more
  154. # idle threads than items in the work queue.
  155. num_threads = len(self._threads)
  156. if num_threads < self._max_workers:
  157. thread_name = '%s_%d' % (self._thread_name_prefix or self,
  158. num_threads)
  159. t = threading.Thread(name=thread_name, target=_worker,
  160. args=(weakref.ref(self, weakref_cb),
  161. self._work_queue,
  162. self._initializer,
  163. self._initargs))
  164. t.daemon = True
  165. t.start()
  166. self._threads.add(t)
  167. _threads_queues[t] = self._work_queue
  168. def _initializer_failed(self):
  169. with self._shutdown_lock:
  170. self._broken = ('A thread initializer failed, the thread pool '
  171. 'is not usable anymore')
  172. # Drain work queue and mark pending futures failed
  173. while True:
  174. try:
  175. work_item = self._work_queue.get_nowait()
  176. except queue.Empty:
  177. break
  178. if work_item is not None:
  179. work_item.future.set_exception(BrokenThreadPool(self._broken))
  180. def shutdown(self, wait=True):
  181. with self._shutdown_lock:
  182. self._shutdown = True
  183. self._work_queue.put(None)
  184. if wait:
  185. for t in self._threads:
  186. t.join()
  187. shutdown.__doc__ = _base.Executor.shutdown.__doc__