diff --git a/Lib/concurrent/futures/__init__.py b/Lib/concurrent/futures/__init__.py index 72de617a5b6..d6ac4b3e0b6 100644 --- a/Lib/concurrent/futures/__init__.py +++ b/Lib/concurrent/futures/__init__.py @@ -17,7 +17,7 @@ wait, as_completed) -__all__ = ( +__all__ = [ 'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED', @@ -31,24 +31,35 @@ 'as_completed', 'ProcessPoolExecutor', 'ThreadPoolExecutor', -) +] + + +try: + import _interpreters +except ImportError: + _interpreters = None + +if _interpreters: + __all__.append('InterpreterPoolExecutor') def __dir__(): - return __all__ + ('__author__', '__doc__') + return __all__ + ['__author__', '__doc__'] def __getattr__(name): - global ProcessPoolExecutor, ThreadPoolExecutor + global ProcessPoolExecutor, ThreadPoolExecutor, InterpreterPoolExecutor if name == 'ProcessPoolExecutor': - from .process import ProcessPoolExecutor as pe - ProcessPoolExecutor = pe - return pe + from .process import ProcessPoolExecutor + return ProcessPoolExecutor if name == 'ThreadPoolExecutor': - from .thread import ThreadPoolExecutor as te - ThreadPoolExecutor = te - return te + from .thread import ThreadPoolExecutor + return ThreadPoolExecutor + + if _interpreters and name == 'InterpreterPoolExecutor': + from .interpreter import InterpreterPoolExecutor + return InterpreterPoolExecutor raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/Lib/concurrent/futures/_base.py b/Lib/concurrent/futures/_base.py index 7d69a5baead..96156f47acc 100644 --- a/Lib/concurrent/futures/_base.py +++ b/Lib/concurrent/futures/_base.py @@ -8,6 +8,8 @@ import threading import time import types +import weakref +from itertools import islice FIRST_COMPLETED = 'FIRST_COMPLETED' FIRST_EXCEPTION = 'FIRST_EXCEPTION' @@ -23,14 +25,6 @@ CANCELLED_AND_NOTIFIED = 'CANCELLED_AND_NOTIFIED' FINISHED = 'FINISHED' -_FUTURE_STATES = [ - PENDING, - RUNNING, - CANCELLED, - CANCELLED_AND_NOTIFIED, - FINISHED -] - _STATE_TO_DESCRIPTION_MAP = { PENDING: "pending", RUNNING: "running", @@ -200,15 +194,15 @@ def as_completed(fs, timeout=None): """An iterator over the given futures that yields each as it completes. Args: - fs: The sequence of Futures (possibly created by different Executors) to - iterate over. - timeout: The maximum number of seconds to wait. If None, then there - is no limit on the wait time. + fs: The sequence of Futures (possibly created by different + Executors) to iterate over. + timeout: The maximum number of seconds to wait. If None, then + there is no limit on the wait time. Returns: - An iterator that yields the given Futures as they complete (finished or - cancelled). If any given Futures are duplicated, they will be returned - once. + An iterator that yields the given Futures as they complete + (finished or cancelled). If any given Futures are duplicated, + they will be returned once. Raises: TimeoutError: If the entire result iterator could not be generated @@ -264,19 +258,20 @@ def wait(fs, timeout=None, return_when=ALL_COMPLETED): """Wait for the futures in the given sequence to complete. Args: - fs: The sequence of Futures (possibly created by different Executors) to - wait upon. - timeout: The maximum number of seconds to wait. If None, then there - is no limit on the wait time. - return_when: Indicates when this function should return. The options - are: + fs: The sequence of Futures (possibly created by different + Executors) to wait upon. + timeout: The maximum number of seconds to wait. If None, then + there is no limit on the wait time. + return_when: Indicates when this function should return. + The options are: FIRST_COMPLETED - Return when any future finishes or is cancelled. FIRST_EXCEPTION - Return when any future finishes by raising an - exception. If no future raises an exception + exception. If no future raises an exception then it is equivalent to ALL_COMPLETED. - ALL_COMPLETED - Return when all futures finish or are cancelled. + ALL_COMPLETED - Return when all futures finish or are + cancelled. Returns: A named 2-tuple of sets. The first set, named 'done', contains the @@ -410,11 +405,12 @@ def add_done_callback(self, fn): Args: fn: A callable that will be called with this future as its only - argument when the future completes or is cancelled. The callable - will always be called by a thread in the same process in which - it was added. If the future has already completed or been - cancelled then the callable will be called immediately. These - callables are called in the order that they were added. + argument when the future completes or is cancelled. The + callable will always be called by a thread in the same + process in which it was added. If the future has already + completed or been cancelled then the callable will be + called immediately. These callables are called in the + order that they were added. """ with self._condition: if self._state not in [CANCELLED, CANCELLED_AND_NOTIFIED, FINISHED]: @@ -429,17 +425,19 @@ def result(self, timeout=None): """Return the result of the call that the future represents. Args: - timeout: The number of seconds to wait for the result if the future - isn't done. If None, then there is no limit on the wait time. + timeout: The number of seconds to wait for the result if the + future isn't done. If None, then there is no limit on the + wait time. Returns: The result of the call that the future represents. Raises: CancelledError: If the future was cancelled. - TimeoutError: If the future didn't finish executing before the given - timeout. - Exception: If the call raised then that exception will be raised. + TimeoutError: If the future didn't finish executing before the + given timeout. + Exception: If the call raised then that exception will be + raised. """ try: with self._condition: @@ -465,17 +463,17 @@ def exception(self, timeout=None): Args: timeout: The number of seconds to wait for the exception if the - future isn't done. If None, then there is no limit on the wait - time. + future isn't done. If None, then there is no limit on the + wait time. Returns: - The exception raised by the call that the future represents or None - if the call completed without raising. + The exception raised by the call that the future represents or + None if the call completed without raising. Raises: CancelledError: If the future was cancelled. - TimeoutError: If the future didn't finish executing before the given - timeout. + TimeoutError: If the future didn't finish executing before the + given timeout. """ with self._condition: @@ -500,22 +498,23 @@ def set_running_or_notify_cancel(self): Should only be used by Executor implementations and unit tests. If the future has been cancelled (cancel() was called and returned - True) then any threads waiting on the future completing (though calls - to as_completed() or wait()) are notified and False is returned. + True) then any threads waiting on the future completing (though + calls to as_completed() or wait()) are notified and False is + returned. If the future was not cancelled then it is put in the running state (future calls to running() will return True) and True is returned. This method should be called by Executor implementations before - executing the work associated with this future. If this method returns - False then the work should not be executed. + executing the work associated with this future. If this method + returns False then the work should not be executed. Returns: False if the Future was cancelled, True otherwise. Raises: - RuntimeError: if this method was already called or if set_result() - or set_exception() was called. + RuntimeError: if this method was already called or if + set_result() or set_exception() was called. """ with self._condition: if self._state == CANCELLED: @@ -572,40 +571,61 @@ class Executor(object): def submit(self, fn, /, *args, **kwargs): """Submits a callable to be executed with the given arguments. - Schedules the callable to be executed as fn(*args, **kwargs) and returns - a Future instance representing the execution of the callable. + Schedules the callable to be executed as fn(*args, **kwargs) and + returns a Future instance representing the execution of the + callable. Returns: A Future representing the given call. """ raise NotImplementedError() - def map(self, fn, *iterables, timeout=None, chunksize=1): + def map(self, fn, *iterables, timeout=None, chunksize=1, buffersize=None): """Returns an iterator equivalent to map(fn, iter). Args: fn: A callable that will take as many arguments as there are passed iterables. - timeout: The maximum number of seconds to wait. If None, then there - is no limit on the wait time. - chunksize: The size of the chunks the iterable will be broken into - before being passed to a child process. This argument is only - used by ProcessPoolExecutor; it is ignored by + timeout: The maximum number of seconds to wait. If None, then + there is no limit on the wait time. + chunksize: The size of the chunks the iterable will be broken + into before being passed to a child process. This argument + is only used by ProcessPoolExecutor; it is ignored by ThreadPoolExecutor. + buffersize: The number of submitted tasks whose results have not + yet been yielded. If the buffer is full, iteration over the + iterables pauses until a result is yielded from the buffer. + If None, all input elements are eagerly collected, and + a task is submitted for each. Returns: - An iterator equivalent to: map(func, *iterables) but the calls may - be evaluated out-of-order. + An iterator equivalent to: map(func, *iterables) but the calls + may be evaluated out-of-order. Raises: - TimeoutError: If the entire result iterator could not be generated - before the given timeout. + TimeoutError: If the entire result iterator could not be + generated before the given timeout. Exception: If fn(*args) raises for any values. """ + if buffersize is not None and not isinstance(buffersize, int): + raise TypeError("buffersize must be an integer or None") + if buffersize is not None and buffersize < 1: + raise ValueError("buffersize must be None or > 0") + if timeout is not None: end_time = timeout + time.monotonic() - fs = [self.submit(fn, *args) for args in zip(*iterables)] + zipped_iterables = zip(*iterables) + if buffersize: + fs = collections.deque( + self.submit(fn, *args) for args in islice(zipped_iterables, buffersize) + ) + else: + fs = [self.submit(fn, *args) for args in zipped_iterables] + + # Use a weak reference to ensure that the executor can be garbage + # collected independently of the result_iterator closure. + executor_weakref = weakref.ref(self) # Yield must be hidden in closure so that the futures are submitted # before the first iterator value is required. @@ -614,6 +634,12 @@ def result_iterator(): # reverse to keep finishing order fs.reverse() while fs: + if ( + buffersize + and (executor := executor_weakref()) + and (args := next(zipped_iterables, None)) + ): + fs.appendleft(executor.submit(fn, *args)) # Careful not to keep a reference to the popped future if timeout is None: yield _result_or_cancel(fs.pop()) @@ -632,8 +658,8 @@ def shutdown(self, wait=True, *, cancel_futures=False): Args: wait: If True then shutdown will not return until all running - futures have finished executing and the resources used by the - executor have been reclaimed. + futures have finished executing and the resources used by + the executor have been reclaimed. cancel_futures: If True then shutdown will cancel all pending futures. Futures that are completed or running will not be cancelled. diff --git a/Lib/concurrent/futures/interpreter.py b/Lib/concurrent/futures/interpreter.py new file mode 100644 index 00000000000..fd3d45144b4 --- /dev/null +++ b/Lib/concurrent/futures/interpreter.py @@ -0,0 +1,123 @@ +"""Implements InterpreterPoolExecutor.""" + +from concurrent import interpreters +import sys +from . import thread as _thread +import traceback + + +def do_call(results, func, args, kwargs): + try: + return func(*args, **kwargs) + except BaseException as exc: + # Send the captured exception out on the results queue, + # but still leave it unhandled for the interpreter to handle. + try: + results.put(exc) + except interpreters.NotShareableError: + # The exception is not shareable. + print('exception is not shareable:', file=sys.stderr) + traceback.print_exception(exc) + results.put(None) + raise # re-raise + + +class WorkerContext(_thread.WorkerContext): + + @classmethod + def prepare(cls, initializer, initargs): + def resolve_task(fn, args, kwargs): + if isinstance(fn, str): + # XXX Circle back to this later. + raise TypeError('scripts not supported') + else: + task = (fn, args, kwargs) + return task + + if initializer is not None: + try: + initdata = resolve_task(initializer, initargs, {}) + except ValueError: + if isinstance(initializer, str) and initargs: + raise ValueError(f'an initializer script does not take args, got {initargs!r}') + raise # re-raise + else: + initdata = None + def create_context(): + return cls(initdata) + return create_context, resolve_task + + def __init__(self, initdata): + self.initdata = initdata + self.interp = None + self.results = None + + def __del__(self): + if self.interp is not None: + self.finalize() + + def initialize(self): + assert self.interp is None, self.interp + self.interp = interpreters.create() + try: + maxsize = 0 + self.results = interpreters.create_queue(maxsize) + + if self.initdata: + self.run(self.initdata) + except BaseException: + self.finalize() + raise # re-raise + + def finalize(self): + interp = self.interp + results = self.results + self.results = None + self.interp = None + if results is not None: + del results + if interp is not None: + interp.close() + + def run(self, task): + try: + return self.interp.call(do_call, self.results, *task) + except interpreters.ExecutionFailed as wrapper: + # Wait for the exception data to show up. + exc = self.results.get() + if exc is None: + # The exception must have been not shareable. + raise # re-raise + raise exc from wrapper + + +class BrokenInterpreterPool(_thread.BrokenThreadPool): + """ + Raised when a worker thread in an InterpreterPoolExecutor failed initializing. + """ + + +class InterpreterPoolExecutor(_thread.ThreadPoolExecutor): + + BROKEN = BrokenInterpreterPool + + @classmethod + def prepare_context(cls, initializer, initargs): + return WorkerContext.prepare(initializer, initargs) + + def __init__(self, max_workers=None, thread_name_prefix='', + initializer=None, initargs=()): + """Initializes a new InterpreterPoolExecutor instance. + + Args: + max_workers: The maximum number of interpreters that can be used + to execute the given calls. + thread_name_prefix: An optional name prefix to give our threads. + initializer: A callable or script used to initialize + each worker interpreter. + initargs: A tuple of arguments to pass to the initializer. + """ + thread_name_prefix = (thread_name_prefix or + (f"InterpreterPoolExecutor-{self._counter()}")) + super().__init__(max_workers, thread_name_prefix, + initializer, initargs) diff --git a/Lib/concurrent/futures/process.py b/Lib/concurrent/futures/process.py index 0dee8303ba2..122be582793 100644 --- a/Lib/concurrent/futures/process.py +++ b/Lib/concurrent/futures/process.py @@ -583,7 +583,7 @@ def _check_system_limits(): raise NotImplementedError(_system_limited) _system_limits_checked = True try: - import multiprocessing.synchronize + import multiprocessing.synchronize # noqa: F401 except ImportError: _system_limited = ( "This Python build lacks multiprocessing.synchronize, usually due " @@ -626,6 +626,14 @@ class BrokenProcessPool(_base.BrokenExecutor): while a future was in the running state. """ +_TERMINATE = "terminate" +_KILL = "kill" + +_SHUTDOWN_CALLBACK_OPERATION = { + _TERMINATE, + _KILL +} + class ProcessPoolExecutor(_base.Executor): def __init__(self, max_workers=None, mp_context=None, @@ -634,19 +642,21 @@ def __init__(self, max_workers=None, mp_context=None, Args: max_workers: The maximum number of processes that can be used to - execute the given calls. If None or not given then as many - worker processes will be created as the machine has processors. - mp_context: A multiprocessing context to launch the workers created - using the multiprocessing.get_context('start method') API. This - object should provide SimpleQueue, Queue and Process. + execute the given calls. If None or not given then as many + worker processes will be created as the machine has + processors. + mp_context: A multiprocessing context to launch the workers + created using the multiprocessing.get_context('start method') + API. This object should provide SimpleQueue, Queue and + Process. initializer: A callable used to initialize worker processes. initargs: A tuple of arguments to pass to the initializer. - max_tasks_per_child: The maximum number of tasks a worker process - can complete before it will exit and be replaced with a fresh - worker process. The default of None means worker process will - live as long as the executor. Requires a non-'fork' mp_context - start method. When given, we default to using 'spawn' if no - mp_context is supplied. + max_tasks_per_child: The maximum number of tasks a worker + process can complete before it will exit and be replaced + with a fresh worker process. The default of None means + worker process will live as long as the executor. Requires + a non-'fork' mp_context start method. When given, we + default to using 'spawn' if no mp_context is supplied. """ _check_system_limits() @@ -810,25 +820,31 @@ def submit(self, fn, /, *args, **kwargs): return f submit.__doc__ = _base.Executor.submit.__doc__ - def map(self, fn, *iterables, timeout=None, chunksize=1): + def map(self, fn, *iterables, timeout=None, chunksize=1, buffersize=None): """Returns an iterator equivalent to map(fn, iter). Args: fn: A callable that will take as many arguments as there are passed iterables. - timeout: The maximum number of seconds to wait. If None, then there - is no limit on the wait time. - chunksize: If greater than one, the iterables will be chopped into - chunks of size chunksize and submitted to the process pool. - If set to one, the items in the list will be sent one at a time. + timeout: The maximum number of seconds to wait. If None, then + there is no limit on the wait time. + chunksize: If greater than one, the iterables will be chopped + into chunks of size chunksize and submitted to the process + pool. If set to one, the items in the list will be sent + one at a time. + buffersize: The number of submitted tasks whose results have not + yet been yielded. If the buffer is full, iteration over the + iterables pauses until a result is yielded from the buffer. + If None, all input elements are eagerly collected, and + a task is submitted for each. Returns: - An iterator equivalent to: map(func, *iterables) but the calls may - be evaluated out-of-order. + An iterator equivalent to: map(func, *iterables) but the calls + may be evaluated out-of-order. Raises: - TimeoutError: If the entire result iterator could not be generated - before the given timeout. + TimeoutError: If the entire result iterator could not be + generated before the given timeout. Exception: If fn(*args) raises for any values. """ if chunksize < 1: @@ -836,7 +852,8 @@ def map(self, fn, *iterables, timeout=None, chunksize=1): results = super().map(partial(_process_chunk, fn), itertools.batched(zip(*iterables), chunksize), - timeout=timeout) + timeout=timeout, + buffersize=buffersize) return _chain_from_iterable_of_lists(results) def shutdown(self, wait=True, *, cancel_futures=False): @@ -860,3 +877,66 @@ def shutdown(self, wait=True, *, cancel_futures=False): self._executor_manager_thread_wakeup = None shutdown.__doc__ = _base.Executor.shutdown.__doc__ + + def _force_shutdown(self, operation): + """Attempts to terminate or kill the executor's workers based off the + given operation. Iterates through all of the current processes and + performs the relevant task if the process is still alive. + + After terminating workers, the pool will be in a broken state + and no longer usable (for instance, new tasks should not be + submitted). + """ + if operation not in _SHUTDOWN_CALLBACK_OPERATION: + raise ValueError(f"Unsupported operation: {operation!r}") + + processes = {} + if self._processes: + processes = self._processes.copy() + + # shutdown will invalidate ._processes, so we copy it right before + # calling. If we waited here, we would deadlock if a process decides not + # to exit. + self.shutdown(wait=False, cancel_futures=True) + + if not processes: + return + + for proc in processes.values(): + try: + if not proc.is_alive(): + continue + except ValueError: + # The process is already exited/closed out. + continue + + try: + if operation == _TERMINATE: + proc.terminate() + elif operation == _KILL: + proc.kill() + except ProcessLookupError: + # The process just ended before our signal + continue + + def terminate_workers(self): + """Attempts to terminate the executor's workers. + Iterates through all of the current worker processes and terminates + each one that is still alive. + + After terminating workers, the pool will be in a broken state + and no longer usable (for instance, new tasks should not be + submitted). + """ + return self._force_shutdown(operation=_TERMINATE) + + def kill_workers(self): + """Attempts to kill the executor's workers. + Iterates through all of the current worker processes and kills + each one that is still alive. + + After killing workers, the pool will be in a broken state + and no longer usable (for instance, new tasks should not be + submitted). + """ + return self._force_shutdown(operation=_KILL) diff --git a/Lib/concurrent/futures/thread.py b/Lib/concurrent/futures/thread.py index 9021dde48ef..909359b6487 100644 --- a/Lib/concurrent/futures/thread.py +++ b/Lib/concurrent/futures/thread.py @@ -44,19 +44,46 @@ def _python_exit(): os.register_at_fork(after_in_child=_threads_queues.clear) +class WorkerContext: + + @classmethod + def prepare(cls, initializer, initargs): + if initializer is not None: + if not callable(initializer): + raise TypeError("initializer must be a callable") + def create_context(): + return cls(initializer, initargs) + def resolve_task(fn, args, kwargs): + return (fn, args, kwargs) + return create_context, resolve_task + + def __init__(self, initializer, initargs): + self.initializer = initializer + self.initargs = initargs + + def initialize(self): + if self.initializer is not None: + self.initializer(*self.initargs) + + def finalize(self): + pass + + def run(self, task): + fn, args, kwargs = task + return fn(*args, **kwargs) + + class _WorkItem: - def __init__(self, future, fn, args, kwargs): + def __init__(self, future, task): self.future = future - self.fn = fn - self.args = args - self.kwargs = kwargs + self.task = task - def run(self): + def run(self, ctx): if not self.future.set_running_or_notify_cancel(): return try: - result = self.fn(*self.args, **self.kwargs) + result = ctx.run(self.task) except BaseException as exc: self.future.set_exception(exc) # Break a reference cycle with the exception 'exc' @@ -67,16 +94,15 @@ def run(self): __class_getitem__ = classmethod(types.GenericAlias) -def _worker(executor_reference, work_queue, initializer, initargs): - if initializer is not None: - try: - initializer(*initargs) - except BaseException: - _base.LOGGER.critical('Exception in initializer:', exc_info=True) - executor = executor_reference() - if executor is not None: - executor._initializer_failed() - return +def _worker(executor_reference, ctx, work_queue): + try: + ctx.initialize() + except BaseException: + _base.LOGGER.critical('Exception in initializer:', exc_info=True) + executor = executor_reference() + if executor is not None: + executor._initializer_failed() + return try: while True: try: @@ -90,7 +116,7 @@ def _worker(executor_reference, work_queue, initializer, initargs): work_item = work_queue.get(block=True) if work_item is not None: - work_item.run() + work_item.run(ctx) # Delete references to object. See GH-60488 del work_item continue @@ -111,6 +137,8 @@ def _worker(executor_reference, work_queue, initializer, initargs): del executor except BaseException: _base.LOGGER.critical('Exception in worker', exc_info=True) + finally: + ctx.finalize() class BrokenThreadPool(_base.BrokenExecutor): @@ -121,11 +149,17 @@ class BrokenThreadPool(_base.BrokenExecutor): class ThreadPoolExecutor(_base.Executor): + BROKEN = BrokenThreadPool + # Used to assign unique thread names when thread_name_prefix is not supplied. _counter = itertools.count().__next__ + @classmethod + def prepare_context(cls, initializer, initargs): + return WorkerContext.prepare(initializer, initargs) + def __init__(self, max_workers=None, thread_name_prefix='', - initializer=None, initargs=()): + initializer=None, initargs=(), **ctxkwargs): """Initializes a new ThreadPoolExecutor instance. Args: @@ -134,6 +168,7 @@ def __init__(self, max_workers=None, thread_name_prefix='', 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. + ctxkwargs: Additional arguments to cls.prepare_context(). """ if max_workers is None: # ThreadPoolExecutor is often used to: @@ -147,8 +182,9 @@ def __init__(self, max_workers=None, thread_name_prefix='', if max_workers <= 0: raise ValueError("max_workers must be greater than 0") - if initializer is not None and not callable(initializer): - raise TypeError("initializer must be a callable") + (self._create_worker_context, + self._resolve_work_item_task, + ) = type(self).prepare_context(initializer, initargs, **ctxkwargs) self._max_workers = max_workers self._work_queue = queue.SimpleQueue() @@ -159,13 +195,11 @@ def __init__(self, max_workers=None, thread_name_prefix='', self._shutdown_lock = threading.Lock() self._thread_name_prefix = (thread_name_prefix or ("ThreadPoolExecutor-%d" % self._counter())) - self._initializer = initializer - self._initargs = initargs def submit(self, fn, /, *args, **kwargs): with self._shutdown_lock, _global_shutdown_lock: if self._broken: - raise BrokenThreadPool(self._broken) + raise self.BROKEN(self._broken) if self._shutdown: raise RuntimeError('cannot schedule new futures after shutdown') @@ -174,7 +208,8 @@ def submit(self, fn, /, *args, **kwargs): 'interpreter shutdown') f = _base.Future() - w = _WorkItem(f, fn, args, kwargs) + task = self._resolve_work_item_task(fn, args, kwargs) + w = _WorkItem(f, task) self._work_queue.put(w) self._adjust_thread_count() @@ -197,9 +232,8 @@ def weakref_cb(_, q=self._work_queue): num_threads) t = threading.Thread(name=thread_name, target=_worker, args=(weakref.ref(self, weakref_cb), - self._work_queue, - self._initializer, - self._initargs)) + self._create_worker_context(), + self._work_queue)) t.start() self._threads.add(t) _threads_queues[t] = self._work_queue @@ -215,7 +249,7 @@ def _initializer_failed(self): except queue.Empty: break if work_item is not None: - work_item.future.set_exception(BrokenThreadPool(self._broken)) + work_item.future.set_exception(self.BROKEN(self._broken)) def shutdown(self, wait=True, *, cancel_futures=False): with self._shutdown_lock: diff --git a/Lib/test/datetimetester.py b/Lib/test/datetimetester.py index b091ced5c29..1b71b71c544 100644 --- a/Lib/test/datetimetester.py +++ b/Lib/test/datetimetester.py @@ -7326,7 +7326,6 @@ def test_update_type_cache(self): """) script_helper.assert_python_ok('-c', script) - @unittest.expectedFailure # TODO: RUSTPYTHON; no concurrent.futures.InterpreterPoolExecutor def test_concurrent_initialization_subinterpreter(self): # gh-136421: Concurrent initialization of _datetime across multiple # interpreters wasn't thread-safe due to its static types. diff --git a/Lib/test/test_concurrent_futures/executor.py b/Lib/test/test_concurrent_futures/executor.py index 7442d3bee52..6a26d695b0b 100644 --- a/Lib/test/test_concurrent_futures/executor.py +++ b/Lib/test/test_concurrent_futures/executor.py @@ -1,8 +1,10 @@ +import itertools import threading import time import unittest import weakref from concurrent import futures +from operator import add from test import support from test.support import Py_GIL_DISABLED @@ -85,6 +87,66 @@ def test_map_timeout(self): self.assertEqual([None, None], results) + def test_map_buffersize_type_validation(self): + for buffersize in ("foo", 2.0): + with self.subTest(buffersize=buffersize): + with self.assertRaisesRegex( + TypeError, + "buffersize must be an integer or None", + ): + self.executor.map(str, range(4), buffersize=buffersize) + + + def test_map_buffersize_value_validation(self): + for buffersize in (0, -1): + with self.subTest(buffersize=buffersize): + with self.assertRaisesRegex( + ValueError, + "buffersize must be None or > 0", + ): + self.executor.map(str, range(4), buffersize=buffersize) + + + def test_map_buffersize(self): + ints = range(4) + for buffersize in (1, 2, len(ints), len(ints) * 2): + with self.subTest(buffersize=buffersize): + res = self.executor.map(str, ints, buffersize=buffersize) + self.assertListEqual(list(res), ["0", "1", "2", "3"]) + + + def test_map_buffersize_on_multiple_iterables(self): + ints = range(4) + for buffersize in (1, 2, len(ints), len(ints) * 2): + with self.subTest(buffersize=buffersize): + res = self.executor.map(add, ints, ints, buffersize=buffersize) + self.assertListEqual(list(res), [0, 2, 4, 6]) + + + def test_map_buffersize_on_infinite_iterable(self): + res = self.executor.map(str, itertools.count(), buffersize=2) + self.assertEqual(next(res, None), "0") + self.assertEqual(next(res, None), "1") + self.assertEqual(next(res, None), "2") + + + def test_map_buffersize_on_empty_iterable(self): + res = self.executor.map(str, [], buffersize=2) + self.assertIsNone(next(res, None)) + + + def test_map_buffersize_when_buffer_is_full(self): + ints = iter(range(4)) + buffersize = 2 + self.executor.map(str, ints, buffersize=buffersize) + self.executor.shutdown(wait=True) # wait for tasks to complete + self.assertEqual( + next(ints), + buffersize, + msg="should have fetched only `buffersize` elements from `ints`.", + ) + + def test_shutdown_race_issue12456(self): # Issue #12456: race condition at shutdown where trying to post a # sentinel in the call queue blocks (the queue is full while processes diff --git a/Lib/test/test_concurrent_futures/test_interpreter_pool.py b/Lib/test/test_concurrent_futures/test_interpreter_pool.py new file mode 100644 index 00000000000..b8f1a0f3a29 --- /dev/null +++ b/Lib/test/test_concurrent_futures/test_interpreter_pool.py @@ -0,0 +1,594 @@ +import _thread +import asyncio +import contextlib +import io +import os +import subprocess +import sys +import textwrap +import time +import unittest +from concurrent.futures.interpreter import BrokenInterpreterPool +from concurrent import interpreters +from concurrent.interpreters import _queues as queues +import _interpreters +from test import support +from test.support import os_helper +from test.support import script_helper +import test.test_asyncio.utils as testasyncio_utils + +from .executor import ExecutorTest, mul +from .util import BaseTestCase, InterpreterPoolMixin, setup_module + + +WINDOWS = sys.platform.startswith('win') + + +@contextlib.contextmanager +def nonblocking(fd): + blocking = os.get_blocking(fd) + if blocking: + os.set_blocking(fd, False) + try: + yield + finally: + if blocking: + os.set_blocking(fd, blocking) + + +def read_file_with_timeout(fd, nbytes, timeout): + with nonblocking(fd): + end = time.time() + timeout + try: + return os.read(fd, nbytes) + except BlockingIOError: + pass + while time.time() < end: + try: + return os.read(fd, nbytes) + except BlockingIOError: + continue + else: + raise TimeoutError('nothing to read') + + +if not WINDOWS: + import select + def read_file_with_timeout(fd, nbytes, timeout): + r, _, _ = select.select([fd], [], [], timeout) + if fd not in r: + raise TimeoutError('nothing to read') + return os.read(fd, nbytes) + + +def noop(): + pass + + +def write_msg(fd, msg): + import os + os.write(fd, msg + b'\0') + + +def read_msg(fd, timeout=10.0): + msg = b'' + ch = read_file_with_timeout(fd, 1, timeout) + while ch != b'\0': + msg += ch + ch = os.read(fd, 1) + return msg + + +def get_current_name(): + return __name__ + + +def fail(exctype, msg=None): + raise exctype(msg) + + +def get_current_interpid(*extra): + interpid, _ = _interpreters.get_current() + return (interpid, *extra) + + +class InterpretersMixin(InterpreterPoolMixin): + + def pipe(self): + r, w = os.pipe() + self.addCleanup(lambda: os.close(r)) + self.addCleanup(lambda: os.close(w)) + return r, w + + +class PickleShenanigans: + """Succeeds with pickle.dumps(), but fails with pickle.loads()""" + def __init__(self, value): + if value == 1: + raise RuntimeError("gotcha") + + def __reduce__(self): + return (self.__class__, (1,)) + + +class InterpreterPoolExecutorTest( + InterpretersMixin, ExecutorTest, BaseTestCase): + + @unittest.expectedFailure + def test_init_script(self): + msg1 = b'step: init' + msg2 = b'step: run' + r, w = self.pipe() + initscript = f""" + import os + msg = {msg2!r} + os.write({w}, {msg1!r} + b'\\0') + """ + script = f""" + os.write({w}, msg + b'\\0') + """ + os.write(w, b'\0') + + executor = self.executor_type(initializer=initscript) + before_init = os.read(r, 100) + fut = executor.submit(script) + after_init = read_msg(r) + fut.result() + after_run = read_msg(r) + + self.assertEqual(before_init, b'\0') + self.assertEqual(after_init, msg1) + self.assertEqual(after_run, msg2) + + @unittest.expectedFailure + def test_init_script_args(self): + with self.assertRaises(ValueError): + self.executor_type(initializer='pass', initargs=('spam',)) + + @unittest.skipIf(sys.platform == "win32", "TODO: RUSTPYTHON; missing os.get_blocking/set_blocking leaves a worker running during cleanup") + def test_init_func(self): + msg = b'step: init' + r, w = self.pipe() + os.write(w, b'\0') + + executor = self.executor_type( + initializer=write_msg, initargs=(w, msg)) + before = os.read(r, 100) + executor.submit(mul, 10, 10) + after = read_msg(r) + + self.assertEqual(before, b'\0') + self.assertEqual(after, msg) + + def test_init_with___main___global(self): + # See https://github.com/python/cpython/pull/133957#issuecomment-2927415311. + text = """if True: + from concurrent.futures import InterpreterPoolExecutor + + INITIALIZER_STATUS = 'uninitialized' + + def init(x): + global INITIALIZER_STATUS + INITIALIZER_STATUS = x + INITIALIZER_STATUS + + def get_init_status(): + return INITIALIZER_STATUS + + if __name__ == "__main__": + exe = InterpreterPoolExecutor(initializer=init, + initargs=('initialized',)) + fut = exe.submit(get_init_status) + print(fut.result()) # 'initialized' + exe.shutdown(wait=True) + print(INITIALIZER_STATUS) # 'uninitialized' + """ + with os_helper.temp_dir() as tempdir: + filename = script_helper.make_script(tempdir, 'my-script', text) + res = script_helper.assert_python_ok(filename) + stdout = res.out.decode('utf-8').strip() + self.assertEqual(stdout.splitlines(), [ + 'initialized', + 'uninitialized', + ]) + + def test_init_closure(self): + count = 0 + def init1(): + assert count == 0, count + def init2(): + nonlocal count + count += 1 + + with contextlib.redirect_stderr(io.StringIO()) as stderr: + with self.executor_type(initializer=init1) as executor: + fut = executor.submit(lambda: None) + self.assertIn('NotShareableError', stderr.getvalue()) + with self.assertRaises(BrokenInterpreterPool): + fut.result() + + with contextlib.redirect_stderr(io.StringIO()) as stderr: + with self.executor_type(initializer=init2) as executor: + fut = executor.submit(lambda: None) + self.assertIn('NotShareableError', stderr.getvalue()) + with self.assertRaises(BrokenInterpreterPool): + fut.result() + + def test_init_instance_method(self): + class Spam: + def initializer(self): + raise NotImplementedError + spam = Spam() + + with contextlib.redirect_stderr(io.StringIO()) as stderr: + with self.executor_type(initializer=spam.initializer) as executor: + fut = executor.submit(lambda: None) + self.assertIn('NotShareableError', stderr.getvalue()) + with self.assertRaises(BrokenInterpreterPool): + fut.result() + + @unittest.expectedFailure + def test_init_exception_in_script(self): + executor = self.executor_type(initializer='raise Exception("spam")') + with executor: + with contextlib.redirect_stderr(io.StringIO()) as stderr: + fut = executor.submit('pass') + with self.assertRaises(BrokenInterpreterPool): + fut.result() + stderr = stderr.getvalue() + self.assertIn('ExecutionFailed: Exception: spam', stderr) + self.assertIn('Uncaught in the interpreter:', stderr) + self.assertIn('The above exception was the direct cause of the following exception:', + stderr) + + def test_init_exception_in_func(self): + executor = self.executor_type(initializer=fail, + initargs=(Exception, 'spam')) + with executor: + with contextlib.redirect_stderr(io.StringIO()) as stderr: + fut = executor.submit(noop) + with self.assertRaises(BrokenInterpreterPool): + fut.result() + stderr = stderr.getvalue() + self.assertIn('ExecutionFailed: Exception: spam', stderr) + self.assertIn('Uncaught in the interpreter:', stderr) + + @unittest.expectedFailure + def test_submit_script(self): + msg = b'spam' + r, w = self.pipe() + script = f""" + import os + os.write({w}, __name__.encode('utf-8') + b'\\0') + """ + executor = self.executor_type() + + fut = executor.submit(script) + res = fut.result() + after = read_msg(r) + + self.assertEqual(after, b'__main__') + self.assertIs(res, None) + + def test_submit_closure(self): + spam = True + def task1(): + return spam + def task2(): + nonlocal spam + spam += 1 + return spam + + executor = self.executor_type() + + fut = executor.submit(task1) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() + + fut = executor.submit(task2) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() + + def test_submit_local_instance(self): + class Spam: + def __init__(self): + self.value = True + + executor = self.executor_type() + fut = executor.submit(Spam) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() + + def test_submit_instance_method(self): + class Spam: + def run(self): + return True + spam = Spam() + + executor = self.executor_type() + fut = executor.submit(spam.run) + with self.assertRaises(_interpreters.NotShareableError): + fut.result() + + def test_submit_func_globals(self): + executor = self.executor_type() + fut = executor.submit(get_current_name) + name = fut.result() + + self.assertEqual(name, __name__) + self.assertNotEqual(name, '__main__') + + @unittest.expectedFailure + def test_submit_exception_in_script(self): + # Scripts are not supported currently. + fut = self.executor.submit('raise Exception("spam")') + with self.assertRaises(Exception) as captured: + fut.result() + self.assertIs(type(captured.exception), Exception) + self.assertEqual(str(captured.exception), 'spam') + cause = captured.exception.__cause__ + self.assertIs(type(cause), interpreters.ExecutionFailed) + for attr in ('__name__', '__qualname__', '__module__'): + self.assertEqual(getattr(cause.excinfo.type, attr), + getattr(Exception, attr)) + self.assertEqual(cause.excinfo.msg, 'spam') + + def test_submit_exception_in_func(self): + fut = self.executor.submit(fail, Exception, 'spam') + with self.assertRaises(Exception) as captured: + fut.result() + self.assertIs(type(captured.exception), Exception) + self.assertEqual(str(captured.exception), 'spam') + cause = captured.exception.__cause__ + self.assertIs(type(cause), interpreters.ExecutionFailed) + for attr in ('__name__', '__qualname__', '__module__'): + self.assertEqual(getattr(cause.excinfo.type, attr), + getattr(Exception, attr)) + self.assertEqual(cause.excinfo.msg, 'spam') + + def test_saturation(self): + blocker = queues.create() + executor = self.executor_type(4) + + for i in range(15 * executor._max_workers): + executor.submit(blocker.get) + self.assertEqual(len(executor._threads), executor._max_workers) + for i in range(15 * executor._max_workers): + blocker.put_nowait(None) + executor.shutdown(wait=True) + + def test_blocking(self): + # There is no guarantee that a worker will be created for every + # submitted task. That's because there's a race between: + # + # * a new worker thread, created when task A was just submitted, + # becoming non-idle when it picks up task A + # * after task B is added to the queue, a new worker thread + # is started only if there are no idle workers + # (the check in ThreadPoolExecutor._adjust_thread_count()) + # + # That means we must not block waiting for *all* tasks to report + # "ready" before we unblock the known-ready workers. + ready = queues.create() + blocker = queues.create() + + def run(taskid, ready, blocker): + # There can't be any globals here. + ready.put_nowait(taskid) + blocker.get() # blocking + + numtasks = 10 + futures = [] + with self.executor_type() as executor: + # Request the jobs. + for i in range(numtasks): + fut = executor.submit(run, i, ready, blocker) + futures.append(fut) + pending = numtasks + while pending > 0: + # Wait for any to be ready. + done = 0 + for _ in range(pending): + try: + ready.get(timeout=1) # blocking + except interpreters.QueueEmpty: + pass + else: + done += 1 + pending -= done + # Unblock the workers. + for _ in range(done): + blocker.put_nowait(None) + + def test_blocking_with_limited_workers(self): + # This is essentially the same as test_blocking, + # but we explicitly force a limited number of workers, + # instead of it happening implicitly sometimes due to a race. + ready = queues.create() + blocker = queues.create() + + def run(taskid, ready, blocker): + # There can't be any globals here. + ready.put_nowait(taskid) + blocker.get() # blocking + + numtasks = 10 + futures = [] + with self.executor_type(4) as executor: + # Request the jobs. + for i in range(numtasks): + fut = executor.submit(run, i, ready, blocker) + futures.append(fut) + pending = numtasks + while pending > 0: + # Wait for any to be ready. + done = 0 + for _ in range(pending): + try: + ready.get(timeout=1) # blocking + except interpreters.QueueEmpty: + pass + else: + done += 1 + pending -= done + # Unblock the workers. + for _ in range(done): + blocker.put_nowait(None) + + @support.requires_gil_enabled("gh-117344: test is flaky without the GIL") + def test_idle_thread_reuse(self): + executor = self.executor_type() + executor.submit(mul, 21, 2).result() + executor.submit(mul, 6, 7).result() + executor.submit(mul, 3, 14).result() + self.assertEqual(len(executor._threads), 1) + executor.shutdown(wait=True) + + def test_pickle_errors_propagate(self): + # GH-125864: Pickle errors happen before the script tries to execute, + # so the queue used to wait infinitely. + fut = self.executor.submit(PickleShenanigans(0)) + expected = interpreters.NotShareableError + with self.assertRaisesRegex(expected, 'args not shareable') as cm: + fut.result() + self.assertRegex(str(cm.exception.__cause__), 'unpickled') + + def test_no_stale_references(self): + # Weak references don't cross between interpreters. + raise unittest.SkipTest('not applicable') + + def test_free_reference(self): + # Weak references don't cross between interpreters. + raise unittest.SkipTest('not applicable') + + @support.requires_subprocess() + def test_import_interpreter_pool_executor(self): + # Test the import behavior normally if _interpreters is unavailable. + code = textwrap.dedent(""" + import sys + # Set it to None to emulate the case when _interpreter is unavailable. + sys.modules['_interpreters'] = None + from concurrent import futures + + try: + futures.InterpreterPoolExecutor + except AttributeError: + pass + else: + print('AttributeError not raised!', file=sys.stderr) + sys.exit(1) + + try: + from concurrent.futures import InterpreterPoolExecutor + except ImportError: + pass + else: + print('ImportError not raised!', file=sys.stderr) + sys.exit(1) + + from concurrent.futures import * + + if 'InterpreterPoolExecutor' in globals(): + print('InterpreterPoolExecutor should not be imported!', + file=sys.stderr) + sys.exit(1) + """) + + cmd = [sys.executable, '-c', code] + p = subprocess.run(cmd, capture_output=True) + self.assertEqual(p.returncode, 0, p.stderr.decode()) + self.assertEqual(p.stdout.decode(), '') + self.assertEqual(p.stderr.decode(), '') + + def test_thread_name_prefix(self): + self.assertStartsWith(self.executor._thread_name_prefix, + "InterpreterPoolExecutor-") + + @unittest.skipUnless(hasattr(_thread, '_get_name'), "missing _thread._get_name") + def test_thread_name_prefix_with_thread_get_name(self): + def get_thread_name(): + import _thread + return _thread._get_name() + + # Some platforms (Linux) are using 16 bytes to store the thread name, + # so only compare the first 15 bytes (without the trailing \n). + self.assertStartsWith(self.executor.submit(get_thread_name).result(), + "InterpreterPoolExecutor-"[:15]) + +class AsyncioTest(InterpretersMixin, testasyncio_utils.TestCase): + + @classmethod + def setUpClass(cls): + # Most uses of asyncio will implicitly call set_event_loop_policy() + # with the default policy if a policy hasn't been set already. + # If that happens in a test, like here, we'll end up with a failure + # when --fail-env-changed is used. That's why the other tests that + # use asyncio are careful to set the policy back to None and why + # we're careful to do so here. We also validate that no other + # tests left a policy in place, just in case. + policy = support.maybe_get_event_loop_policy() + assert policy is None, policy + cls.addClassCleanup(lambda: asyncio.events._set_event_loop_policy(None)) + + def setUp(self): + super().setUp() + self.loop = asyncio.new_event_loop() + self.set_event_loop(self.loop) + + self.executor = self.executor_type() + self.addCleanup(lambda: self.executor.shutdown()) + + def tearDown(self): + if not self.loop.is_closed(): + testasyncio_utils.run_briefly(self.loop) + + self.doCleanups() + support.gc_collect() + super().tearDown() + + def test_run_in_executor(self): + unexpected, _ = _interpreters.get_current() + + func = get_current_interpid + fut = self.loop.run_in_executor(self.executor, func, 'yo') + interpid, res = self.loop.run_until_complete(fut) + + self.assertEqual(res, 'yo') + self.assertNotEqual(interpid, unexpected) + + def test_run_in_executor_cancel(self): + executor = self.executor_type() + + called = False + + def patched_call_soon(*args): + nonlocal called + called = True + + func = time.sleep + fut = self.loop.run_in_executor(self.executor, func, 0.05) + fut.cancel() + self.loop.run_until_complete( + self.loop.shutdown_default_executor()) + self.loop.close() + self.loop.call_soon = patched_call_soon + self.loop.call_soon_threadsafe = patched_call_soon + time.sleep(0.4) + self.assertFalse(called) + + def test_default_executor(self): + unexpected, _ = _interpreters.get_current() + + self.loop.set_default_executor(self.executor) + fut = self.loop.run_in_executor(None, get_current_interpid) + interpid, = self.loop.run_until_complete(fut) + + self.assertNotEqual(interpid, unexpected) + + +def setUpModule(): + setup_module() + + +if __name__ == "__main__": + unittest.main() diff --git a/Lib/test/test_concurrent_futures/test_process_pool.py b/Lib/test/test_concurrent_futures/test_process_pool.py index 5d4e9677f5c..e99abaf6523 100644 --- a/Lib/test/test_concurrent_futures/test_process_pool.py +++ b/Lib/test/test_concurrent_futures/test_process_pool.py @@ -1,13 +1,16 @@ import os +import queue import sys import threading import time import unittest +import unittest.mock from concurrent import futures from concurrent.futures.process import BrokenProcessPool from test import support from test.support import hashlib_helper +from test.test_importlib.metadata.fixtures import parameterize from .executor import ExecutorTest, mul from .util import ( @@ -23,6 +26,23 @@ def __del__(self): self.event.set() +TERMINATE_WORKERS = futures.ProcessPoolExecutor.terminate_workers.__name__ +KILL_WORKERS = futures.ProcessPoolExecutor.kill_workers.__name__ +FORCE_SHUTDOWN_PARAMS = [ + dict(function_name=TERMINATE_WORKERS), + dict(function_name=KILL_WORKERS), +] + +def _put_wait_put(queue, event): + """ Used as part of test_terminate_workers """ + queue.put('started') + event.wait() + + # We should never get here since the event will not get set + queue.put('finished') + + + class ProcessPoolExecutorTest(ExecutorTest): @unittest.skipUnless(sys.platform=='win32', 'Windows-only process limit') @@ -219,6 +239,108 @@ def mock_start_new_thread(func, *args, **kwargs): executor.shutdown() + def test_terminate_workers(self): + mock_fn = unittest.mock.Mock() + with self.executor_type(max_workers=1) as executor: + executor._force_shutdown = mock_fn + executor.terminate_workers() + + mock_fn.assert_called_once_with(operation=futures.process._TERMINATE) + + def test_kill_workers(self): + mock_fn = unittest.mock.Mock() + with self.executor_type(max_workers=1) as executor: + executor._force_shutdown = mock_fn + executor.kill_workers() + + mock_fn.assert_called_once_with(operation=futures.process._KILL) + + def test_force_shutdown_workers_invalid_op(self): + with self.executor_type(max_workers=1) as executor: + self.assertRaises(ValueError, + executor._force_shutdown, + operation='invalid operation'), + + @parameterize(*FORCE_SHUTDOWN_PARAMS) + def test_force_shutdown_workers(self, function_name): + manager = self.get_context().Manager() + q = manager.Queue() + e = manager.Event() + + with self.executor_type(max_workers=1) as executor: + executor.submit(_put_wait_put, q, e) + + # We should get started, but not finished since we'll terminate the + # workers just after and never set the event. + self.assertEqual(q.get(timeout=support.SHORT_TIMEOUT), 'started') + + worker_process = list(executor._processes.values())[0] + + Mock = unittest.mock.Mock + worker_process.terminate = Mock(wraps=worker_process.terminate) + worker_process.kill = Mock(wraps=worker_process.kill) + + getattr(executor, function_name)() + worker_process.join() + + if function_name == TERMINATE_WORKERS: + worker_process.terminate.assert_called() + elif function_name == KILL_WORKERS: + worker_process.kill.assert_called() + else: + self.fail(f"Unknown operation: {function_name}") + + self.assertRaises(queue.Empty, q.get, timeout=0.01) + + @parameterize(*FORCE_SHUTDOWN_PARAMS) + def test_force_shutdown_workers_dead_workers(self, function_name): + with self.executor_type(max_workers=1) as executor: + future = executor.submit(os._exit, 1) + self.assertRaises(BrokenProcessPool, future.result) + + # even though the pool is broken, this shouldn't raise + getattr(executor, function_name)() + + @parameterize(*FORCE_SHUTDOWN_PARAMS) + def test_force_shutdown_workers_not_started_yet(self, function_name): + ctx = self.get_context() + with unittest.mock.patch.object(ctx, 'Process') as mock_process: + with self.executor_type(max_workers=1, mp_context=ctx) as executor: + # The worker has not been started yet, terminate/kill_workers + # should basically no-op + getattr(executor, function_name)() + + mock_process.return_value.kill.assert_not_called() + mock_process.return_value.terminate.assert_not_called() + + @parameterize(*FORCE_SHUTDOWN_PARAMS) + def test_force_shutdown_workers_stops_pool(self, function_name): + with self.executor_type(max_workers=1) as executor: + task = executor.submit(time.sleep, 0) + self.assertIsNone(task.result()) + + worker_process = list(executor._processes.values())[0] + getattr(executor, function_name)() + + self.assertRaises(RuntimeError, executor.submit, time.sleep, 0) + + # A signal sent, is not a signal reacted to. + # So wait a moment here for the process to die. + # If we don't, every once in a while we may get an ENV CHANGE + # error since the process would be alive immediately after the + # test run.. and die a moment later. + worker_process.join(support.SHORT_TIMEOUT) + + # Oddly enough, even though join completes, sometimes it takes a + # moment for the process to actually be marked as dead. + # ... that seems a bit buggy. + # We need it dead before ending the test to ensure it doesn't + # get marked as an ENV CHANGE due to living child process. + for _ in support.sleeping_retry(support.SHORT_TIMEOUT): + if not worker_process.is_alive(): + break + + create_executor_tests(globals(), ProcessPoolExecutorTest, executor_mixins=(ProcessPoolForkMixin, ProcessPoolForkserverMixin, diff --git a/Lib/test/test_concurrent_futures/util.py b/Lib/test/test_concurrent_futures/util.py index e85ef3b1c91..1ea5a16d8b9 100644 --- a/Lib/test/test_concurrent_futures/util.py +++ b/Lib/test/test_concurrent_futures/util.py @@ -83,6 +83,14 @@ def create_event(self): return threading.Event() +@support.skip_if_sanitizer("gh-129824: data races in InterpreterPool tests", thread=True) +class InterpreterPoolMixin(ExecutorMixin): + executor_type = futures.InterpreterPoolExecutor + + def create_event(self): + self.skipTest("InterpreterPoolExecutor doesn't support events") + + class ProcessPoolForkMixin(ExecutorMixin): executor_type = futures.ProcessPoolExecutor ctx = "fork" diff --git a/Lib/test/test_pyrepl/test_pyrepl.py b/Lib/test/test_pyrepl/test_pyrepl.py index 887d11acdb1..4b7579d821c 100644 --- a/Lib/test/test_pyrepl/test_pyrepl.py +++ b/Lib/test/test_pyrepl/test_pyrepl.py @@ -1039,7 +1039,6 @@ def test_builtin_completion_top_level(self): output = reader.readline() self.assertEqual(output, expected) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_relative_import_completions(self): cases = ( (None, "from .readl\t\n", "from .readl"), diff --git a/crates/vm/src/vm/crossinterp.rs b/crates/vm/src/vm/crossinterp.rs index f2b91dcc8c5..2b799bf0031 100644 --- a/crates/vm/src/vm/crossinterp.rs +++ b/crates/vm/src/vm/crossinterp.rs @@ -2,6 +2,7 @@ //! //! Interpreters do not share `PyObject` graphs. Shareable values are converted //! to an interpreter-neutral payload and rebuilt in the destination. +// cspell:ignore unpickler #[cfg(not(feature = "threading"))] use crate::protocol::BufferDescriptor; @@ -140,7 +141,14 @@ pub enum SharedValue { /// Marshalled code of a stateless function; rebuilt against `__main__`. Function(Vec), /// `pickle.dumps` output, used by [`Fallback::Full`]. - Pickled(Vec), + Pickled(PickledData), +} + +/// Pickled bytes and the originating script, without interpreter-owned references. +#[derive(Debug, Clone)] +pub struct PickledData { + bytes: Vec, + mainfile: Option, } impl SharedValue { @@ -360,20 +368,144 @@ fn memoryview_from_buffer(buffer: PyBuffer, vm: &VirtualMachine) -> PyResult { Ok(mv.into_pyobject(vm)) } -fn pickle_dumps(obj: &PyObject, vm: &VirtualMachine) -> PyResult> { +fn pickle_dumps(obj: &PyObject, vm: &VirtualMachine) -> PyResult { let dumps = vm.import("pickle", 0)?.get_attr("dumps", vm)?; let bytes = dumps.call((obj.to_owned(),), vm)?; let bytes = bytes .downcast::() .map_err(|_| vm.new_type_error("pickle.dumps() did not return bytes"))?; - Ok(bytes.as_bytes().to_vec()) + // Missing or non-string __file__ does not prevent ordinary pickling. + let mainfile = vm + .main_namespace() + .ok() + .and_then(|ns| ns.get_item_opt("__file__", vm).ok().flatten()) + .and_then(|file| { + file.downcast_ref::() + .and_then(|s| s.to_str().map(str::to_owned)) + }); + Ok(PickledData { + bytes: bytes.as_bytes().to_vec(), + mainfile, + }) } -fn pickle_loads(data: &[u8], vm: &VirtualMachine) -> PyResult { +fn pickle_loads(data: &PickledData, vm: &VirtualMachine) -> PyResult { let loads = vm.import("pickle", 0)?.get_attr("loads", vm)?; - loads - .call((vm.ctx.new_bytes(data.to_vec()),), vm) - .map_err(|cause| not_shareable_error_from(vm, "object could not be unpickled", cause)) + let bytes = vm.ctx.new_bytes(data.bytes.clone()); + let cause = match loads.call((bytes.clone(),), vm) { + Ok(obj) => return Ok(obj), + Err(cause) => cause, + }; + let missing_main_attr = cause.fast_isinstance(vm.ctx.exceptions.attribute_error) + && cause + .args() + .as_slice() + .first() + .and_then(|arg| arg.downcast_ref::()) + .is_some_and(|msg| { + msg.as_bytes() + .starts_with(b"module '__main__' has no attribute '") + }); + if missing_main_attr && let Some(mainfile) = &data.mainfile { + let retry = || -> PyResult { + let isolated = isolated_main(mainfile, vm)?; + let file = vm + .import("io", 0)? + .get_attr("BytesIO", vm)? + .call((bytes,), vm)?; + let unpickler = vm + .import("pickle", 0)? + .get_attr("Unpickler", vm)? + .call((file,), vm)?; + let original = unpickler.get_attr("find_class", vm)?; + let state = vm.ctx.new_tuple(vec![isolated, original]); + const FIND_CLASS: crate::function::PyMethodDef = + crate::function::PyMethodDef::new_const( + "find_class", + find_main_class, + crate::function::PyMethodFlags::empty(), + crate::function::ItemDoc::NONE, + ); + unpickler.set_attr( + "find_class", + FIND_CLASS.build_bound_function(&vm.ctx, state.into()), + vm, + )?; + // Resolve globals on this unpickler only: replacing sys.modules + // would let concurrent consumers observe or restore the wrong module. + let _clear = scopeguard::guard((), |()| { + let _ = unpickler.del_attr("find_class", vm); + }); + vm.call_method(&unpickler, "load", ()) + }; + if let Ok(obj) = retry() { + return Ok(obj); + } + } + // Like CPython, retain the original unpickling error if the retry fails. + Err(not_shareable_error_from( + vm, + "object could not be unpickled", + cause, + )) +} + +fn find_main_class( + state: PyRef, + module: PyRef, + name: PyRef, + vm: &VirtualMachine, +) -> PyResult { + let [isolated, original] = state.as_slice() else { + unreachable!() + }; + if module.as_bytes() != b"__main__" { + return original.call((module, name), vm); + } + vm.audit("pickle.find_class", || (module, name.clone()))?; + let mut value = isolated.clone(); + if let Some(name) = name.to_str() { + for part in name.split('.') { + value = value.get_attr(&vm.ctx.new_str(part), vm)?; + } + Ok(value) + } else { + value.get_attr(&name, vm) + } +} + +fn isolated_main(mainfile: &str, vm: &VirtualMachine) -> PyResult { + let module = vm.import("_interpreters", 0)?; + let state = module + .dict() + .ok_or_else(|| vm.new_runtime_error("missing _interpreters namespace"))?; + if let Some(loaded) = state.get_item_opt("_cached_main", vm)? { + return Ok(loaded); + } + // Use the import machinery's reentrant, deadlock-aware per-module lock + // while executing the script, including detaching threads that wait for it. + let manager = vm + .import("_frozen_importlib", 0)? + .get_attr("_ModuleLockManager", vm)? + .call(("",), vm)?; + vm.call_method(&manager, "__enter__", ())?; + let _unlock = scopeguard::guard((), |()| { + let _ = vm.call_method( + &manager, + "__exit__", + (vm.ctx.none(), vm.ctx.none(), vm.ctx.none()), + ); + }); + if let Some(loaded) = state.get_item_opt("_cached_main", vm)? { + return Ok(loaded); + } + let loaded = vm.new_module("__main__", vm.ctx.new_dict(), None); + let run_path = vm.import("runpy", 0)?.get_attr("run_path", vm)?; + // Avoid executing the script's `if __name__ == '__main__'` entry point. + let namespace = run_path.call((mainfile, vm.ctx.none(), ""), vm)?; + loaded.dict().merge_object(namespace, vm)?; + // The native module owns this cache so normal module teardown releases it. + state.setdefault(vm.ctx.new_str("_cached_main").into(), loaded.into(), vm) } fn not_shareable_error(vm: &VirtualMachine, msg: impl Into) -> PyBaseExceptionRef { diff --git a/crates/vm/src/vm/interpreter.rs b/crates/vm/src/vm/interpreter.rs index 93f60edf058..512525519d3 100644 --- a/crates/vm/src/vm/interpreter.rs +++ b/crates/vm/src/vm/interpreter.rs @@ -162,6 +162,8 @@ where audit_hooks: PyMutex::default(), codec_registry, finalizing: AtomicBool::new(false), + #[cfg(feature = "threading")] + finalizing_thread_ident: AtomicCell::new(0), warnings, override_frozen_modules: AtomicCell::new(0), before_forkers: PyMutex::default(), @@ -683,7 +685,12 @@ impl Interpreter { // running). Their `_ThreadHandle` stays not-done so `is_alive()` // is still true for `join()` during the GC that follows. #[cfg(feature = "threading")] - vm.state.stop_the_world.stop_the_world(&vm.state); + { + vm.state.stop_the_world.stop_the_world(&vm.state); + vm.state + .finalizing_thread_ident + .store(crate::stdlib::_thread::get_ident()); + } vm.state.finalizing.store(true, Ordering::Release); #[cfg(feature = "threading")] { diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index dbdd5ba7b8d..e442e01d6b6 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -825,6 +825,9 @@ pub struct PyGlobalState { pub(crate) audit_hooks: PyMutex>, pub codec_registry: CodecsRegistry, pub finalizing: AtomicBool, + /// The thread performing finalization, which need not be the process main thread. + #[cfg(feature = "threading")] + pub(crate) finalizing_thread_ident: AtomicCell, pub warnings: WarningsState, pub override_frozen_modules: AtomicCell, pub before_forkers: PyMutex>, @@ -3563,7 +3566,9 @@ impl VirtualMachine { /// platforms where signals are not supported. pub fn check_signals(&self) -> PyResult<()> { #[cfg(feature = "threading")] - if self.state.finalizing.load(Ordering::Acquire) && !self.is_main_thread() { + if self.state.finalizing.load(Ordering::Acquire) + && stdlib::_thread::get_ident() != self.state.finalizing_thread_ident.load() + { // `_PyThreadState_MustExit` → `_PyThreadState_HangThread`. // Do not return SystemExit: that would mark the handle done and // make `Thread.is_alive()` false for a daemon still forced off diff --git a/extra_tests/snippets/stdlib_subinterpreters.py b/extra_tests/snippets/stdlib_subinterpreters.py index d30ae9aaa07..3ca0be27411 100644 --- a/extra_tests/snippets/stdlib_subinterpreters.py +++ b/extra_tests/snippets/stdlib_subinterpreters.py @@ -1,3 +1,8 @@ +import subprocess +import sys +import tempfile +from pathlib import Path + import _interpchannels import _interpreters @@ -21,3 +26,94 @@ def roundtrip_buffer(view): assert received.strides == view.strides assert received.format == view.format assert received.itemsize == view.itemsize + + +# A subinterpreter's finalizing thread need not be the process main thread. +# Keep a regression bounded: the old eval-breaker parked this worker forever. +subprocess.run( + [ + sys.executable, + "-c", + """ +from concurrent import interpreters +from threading import Thread + +completed = [] + +def worker(): + interp = interpreters.create() + interp.exec(''' +import gc, weakref +gc.disable() +class C: + pass +item = C() +item.cycle = item +ref = weakref.ref(item, lambda _: None) +del item +''') + interp.close() + completed.append(True) + +thread = Thread(target=worker) +thread.start() +thread.join() +assert completed == [True] +""", + ], + check=True, + timeout=30, +) + + +# Script globals persist in the isolated copy, without replacing worker __main__. +with tempfile.TemporaryDirectory() as directory: + script = Path(directory) / "pool_main.py" + script.write_text( + """ +from concurrent.futures import InterpreterPoolExecutor + +counter = 0 + +def increment(): + global counter + import __main__ + counter += 1 + return counter, hasattr(__main__, 'increment') + +def inspect_state(): + import __main__ + return counter, hasattr(__main__, 'increment') + +def consume(queue): + from threading import Barrier, Thread + barrier = Barrier(2) + results = [] + def receive(): + barrier.wait() + results.append(queue.get()()) + threads = [Thread(target=receive) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + return tuple(results) + +if __name__ == '__main__': + from concurrent import interpreters + queue = interpreters.create_queue() + queue.put(inspect_state) + queue.put(inspect_state) + interp = interpreters.create() + try: + assert interp.call(consume, queue) == ((0, False), (0, False)) + finally: + interp.close() + with InterpreterPoolExecutor(max_workers=1) as pool: + assert pool.submit(increment).result() == (1, False) + assert pool.submit(increment).result() == (2, False) + assert counter == 0 +""", + encoding="utf-8", + ) + subprocess.run([sys.executable, str(script)], check=True, timeout=30)