bounded_pool_executor_inferface.py

import threading
from abc import ABC
from concurrent.futures import Executor, Future
from multiprocessing import synchronize
from typing import Optional, Sequence, Union

__all__: Sequence[str] = ("BoundedPoolExecutor",)


class BoundedPoolExecutor(ABC, Executor):
    """A workpool that limits the number of active work items using an internal semaphore.

    Note: implementations must also implement concurrent.futures.Executor !
    """

    semaphore: Optional[Union[threading.Lock, synchronize.SemLock]] = None

    def _acquire(self) -> None:
        """Acquires a lock on the internal semaphore."""
        assert self.semaphore is not None
        self.semaphore.acquire()

    def _release(self, _=None) -> None:
        """Releases the internal semaphore. Ignores input: necessary to accept one input for future callback."""
        assert self.semaphore is not None
        self.semaphore.release()

    def submit(self, fn, *args, **kwargs) -> Future:
        """Submits the function & its input for Asynchronous execution.

        Will wait to run until there is a free worker: this `submit` will only work on
        """
        self._acquire()
        try:
            # from concurrent.futures.Executor
            future: Future = super().submit(fn, *args, **kwargs)  # type: ignore
        except Exception:  # noqa
            self._release()
            raise
        else:
            future.add_done_callback(self._release)
            return future