deadpool-executor 2026.6.1__py3-none-any.whl

deadpool.py

@@ -0,0 +1,1022 @@
+"""
+Deadpool
+========
+
+Important design considerations:
+
+Backpressure
+------------
+
+To allow backpressure when submitting work to the pool, we make
+the ``submit`` method block when the number of pending tasks is
+greater than the ``max_workers`` parameter. This has consequences,
+basically it means the main thread is blocked and nothing else
+can happen until it unblocks by getting space in the queue.
+
+Deadpool itself needs to do actions around job management, so
+this is why we have a separate "supervisor" thread for each
+worker process.
+
+"""
+
+import concurrent.futures
+import ctypes
+import logging
+import multiprocessing as mp
+import os
+import pickle
+import signal
+import sys
+import threading
+import traceback
+import typing
+import weakref
+import atexit
+import json
+from concurrent.futures import CancelledError, Executor, InvalidStateError, as_completed
+from dataclasses import dataclass, field
+from multiprocessing.connection import Connection
+from queue import Empty, PriorityQueue, Queue, SimpleQueue
+from typing import Callable, Optional, Tuple
+from collections.abc import Mapping
+from functools import partial
+
+import psutil
+from setproctitle import setproctitle
+
+__version__ = "2026.6.1"
+__all__ = [
+    "Deadpool",
+    "Future",
+    "CancelledError",
+    "TimeoutError",
+    "ProcessError",
+    "PoolClosed",
+    "as_completed",
+]
+logger = logging.getLogger("deadpool")
+
+
+# Does not work. Hangs the process on exit.
+# There currently isn't an official way to clean up the
+# resource tracker process. It is an open issue on the
+# Python issue tracker.
+# @atexit.register
+# def stop_resource_tracker():
+#     from multiprocessing import resource_tracker
+#     tracker = resource_tracker._resource_tracker
+#     try:
+#         import time
+#         time.sleep(5)
+#         tracker._stop()
+#     except Exception:
+#         logger.info("Error stopping the multiprocessing resource tracker")
+
+
+@dataclass
+class Stat:
+    lock: threading.Lock
+    value: int = 0
+
+    def increment(self, value: int = 1):
+        with self.lock:
+            self.value += value
+
+    def set(self, value: int = 0):
+        self.value = value
+
+
+class Statistics:
+    def __init__(self):
+        self._lock = threading.Lock()
+
+        self.tasks_received = Stat(self._lock, 0)
+        self.tasks_launched = Stat(self._lock, 0)
+        self.tasks_failed = Stat(self._lock, 0)
+        self.worker_processes_created = Stat(self._lock, 0)
+        self.max_workers_busy_concurrently = Stat(self._lock, 0)
+
+    def reset_counters(self):
+        self.tasks_received.set()
+        self.tasks_launched.set()
+        self.tasks_failed.set()
+        self.worker_processes_created.set()
+        self.max_workers_busy_concurrently.set()
+
+    def to_dict(self) -> dict[str, typing.Any]:
+        return {
+            "tasks_received": self.tasks_received.value,
+            "tasks_launched": self.tasks_launched.value,
+            "tasks_failed": self.tasks_failed.value,
+            "worker_processes_created": self.worker_processes_created.value,
+            "max_workers_busy_concurrently": self.max_workers_busy_concurrently.value,
+        }
+
+
+@dataclass(order=True)
+class PrioritizedItem:
+    priority: int
+    item: typing.Any = field(compare=False)
+
+
+@dataclass(init=False)
+class WorkerProcess:
+    process: mp.Process
+    connection_receive_msgs_from_process: Connection
+    connection_send_msgs_to_process: Connection
+    # Stats
+    tasks_ran_counter: int
+    # Controls
+    # If the subprocess RSS memory is above this threshold,
+    # ask the system allocator to release unused memory back
+    # to the OS.
+    malloc_trim_rss_memory_threshold_bytes: Optional[int] = None
+    ok: bool = True
+
+    def __init__(
+        self,
+        initializer=None,
+        initargs=(),
+        finalizer=None,
+        finargs=(),
+        daemon=True,
+        mp_context="forkserver",
+        malloc_trim_rss_memory_threshold_bytes=None,
+    ):
+        # For the process to send info OUT OF the process
+        conn_receiver, conn_sender = mp.Pipe(duplex=False)
+        # For sending work INTO the process
+        conn_receiver2, conn_sender2 = mp.Pipe(duplex=False)
+        p = mp_context.Process(
+            daemon=daemon,
+            target=raw_runner2,
+            args=(
+                conn_sender,
+                conn_receiver2,
+                os.getpid(),
+                initializer,
+                initargs,
+                finalizer,
+                finargs,
+                malloc_trim_rss_memory_threshold_bytes,
+            ),
+        )
+
+        p.start()
+        self.process = p
+        self.connection_receive_msgs_from_process = conn_receiver
+        self.connection_send_msgs_to_process = conn_sender2
+        self.tasks_ran_counter = 0
+        self.ok = True
+
+    def __hash__(self):
+        return hash(self.process.pid)
+
+    @property
+    def pid(self):
+        return self.process.pid
+
+    def get_rss_bytes(self) -> int:
+        return psutil.Process(pid=self.pid).memory_info().rss
+
+    def submit_job(self, job):
+        self.tasks_ran_counter += 1
+        self.connection_send_msgs_to_process.send(job)
+
+    def shutdown(self, wait=True):
+        if not self.process.is_alive():
+            return
+
+        self.connection_receive_msgs_from_process.close()
+
+        if self.connection_send_msgs_to_process.writable:  # pragma: no branch
+            try:
+                self.connection_send_msgs_to_process.send(None)
+            except BrokenPipeError:  # pragma: no cover
+                pass
+            else:
+                self.connection_send_msgs_to_process.close()
+
+        if wait:
+            self.process.join()
+
+    def is_alive(self):
+        return self.process.is_alive()
+
+    def format_death_message(self, join_timeout: float = 0.1) -> str:
+        # When a worker dies from a signal, the parent sees EOF on the pipe
+        # before `exitcode` becomes non-None — so the obvious `exitcode` check
+        # right after EOFError reports None, and we lose the signal name.
+        # See issue #331.
+        #
+        # Ordering on Linux: `do_exit()` runs `__exit_files()` (closes the
+        # child's fds -> EOF on our pipe) *before* `exit_notify()` sets
+        # TASK_ZOMBIE and sends SIGCHLD. Only after the latter can
+        # `waitpid(pid, WNOHANG)` report a status, which is what populates
+        # `Process.exitcode` via `Popen.poll()`. The gap is typically
+        # microseconds but is unbounded under pathological load.
+        #
+        # `Process.join(timeout=X)` bridges the gap cleanly: it waits on the
+        # sentinel fd (closed in the same `__exit_files()` call, so already
+        # readable by the time we get here) and then calls blocking
+        # `waitpid(pid)`, returning as soon as the child reaches TASK_ZOMBIE.
+        # The timeout is a ceiling for pathological cases — and for the edge
+        # case where a worker closed its data pipe voluntarily without
+        # exiting, in which case we'd otherwise block until the timeout.
+        #
+        # Refs:
+        #   - cpython Lib/multiprocessing/popen_fork.py (Popen.wait/poll)
+        #   - cpython Lib/multiprocessing/process.py (Process.exitcode/join)
+        #   - Linux do_exit() ordering: __exit_files() precedes exit_notify()
+        #   - https://docs.python.org/3/library/multiprocessing.html
+        #     (Process.exitcode, Process.join semantics)
+        proc = self.process
+        if proc.exitcode is None:
+            proc.join(timeout=join_timeout)
+
+        exitcode = proc.exitcode
+        if exitcode is None:
+            return "Worker process died unexpectedly"
+
+        try:
+            signame = signal.strsignal(-exitcode)
+        except (ValueError, TypeError):
+            signame = "Unknown"
+
+        return (
+            f"Subprocess {self.pid} completed unexpectedly with "
+            f"exitcode {exitcode} ({signame})"
+        )
+
+    def results_are_available(self, block_for: float = 0.2):
+        return self.connection_receive_msgs_from_process.poll(timeout=block_for)
+
+    def get_results(self):
+        return self.connection_receive_msgs_from_process.recv()
+
+
+class Future(concurrent.futures.Future):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._pid: Optional[int] = None
+        self.pid_callback = None
+
+    @property
+    def pid(self):
+        return self._pid
+
+    @pid.setter
+    def pid(self, value):
+        self._pid = value
+        if self.pid_callback:
+            try:
+                self.pid_callback(self)
+            except Exception:  # pragma: no cover
+                logger.exception("Error calling pid_callback")
+
+    def add_pid_callback(self, fn):
+        self.pid_callback = fn
+
+    def cancel_and_kill_if_running(self, sig=signal.SIGKILL):
+        self.cancel()
+        if self.pid:
+            try:
+                kill_proc_tree(self.pid, sig=sig)
+            except Exception as e:  # pragma: no cover
+                logger.warning(f"Got error killing pid {self.pid}: {e}")
+
+
+class TimeoutError(concurrent.futures.TimeoutError): ...
+
+
+class ProcessError(mp.ProcessError): ...
+
+
+class PoolClosed(Exception): ...
+
+
+class Deadpool(Executor):
+    def __init__(
+        self,
+        max_workers: Optional[int] = None,
+        min_workers: Optional[int] = None,
+        max_tasks_per_child: Optional[int] = None,
+        max_worker_memory_bytes: Optional[int] = None,
+        mp_context=None,
+        initializer=None,
+        initargs=(),
+        finalizer=None,
+        finalargs=(),
+        max_backlog=1000,
+        shutdown_wait: Optional[bool] = None,
+        shutdown_cancel_futures: Optional[bool] = None,
+        daemon=True,
+        malloc_trim_rss_memory_threshold_bytes: Optional[int] = None,
+        propagate_environ: Optional[Mapping] = None,
+    ) -> None:
+        """The pool.
+
+        :param propagate_environ: A mapping of environment variables to
+            propagate to the worker processes. This is useful for
+            setting up the environment in the worker processes. Subprocesses
+            will inherit the environment of the parent process, but crucially,
+            they will not inherit any changes made to the environment after
+            the subprocess is created (via `os.environ`). This parameter
+            allows you to specify a mapping of environment variables to
+            propagate to the worker processes. The worker processes will
+            receive these environment variables at the time they are created.
+            There are two important points: firstly, these env vars will
+            be set before the initializer is run, so the initializer can
+            use them. Secondly, these are applied only when the worker
+            process is created, which means that you can dynamically change the
+            values of the dict supplied here, and they will be used in
+            new worker processes as they are created. (The new parameters
+            will not be seen by existing worker processes.)
+
+        """
+        super().__init__()
+
+        if not mp_context:
+            mp_context = "forkserver"
+
+        if isinstance(mp_context, str):
+            mp_context = mp.get_context(mp_context)
+
+        # This is stored (instead of immediately currying the `initializer`)
+        # for a very important reason, which you can read about in the
+        # `add_worker_to_pool` method.
+        self.propagate_environ = propagate_environ
+        self.ctx = mp_context
+        self.initializer = initializer
+        self.initargs = initargs
+        self.finitializer = finalizer
+        self.finitargs = finalargs
+        self.pool_size = max_workers or len(os.sched_getaffinity(0))
+        if min_workers is None:
+            self.min_workers = self.pool_size
+        else:
+            self.min_workers = min_workers
+
+        self.max_tasks_per_child = max_tasks_per_child
+        self.max_worker_memory_bytes = max_worker_memory_bytes
+        self.submitted_jobs: PriorityQueue[PrioritizedItem] = PriorityQueue(
+            maxsize=max_backlog
+        )
+        self.running_jobs = Queue(maxsize=self.pool_size)
+        self.running_futs = weakref.WeakSet()
+        self.existing_workers = weakref.WeakSet()
+        # Lock protecting busy_workers, existing_workers, and
+        # running_futs for thread-safety without the GIL.
+        self._workers_lock = threading.Lock()
+        self.closed = False
+        self.shutdown_wait = shutdown_wait
+        self.shutdown_cancel_futures = shutdown_cancel_futures
+        self.daemon = daemon
+        self.malloc_trim_rss_memory_threshold_bytes = (
+            malloc_trim_rss_memory_threshold_bytes
+        )
+        self._statistics = Statistics()
+
+        # TODO: overcommit
+        self.workers: SimpleQueue[WorkerProcess] = SimpleQueue()
+        for _ in range(self.pool_size):
+            self.add_worker_to_pool()
+        # When a worker is running a job, it will be removed from
+        # the workers queue, and added to the busy_workers set.
+        # When a worker successfully completes a job, it will be
+        # added back to the workers queue, and removed from the
+        # busy_workers set.
+        self.busy_workers = set()  #  weakref.WeakSet()
+
+        # THE ONLY ACTIVE, PERSISTENT STATE IN DEADPOOL IS THIS THREAD
+        # BELOW. PROTECT IT AT ALL COSTS.
+        self.runner_thread = threading.Thread(
+            target=self.runner, name="deadpool.runner", daemon=True
+        )
+        self.runner_thread.start()
+
+    def get_statistics(self) -> dict[str, typing.Any]:
+        stats = self._statistics.to_dict()
+
+        # These are not counters; they are determined at the time of the
+        # call based on the state of the worker processes.
+        with self._workers_lock:
+            stats["worker_processes_still_alive"] = len(self.existing_workers)
+            stats["worker_processes_busy"] = len(self.busy_workers)
+        stats["worker_processes_idle"] = self.workers.qsize()
+
+        return stats
+
+    def add_worker_to_pool(self):
+        if self.propagate_environ:
+            # By constructing here, late, we allow the user to make
+            # changes dynamically to the configured env vars and these
+            # will be reflected in the worker processes as they are
+            # added to the pool. This has a large number of interesting
+            # applications, such as dynamically changing the logging
+            # level of the worker processes, or changing the location
+            # of a file that the worker processes need to read, or
+            # changing timeouts and so on. All the user needs to do
+            # is update the value on the Deadpool instance itself.
+            initializer = partial(
+                initializer_environ_propagator,
+                dict(self.propagate_environ),
+                original_initializer=self.initializer,
+            )
+        else:
+            initializer = self.initializer
+
+        worker = WorkerProcess(
+            initializer=initializer,
+            initargs=self.initargs,
+            finalizer=self.finitializer,
+            finargs=self.finitargs,
+            mp_context=self.ctx,
+            daemon=self.daemon,
+            malloc_trim_rss_memory_threshold_bytes=self.malloc_trim_rss_memory_threshold_bytes,
+        )
+        self.workers.put(worker)
+        self._statistics.worker_processes_created.increment()
+        with self._workers_lock:
+            self.existing_workers.add(worker)
+
+    def clear_workers(self):
+        """Clear all workers from the pool.
+
+        Typically they will all get added back according to the
+        rules for `max_workers` and so on. One neat reason to do
+        this is to have new settings take effect, such as a new
+        environment variable that needs to be set in the workers.
+        """
+        while not self.workers.empty():
+            worker = self.workers.get()
+            worker.shutdown(wait=False)
+
+    def runner(self):
+        while True:
+            # This will block if the queue of running jobs is full.
+            self.running_jobs.put(None)
+
+            priority_job = self.submitted_jobs.get()
+            job = priority_job.item
+            if job is None:
+                # This is for the `None` that terminates the while loop.
+                self.submitted_jobs.task_done()
+                self.running_jobs.get()
+                # TODO: this probably isn't necessary, since cleanup is happening
+                # in the shutdown method anyway.
+                cancel_all_futures_on_queue(self.submitted_jobs)
+                logger.debug("Got shutdown event, leaving runner.")
+                return
+
+            *_, fut = job
+            if fut.done():
+                # This shouldn't really be possible, but if the associated future
+                # for this job has somehow already been marked as done (e.g. if
+                # the caller decided to cancel it themselves) then just skip the
+                # whole job.
+                self.submitted_jobs.task_done()
+                self.running_jobs.get()
+                continue
+
+            t = threading.Thread(target=self.run_task, args=job, daemon=True)
+            self._statistics.tasks_launched.increment()
+            t.start()
+
+    def get_process(self) -> WorkerProcess:
+        with self._workers_lock:
+            bw = len(self.busy_workers)
+        mw = self.pool_size
+        qs = self.workers.qsize()
+
+        total_workers = bw + qs
+        if total_workers < mw and qs == 0:
+            self.add_worker_to_pool()
+
+        wp = self.workers.get()
+        with self._workers_lock:
+            self.busy_workers.add(wp)
+            busy_count = len(self.busy_workers)
+        with self._statistics.max_workers_busy_concurrently.lock:
+            if busy_count > self._statistics.max_workers_busy_concurrently.value:
+                self._statistics.max_workers_busy_concurrently.value = busy_count
+
+        return wp
+
+    def done_with_process(self, wp: WorkerProcess):
+        # This worker is done with its job and is no longer busy.
+        with self._workers_lock:
+            self.busy_workers.remove(wp)
+            count_workers_busy = len(self.busy_workers)
+        count_workers_idle = self.workers.qsize()
+        backlog_size = self.submitted_jobs.qsize()
+
+        # The `1` is for `wp` itself.
+        total_workers = count_workers_busy + count_workers_idle + 1
+        there_are_more_workers_than_min = total_workers > self.min_workers
+        task_backlog_is_empty = backlog_size == 0
+
+        # if there_are_more_workers_than_min and (there_are_idle_workers or task_backlog_is_empty):
+        if there_are_more_workers_than_min and task_backlog_is_empty:
+            # We have more workers than the minimum, and there is no backlog of
+            # tasks. This implies any tasks currently in play have already been picked
+            # up by workers in the pool, or the pool is idle. We can safely remove
+            # this worker from the pool.
+            wp.shutdown(wait=False)
+            return
+
+        if not wp.is_alive():
+            self.add_worker_to_pool()
+            return
+
+        if not wp.ok:
+            self.add_worker_to_pool()
+            return
+
+        if self.max_tasks_per_child is not None:
+            if wp.tasks_ran_counter >= self.max_tasks_per_child:
+                logger.debug(f"Worker {wp.pid} hit max tasks per child.")
+                wp.shutdown(wait=False)
+                self.add_worker_to_pool()
+                return
+
+        if self.max_worker_memory_bytes is not None:
+            mem = wp.get_rss_bytes()
+            logger.debug(f"Worker {wp.pid} has {mem} bytes of RSS memory.")
+            if mem >= self.max_worker_memory_bytes:
+                logger.debug(f"Worker {wp.pid} hit max memory threshold.")
+                wp.shutdown(wait=False)
+                self.add_worker_to_pool()
+                return
+
+        self.workers.put(wp)
+
+    def run_task(self, fn, args, kwargs, timeout, fut: Future):
+        try:
+            retry_count = 10
+            while retry_count > 0:
+                retry_count -= 1
+                worker: WorkerProcess = self.get_process()
+                try:
+                    worker.submit_job((fn, args, kwargs, timeout))
+                    break
+                except (pickle.PicklingError, AttributeError) as e:
+                    # If the user passed in a function or params that can't
+                    # be pickled, use the future to communicate the error.
+                    # Note that in this scenario, there is nothing wrong
+                    # with the worker process itself, so we don't need to
+                    # shut it down.
+                    fut.set_exception(e)
+                    self.done_with_process(worker)
+                    return
+                except BrokenPipeError:
+                    # This likely comes from trying to send a job over a pipe
+                    # that has been closed. This is a serious problem, and
+                    # we should shut down the worker process and get rid of
+                    # it. We're going to loop back around and try again with
+                    # a new worker.
+                    # TODO: it seems that this might be expected in situations
+                    #  where the worker process often OOMs. As such, not sure
+                    #  whether logging at warning level is appropriate.
+                    logger.warning(f"BrokenPipeError on {worker.pid}, retrying.")
+                    worker.ok = False
+                    self.done_with_process(worker)
+                    # TODO: probably this should be moved into the `done_with_process`
+                    #  and can act on the `worker.ok` flag.
+                    kill_proc_tree(worker.pid, sig=signal.SIGKILL)
+            else:  # pragma: no cover
+                # If we get here, we've tried to submit the job to a worker
+                # process multiple times and failed each time. We're giving
+                # up.
+                logger.error("Failed to submit job to worker")
+                fut.set_exception(ProcessError("Failed to submit job to worker"))
+                return
+
+            fut.pid = worker.pid
+            with self._workers_lock:
+                self.running_futs.add(fut)
+
+            while True:
+                if worker.results_are_available():
+                    try:
+                        results = worker.get_results()
+                    except EOFError:
+                        self._statistics.tasks_failed.increment()
+                        if not fut.done():
+                            try:
+                                fut.set_exception(
+                                    ProcessError(worker.format_death_message())
+                                )
+                            except InvalidStateError:
+                                pass
+                    except BaseException as e:
+                        self._statistics.tasks_failed.increment()
+                        logger.debug(f"Unexpected exception from worker: {e}")
+                        if not fut.done():
+                            try:
+                                fut.set_exception(e)
+                            except InvalidStateError:
+                                pass
+                    else:
+                        if isinstance(results, BaseException):
+                            self._statistics.tasks_failed.increment()
+                            if not fut.done():
+                                try:
+                                    fut.set_exception(results)
+                                except InvalidStateError:
+                                    pass
+                        else:
+                            if not fut.done():
+                                try:
+                                    fut.set_result(results)
+                                except InvalidStateError:
+                                    pass
+
+                        if isinstance(results, TimeoutError):
+                            self._statistics.tasks_failed.increment()
+                            logger.debug(
+                                f"TimeoutError on {worker.pid}, setting ok=False"
+                            )
+                            worker.ok = False
+                    break
+                elif not worker.is_alive():
+                    self._statistics.tasks_failed.increment()
+                    logger.debug(f"p is no longer alive: {worker.process}")
+                    if not fut.done():
+                        # It is possible that fut has already had a result set on
+                        # it. If that's the case we'll do nothing. Otherwise, put
+                        # an exception reporting the unexpected situation.
+                        try:
+                            fut.set_exception(
+                                ProcessError(worker.format_death_message())
+                            )
+                        except InvalidStateError:  # pragma: no cover
+                            # We still have to catch this even though there is a
+                            # check for `fut.done()`, simply due to an possible
+                            # race between the done check and the set_exception call.
+                            pass
+
+                    break
+                else:
+                    pass  # pragma: no cover
+
+            self.done_with_process(worker)
+        finally:
+            self.submitted_jobs.task_done()
+
+            if not fut.done():  # pragma: no cover
+                fut.set_exception(ProcessError("Somehow no result got set on fut."))
+
+            try:
+                self.running_jobs.get_nowait()
+            except Empty:  # pragma: no cover
+                logger.warning("Weird error, did not expect running jobs to be empty")
+
+    def submit(
+        self,
+        fn: Callable,
+        /,
+        *args,
+        deadpool_timeout=None,
+        deadpool_priority=0,
+        **kwargs,
+    ) -> Future:
+        if deadpool_priority < 0:  # pragma: no cover
+            raise ValueError(
+                f"Parameter deadpool_priority must be >= 0, but was {deadpool_priority}"
+            )
+
+        if self.closed:
+            raise PoolClosed("The pool is closed. No more tasks can be submitted.")
+
+        fut = Future()
+        self.submitted_jobs.put(
+            PrioritizedItem(
+                priority=deadpool_priority,
+                item=(fn, args, kwargs, deadpool_timeout, fut),
+            )
+        )
+        self._statistics.tasks_received.increment()
+        return fut
+
+    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
+        if self.closed:
+            return
+
+        logger.debug(f"shutdown: {wait=} {cancel_futures=}")
+
+        # No more new tasks can be submitted
+        self.closed = True
+
+        if cancel_futures:
+            cancel_all_futures_on_queue(self.submitted_jobs)
+
+        if wait:
+            # The None sentinel will pop last
+            shutdown_priority = sys.maxsize
+        else:
+            # The None sentinel will pop first
+            shutdown_priority = -1
+
+        try:
+            self.submitted_jobs.put(
+                PrioritizedItem(priority=shutdown_priority, item=None),
+                timeout=2.0,
+            )
+        except TimeoutError:  # pragma: no cover
+            logger.warning(
+                "Timed out putting None on the submit queue. This "
+                "should not be possible "
+                "and might be a bug in deadpool."
+            )
+
+        # Up till this point, all the pending work that has been
+        # submitted, but not yet started, has been cancelled. The
+        # runner loop has also been stopped (with the None sentinel).
+        # The only thing left to do is decide whether or not to
+        # actively kill processes that are still running. We presume
+        # that if the user is asking for cancellation and doesn't
+        # want to wait, that she probably wants us to also stop
+        # running processes.
+        if (not wait) and cancel_futures:
+            with self._workers_lock:
+                running_futs = list(self.running_futs)
+            for fut in running_futs:
+                fut.cancel_and_kill_if_running()
+
+        logger.debug("waiting for submitted_jobs to join...")
+        self.submitted_jobs.join()
+
+        super().shutdown(wait, cancel_futures=cancel_futures)
+
+        # We can now remove all other processes hanging around
+        # in the background.
+        while not self.workers.empty():
+            try:
+                worker = self.workers.get_nowait()
+                worker.shutdown()
+            except Empty:  # pragma: no cover
+                break
+
+        # There may be a few processes left in the
+        # `busy_workers` queue. Shut them down too.
+        with self._workers_lock:
+            remaining = list(self.busy_workers)
+            self.busy_workers.clear()
+        for worker in remaining:
+            worker.shutdown()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        kwargs = {}
+        if self.shutdown_wait is not None:
+            kwargs["wait"] = self.shutdown_wait
+
+        if self.shutdown_cancel_futures is not None:
+            kwargs["cancel_futures"] = self.shutdown_cancel_futures
+
+        self.shutdown(**kwargs)
+        self.runner_thread.join()
+        return False
+
+
+def cancel_all_futures_on_queue(q: Queue):
+    while True:
+        try:
+            priority_item = q.get_nowait()
+            q.task_done()
+            job = priority_item.item
+            *_, fut = job
+            fut.cancel()
+        except Empty:
+            break
+
+
+# Taken from
+# https://psutil.readthedocs.io/en/latest/index.html?highlight=children#kill-process-tree
+def kill_proc_tree(
+    pid,
+    sig=signal.SIGTERM,
+    include_parent=True,
+    timeout=None,
+    on_terminate=None,
+    allow_kill_self=False,
+):
+    """Kill a process tree (including grandchildren) with signal
+    "sig" and return a (gone, still_alive) tuple.
+    "on_terminate", if specified, is a callback function which is
+    called as soon as a child terminates.
+    """
+    if not allow_kill_self and pid == os.getpid():
+        raise ValueError("Won't kill myself")
+
+    try:
+        parent = psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        return
+
+    children = parent.children(recursive=True)
+    if include_parent:
+        children.append(parent)
+
+    for p in children:
+        try:
+            p.send_signal(sig)
+        except psutil.NoSuchProcess:  # pragma: no cover
+            pass
+
+    gone, alive = psutil.wait_procs(children, timeout=timeout, callback=on_terminate)
+    return (gone, alive)
+
+
+def raw_runner2(
+    conn: Connection,
+    conn_receiver: Connection,
+    parent_pid,
+    initializer,
+    initargs,
+    finitializer: Optional[Callable] = None,
+    finitargs: Optional[Tuple] = None,
+    mem_clear_threshold_bytes: Optional[int] = None,
+    kill_proc_tree=kill_proc_tree,
+):
+    setproctitle("deadpool.worker")
+    # This event is used to signal that the "parent"
+    # monitor thread should be deactivated.
+    evt = threading.Event()
+
+    def self_destruct_if_parent_disappers():
+        """Poll every 5 seconds to see whether the parent is still
+        alive.
+        """
+        while True:
+            if evt.wait(2.0):
+                return
+
+            if not psutil.pid_exists(parent_pid):
+                logger.warning(f"Parent {parent_pid} is gone, self-destructing.")
+                evt.set()
+                atexit._run_exitfuncs()
+                kill_proc_tree(
+                    pid, sig=signal.SIGKILL, allow_kill_self=True
+                )  # pragma: no cover
+                return  # pragma: no cover
+
+    tparent = threading.Thread(target=self_destruct_if_parent_disappers, daemon=True)
+    tparent.start()
+
+    def deactivate_parentless_self_destruct():
+        evt.set()
+
+    proc = psutil.Process()
+    pid = proc.pid
+
+    def conn_send_safe(obj):
+        try:
+            conn.send(obj)
+        except BrokenPipeError:  # pragma: no cover
+            logger.debug("Pipe not usable")
+        except BaseException:  # pragma: no cover
+            logger.exception("Unexpected pipe error")
+
+    def timed_out():
+        """Action to fire when the timeout given to ``threading.Timer``
+        is reached. It kills this process with SIGKILL."""
+        # First things first. Set a self-destruct timer for ourselves.
+        # If we don't finish up in time, boom.
+        deactivate_parentless_self_destruct()
+        conn_send_safe(TimeoutError(f"Process {pid} timed out, self-destructing."))
+        # kill_proc_tree_in_process_daemon(pid, signal.SIGKILL)
+        atexit._run_exitfuncs()
+        kill_proc_tree(
+            pid, sig=signal.SIGKILL, allow_kill_self=True
+        )  # pragma: no cover
+
+    if initializer:
+        initargs = initargs or ()
+        try:
+            initializer(*initargs)
+        except Exception:
+            logger.exception("Initializer failed")
+
+    while True:
+        # Wait for some work.
+        try:
+            logger.debug("Waiting for work...")
+            job = conn_receiver.recv()
+            logger.debug("Got a job")
+        except EOFError:
+            logger.debug("Received EOF, exiting.")
+            break
+        except KeyboardInterrupt:  # pragma: no cover
+            logger.debug("Received KeyboardInterrupt, exiting.")
+            break
+        except BaseException:  # pragma: no cover
+            logger.exception("Received unexpected exception, exiting.")
+            break
+
+        if job is None:
+            logger.debug("Received None, exiting.")
+            break
+
+        # Real work, unpack.
+        fn, args, kwargs, timeout = job
+
+        if timeout:
+            t = threading.Timer(timeout, timed_out)
+            t.start()
+            deactivate_timer = lambda: t.cancel()  # noqa: E731
+        else:
+            deactivate_timer = lambda: None  # noqa: E731
+
+        try:
+            results = fn(*args, **kwargs)
+        except BaseException as e:
+            # Check whether the exception can be pickled. If not we're going
+            # to wrap it. Why do this? It turns out that mp.Connection.send
+            # will try to pickle the exception, and if it can't, it will
+            # lose its mind. I've gotten segfaults in Python with this.
+            try:  # pragma: no cover
+                pickle.dumps(e)
+            except Exception as pickle_error:
+                msg = (
+                    f"An exception occurred but pickling it failed. "
+                    f"The original exception is presented here as a string with "
+                    f"traceback.\n{e}\n{traceback.format_exception(e)}\n\n"
+                    f"The reason for the pickling failure is the following:\n"
+                    f"{traceback.format_exception(pickle_error)}"
+                )
+                e = ProcessError(msg)
+
+            # Because we can't retain the traceback (can't be pickled by default,
+            # an external library like "tblib" would be needed), we're going to
+            # render the traceback to a string and add that to the exception
+            # text. This approach also works for when deadpool can be distributed
+            # across multiple machines, since the traceback is a string.
+            traceback_str = "".join(
+                traceback.format_exception(type(e), e, e.__traceback__)
+            )
+            # Modify the exception's args to include the traceback
+            # This changes the string representation of the exception
+            e.args = (f"{e}\n{traceback_str}",) + e.args[1:]
+            conn_send_safe(e)
+        else:
+            conn_send_safe(results)
+        finally:
+            deactivate_timer()
+
+            if mem_clear_threshold_bytes is not None:
+                mem = proc.memory_info().rss
+                if mem > mem_clear_threshold_bytes:
+                    trim_memory()
+
+    if finitializer:
+        finitargs = finitargs or ()
+        try:
+            finitializer(*finitargs)
+        except BaseException:
+            logger.exception("finitializer failed")
+
+    # We've reached the end of this function which means this
+    # process must exit. However, we started a couple threads
+    # in here and they don't magically exit. Additional
+    # synchronization controls are needed to tell the threads
+    # to exit, which we don't have. However, we do have a kill
+    # switch. Since this process worker will process no more
+    # work, and since we've already fun the finalizer, we may
+    # as well just nuke it. That will remove its memory space
+    # and all its threads too.
+    deactivate_parentless_self_destruct()
+    logger.debug(f"Deleting worker {pid=}")
+    atexit._run_exitfuncs()
+    kill_proc_tree(pid, sig=signal.SIGKILL, allow_kill_self=True)  # pragma: no cover
+
+
+def kill_proc_tree_in_process_daemon(pid, sig):  # pragma: no cover
+    mp.Process(target=kill_proc_tree, args=(pid, sig), daemon=True).start()
+
+
+def trim_memory() -> None:
+    """Tell malloc to give all the unused memory back to the OS."""
+    if sys.platform == "linux":
+        libc = ctypes.CDLL("libc.so.6")
+        libc.malloc_trim(0)
+
+
+def initializer_environ_propagator(
+    environ: dict,
+    original_initializer: Optional[Callable] = None,
+    initargs=(),
+):
+    """Wrap the original initializer with one that sets the
+    environment variables in the given dict."""
+
+    # Quite important that we run this first, so that the
+    # environment variables are set before the original
+    # initializer runs. This allows the original initializer
+    # to use the environment variables.
+    os.environ.update(environ or {})
+    if original_initializer:
+        original_initializer(*(initargs or ()))