thread.py

# Copyright 2009 Brian Quinlan. All Rights Reserved.
# Licensed to PSF under a Contributor Agreement.

"""Implements ThreadPoolExecutor."""

__author__='Brian Quinlan (brian@sweetapp.com)'

importatexit
fromconcurrent.futuresimport_base
importitertools
importqueue
importthreading
importweakref
importos

# Workers are created as daemon threads. This is done to allow the interpreter
# to exit when there are still idle threads in a ThreadPoolExecutor's thread
# pool (i.e. shutdown() was not called). However, allowing workers to die with
# the interpreter has two undesirable properties:
# - The workers would still be running during interpreter shutdown,
# meaning that they would fail in unpredictable ways.
# - The workers could be killed while evaluating a work item, which could
# be bad if the callable being evaluated has external side-effects e.g.
# writing to a file.
#
# To work around this problem, an exit handler is installed which tells the
# workers to exit when their work queues are empty and then waits until the
# threads finish.

_threads_queues=weakref.WeakKeyDictionary()
_shutdown=False

def_python_exit():
global_shutdown
_shutdown=True
items=list(_threads_queues.items())
fort, qinitems:
q.put(None)
fort, qinitems:
t.join()

atexit.register(_python_exit)


class_WorkItem(object):
def__init__(self, future, fn, args, kwargs):
self.future=future
self.fn=fn
self.args=args
self.kwargs=kwargs

defrun(self):
ifnotself.future.set_running_or_notify_cancel():
return

try:
result=self.fn(*self.args, **self.kwargs)
exceptBaseExceptionasexc:
self.future.set_exception(exc)
# Break a reference cycle with the exception 'exc'
self=None
else:
self.future.set_result(result)


def_worker(executor_reference, work_queue, initializer, initargs):
ifinitializerisnotNone:
try:
initializer(*initargs)
exceptBaseException:
_base.LOGGER.critical('Exception in initializer:', exc_info=True)
executor=executor_reference()
ifexecutorisnotNone:
executor._initializer_failed()
return
try:
whileTrue:
work_item=work_queue.get(block=True)
ifwork_itemisnotNone:
work_item.run()
# Delete references to object. See issue16284
delwork_item
continue
executor=executor_reference()
# Exit if:
# - The interpreter is shutting down OR
# - The executor that owns the worker has been collected OR
# - The executor that owns the worker has been shutdown.
if_shutdownorexecutorisNoneorexecutor._shutdown:
# Flag the executor as shutting down as early as possible if it
# is not gc-ed yet.
ifexecutorisnotNone:
executor._shutdown=True
# Notice other workers
work_queue.put(None)
return
delexecutor
exceptBaseException:
_base.LOGGER.critical('Exception in worker', exc_info=True)


classBrokenThreadPool(_base.BrokenExecutor):
"""
 Raised when a worker thread in a ThreadPoolExecutor failed initializing.
 """


classThreadPoolExecutor(_base.Executor):

# Used to assign unique thread names when thread_name_prefix is not supplied.
_counter=itertools.count().__next__

def__init__(self, max_workers=None, thread_name_prefix='',
initializer=None, initargs=()):
"""Initializes a new ThreadPoolExecutor instance.

 Args:
 max_workers: The maximum number of threads that can be used to
 execute the given calls.
 thread_name_prefix: An optional name prefix to give our threads.
 initializer: A callable used to initialize worker threads.
 initargs: A tuple of arguments to pass to the initializer.
 """
ifmax_workersisNone:
# Use this number because ThreadPoolExecutor is often
# used to overlap I/O instead of CPU work.
max_workers= (os.cpu_count() or1) *5
ifmax_workers<=0:
raiseValueError("max_workers must be greater than 0")

ifinitializerisnotNoneandnotcallable(initializer):
raiseTypeError("initializer must be a callable")

self._max_workers=max_workers
self._work_queue=queue.SimpleQueue()
self._threads=set()
self._broken=False
self._shutdown=False
self._shutdown_lock=threading.Lock()
self._thread_name_prefix= (thread_name_prefixor
 ("ThreadPoolExecutor-%d"%self._counter()))
self._initializer=initializer
self._initargs=initargs

defsubmit(*args, **kwargs):
iflen(args) >=2:
self, fn, *args=args
elifnotargs:
raiseTypeError("descriptor 'submit' of 'ThreadPoolExecutor' object "
"needs an argument")
elif'fn'inkwargs:
fn=kwargs.pop('fn')
self, *args=args
else:
raiseTypeError('submit expected at least 1 positional argument, '
'got %d'% (len(args)-1))

withself._shutdown_lock:
ifself._broken:
raiseBrokenThreadPool(self._broken)

ifself._shutdown:
raiseRuntimeError('cannot schedule new futures after shutdown')
if_shutdown:
raiseRuntimeError('cannot schedule new futures after '
'interpreter shutdown')

f=_base.Future()
w=_WorkItem(f, fn, args, kwargs)

self._work_queue.put(w)
self._adjust_thread_count()
returnf
submit.__doc__=_base.Executor.submit.__doc__

def_adjust_thread_count(self):
# When the executor gets lost, the weakref callback will wake up
# the worker threads.
defweakref_cb(_, q=self._work_queue):
q.put(None)
# TODO(bquinlan): Should avoid creating new threads if there are more
# idle threads than items in the work queue.
num_threads=len(self._threads)
ifnum_threads<self._max_workers:
thread_name='%s_%d'% (self._thread_name_prefixorself,
num_threads)
t=threading.Thread(name=thread_name, target=_worker,
args=(weakref.ref(self, weakref_cb),
self._work_queue,
self._initializer,
self._initargs))
t.daemon=True
t.start()
self._threads.add(t)
_threads_queues[t] =self._work_queue

def_initializer_failed(self):
withself._shutdown_lock:
self._broken= ('A thread initializer failed, the thread pool '
'is not usable anymore')
# Drain work queue and mark pending futures failed
whileTrue:
try:
work_item=self._work_queue.get_nowait()
exceptqueue.Empty:
break
ifwork_itemisnotNone:
work_item.future.set_exception(BrokenThreadPool(self._broken))

defshutdown(self, wait=True):
withself._shutdown_lock:
self._shutdown=True
self._work_queue.put(None)
ifwait:
fortinself._threads:
t.join()
shutdown.__doc__=_base.Executor.shutdown.__doc__