thread.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. # Copyright 2009 Brian Quinlan. All Rights Reserved.
  2. # Licensed to PSF under a Contributor Agreement.
  3. """Implements ThreadPoolExecutor."""
  4. from __future__ import with_statement
  5. import atexit
  6. import threading
  7. import weakref
  8. import sys
  9. from concurrent.futures import _base
  10. try:
  11. import queue
  12. except ImportError:
  13. import Queue as queue
  14. __author__ = 'Brian Quinlan (brian@sweetapp.com)'
  15. # Workers are created as daemon threads. This is done to allow the interpreter
  16. # to exit when there are still idle threads in a ThreadPoolExecutor's thread
  17. # pool (i.e. shutdown() was not called). However, allowing workers to die with
  18. # the interpreter has two undesirable properties:
  19. # - The workers would still be running during interpretor shutdown,
  20. # meaning that they would fail in unpredictable ways.
  21. # - The workers could be killed while evaluating a work item, which could
  22. # be bad if the callable being evaluated has external side-effects e.g.
  23. # writing to a file.
  24. #
  25. # To work around this problem, an exit handler is installed which tells the
  26. # workers to exit when their work queues are empty and then waits until the
  27. # threads finish.
  28. _thread_references = set()
  29. _shutdown = False
  30. def _python_exit():
  31. global _shutdown
  32. _shutdown = True
  33. for thread_reference in _thread_references:
  34. thread = thread_reference()
  35. if thread is not None:
  36. thread.join()
  37. def _remove_dead_thread_references():
  38. """Remove inactive threads from _thread_references.
  39. Should be called periodically to prevent memory leaks in scenarios such as:
  40. >>> while True:
  41. ... t = ThreadPoolExecutor(max_workers=5)
  42. ... t.map(int, ['1', '2', '3', '4', '5'])
  43. """
  44. for thread_reference in set(_thread_references):
  45. if thread_reference() is None:
  46. _thread_references.discard(thread_reference)
  47. atexit.register(_python_exit)
  48. class _WorkItem(object):
  49. def __init__(self, future, fn, args, kwargs):
  50. self.future = future
  51. self.fn = fn
  52. self.args = args
  53. self.kwargs = kwargs
  54. def run(self):
  55. if not self.future.set_running_or_notify_cancel():
  56. return
  57. try:
  58. result = self.fn(*self.args, **self.kwargs)
  59. except BaseException:
  60. e = sys.exc_info()[1]
  61. self.future.set_exception(e)
  62. else:
  63. self.future.set_result(result)
  64. def _worker(executor_reference, work_queue):
  65. try:
  66. while True:
  67. try:
  68. work_item = work_queue.get(block=True, timeout=0.1)
  69. except queue.Empty:
  70. executor = executor_reference()
  71. # Exit if:
  72. # - The interpreter is shutting down OR
  73. # - The executor that owns the worker has been collected OR
  74. # - The executor that owns the worker has been shutdown.
  75. if _shutdown or executor is None or executor._shutdown:
  76. return
  77. del executor
  78. else:
  79. work_item.run()
  80. except BaseException:
  81. _base.LOGGER.critical('Exception in worker', exc_info=True)
  82. class ThreadPoolExecutor(_base.Executor):
  83. def __init__(self, max_workers):
  84. """Initializes a new ThreadPoolExecutor instance.
  85. Args:
  86. max_workers: The maximum number of threads that can be used to
  87. execute the given calls.
  88. """
  89. _remove_dead_thread_references()
  90. self._max_workers = max_workers
  91. self._work_queue = queue.Queue()
  92. self._threads = set()
  93. self._shutdown = False
  94. self._shutdown_lock = threading.Lock()
  95. def submit(self, fn, *args, **kwargs):
  96. with self._shutdown_lock:
  97. if self._shutdown:
  98. raise RuntimeError('cannot schedule new futures after shutdown')
  99. f = _base.Future()
  100. w = _WorkItem(f, fn, args, kwargs)
  101. self._work_queue.put(w)
  102. self._adjust_thread_count()
  103. return f
  104. submit.__doc__ = _base.Executor.submit.__doc__
  105. def _adjust_thread_count(self):
  106. # TODO(bquinlan): Should avoid creating new threads if there are more
  107. # idle threads than items in the work queue.
  108. if len(self._threads) < self._max_workers:
  109. t = threading.Thread(target=_worker,
  110. args=(weakref.ref(self), self._work_queue))
  111. t.daemon = True
  112. t.start()
  113. self._threads.add(t)
  114. _thread_references.add(weakref.ref(t))
  115. def shutdown(self, wait=True):
  116. with self._shutdown_lock:
  117. self._shutdown = True
  118. if wait:
  119. for t in self._threads:
  120. t.join()
  121. shutdown.__doc__ = _base.Executor.shutdown.__doc__