thread.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  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 queue
  8. import threading
  9. import weakref
  10. import os
  11. # Workers are created as daemon threads. This is done to allow the interpreter
  12. # to exit when there are still idle threads in a ThreadPoolExecutor's thread
  13. # pool (i.e. shutdown() was not called). However, allowing workers to die with
  14. # the interpreter has two undesirable properties:
  15. # - The workers would still be running during interpretor shutdown,
  16. # meaning that they would fail in unpredictable ways.
  17. # - The workers could be killed while evaluating a work item, which could
  18. # be bad if the callable being evaluated has external side-effects e.g.
  19. # writing to a file.
  20. #
  21. # To work around this problem, an exit handler is installed which tells the
  22. # workers to exit when their work queues are empty and then waits until the
  23. # threads finish.
  24. _threads_queues = weakref.WeakKeyDictionary()
  25. _shutdown = False
  26. def _python_exit():
  27. global _shutdown
  28. _shutdown = True
  29. items = list(_threads_queues.items())
  30. for t, q in items:
  31. q.put(None)
  32. for t, q in items:
  33. t.join()
  34. atexit.register(_python_exit)
  35. class _WorkItem(object):
  36. def __init__(self, future, fn, args, kwargs):
  37. self.future = future
  38. self.fn = fn
  39. self.args = args
  40. self.kwargs = kwargs
  41. def run(self):
  42. if not self.future.set_running_or_notify_cancel():
  43. return
  44. try:
  45. result = self.fn(*self.args, **self.kwargs)
  46. except BaseException as e:
  47. self.future.set_exception(e)
  48. else:
  49. self.future.set_result(result)
  50. def _worker(executor_reference, work_queue):
  51. try:
  52. while True:
  53. work_item = work_queue.get(block=True)
  54. if work_item is not None:
  55. work_item.run()
  56. # Delete references to object. See issue16284
  57. del work_item
  58. continue
  59. executor = executor_reference()
  60. # Exit if:
  61. # - The interpreter is shutting down OR
  62. # - The executor that owns the worker has been collected OR
  63. # - The executor that owns the worker has been shutdown.
  64. if _shutdown or executor is None or executor._shutdown:
  65. # Notice other workers
  66. work_queue.put(None)
  67. return
  68. del executor
  69. except BaseException:
  70. _base.LOGGER.critical('Exception in worker', exc_info=True)
  71. class ThreadPoolExecutor(_base.Executor):
  72. def __init__(self, max_workers=None):
  73. """Initializes a new ThreadPoolExecutor instance.
  74. Args:
  75. max_workers: The maximum number of threads that can be used to
  76. execute the given calls.
  77. """
  78. if max_workers is None:
  79. # Use this number because ThreadPoolExecutor is often
  80. # used to overlap I/O instead of CPU work.
  81. max_workers = (os.cpu_count() or 1) * 5
  82. if max_workers <= 0:
  83. raise ValueError("max_workers must be greater than 0")
  84. self._max_workers = max_workers
  85. self._work_queue = queue.Queue()
  86. self._threads = set()
  87. self._shutdown = False
  88. self._shutdown_lock = threading.Lock()
  89. def submit(self, fn, *args, **kwargs):
  90. with self._shutdown_lock:
  91. if self._shutdown:
  92. raise RuntimeError('cannot schedule new futures after shutdown')
  93. f = _base.Future()
  94. w = _WorkItem(f, fn, args, kwargs)
  95. self._work_queue.put(w)
  96. self._adjust_thread_count()
  97. return f
  98. submit.__doc__ = _base.Executor.submit.__doc__
  99. def _adjust_thread_count(self):
  100. # When the executor gets lost, the weakref callback will wake up
  101. # the worker threads.
  102. def weakref_cb(_, q=self._work_queue):
  103. q.put(None)
  104. # TODO(bquinlan): Should avoid creating new threads if there are more
  105. # idle threads than items in the work queue.
  106. if len(self._threads) < self._max_workers:
  107. t = threading.Thread(target=_worker,
  108. args=(weakref.ref(self, weakref_cb),
  109. self._work_queue))
  110. t.daemon = True
  111. t.start()
  112. self._threads.add(t)
  113. _threads_queues[t] = self._work_queue
  114. def shutdown(self, wait=True):
  115. with self._shutdown_lock:
  116. self._shutdown = True
  117. self._work_queue.put(None)
  118. if wait:
  119. for t in self._threads:
  120. t.join()
  121. shutdown.__doc__ = _base.Executor.shutdown.__doc__