| from concurrent.futures import Executor |
| from threading import Thread, Lock |
| from collections import namedtuple, deque |
| from functools import partial |
| import logging |
| import weakref |
| |
| from .common import MAX_TIMEOUT |
| from .wrap import CanCustomizeBind |
| from .map import MapFuture |
| from .helpers import executor_loop, ShutdownHelper |
| from .event import get_event, is_shutdown |
| from .logwrap import LogWrapper |
| |
| from .metrics import metrics, track_future |
| |
| |
| class ThrottleFuture(MapFuture): |
| def __init__(self, executor): |
| self._executor = executor |
| super(ThrottleFuture, self).__init__(delegate=None, map_fn=lambda x: x) |
| self.add_done_callback(self._clear_executor) |
| |
| def _me_cancel(self): |
| if self._delegate: |
| return self._delegate.cancel() |
| executor = self._executor |
| return executor and executor._do_cancel(self) |
| |
| @classmethod |
| def _clear_executor(cls, future): |
| future._executor = None |
| |
| |
| ThrottleJob = namedtuple("ThrottleJob", ["future", "fn", "args", "kwargs"]) |
| |
| |
| class AtomicInt(object): |
| def __init__(self): |
| self.value = 0 |
| self.lock = Lock() |
| |
| def decr(self): |
| with self.lock: |
| self.value -= 1 |
| |
| def incr(self): |
| with self.lock: |
| self.value += 1 |
| |
| |
| class ThrottleExecutor(CanCustomizeBind, Executor): |
| """An executor which delegates to another executor while enforcing |
| a limit on the number of futures running concurrently. |
| |
| - Callables are submitted to the delegate executor, from a different |
| thread than the calling thread. |
| |
| - Where `count` is used to initialize this executor, if there |
| are already `count` futures submitted to the delegate executor and not |
| yet :meth:`~concurrent.futures.Future.done`, additional callables will |
| either be queued or will block on submit, and will only be submitted |
| to the delegate executor once there are less than `count` futures in |
| progress. |
| |
| .. versionadded:: 1.9.0 |
| """ |
| |
| def __init__(self, delegate, count, logger=None, name="default", block=False): |
| """ |
| Parameters: |
| delegate (~concurrent.futures.Executor): |
| an executor to which callables will be submitted |
| |
| count (int, callable): |
| int: |
| maximum number of concurrently running futures. |
| callable: |
| a callable which returns an ``int`` (or ``None``, to indicate |
| no throttling). |
| |
| The callable will be invoked each time this executor needs |
| to decide whether to throttle futures; this may be used |
| to implement dynamic throttling. |
| |
| .. versionadded:: 2.5.0 |
| |
| block (bool) |
| If ``True``, calls to ``submit()`` on this executor may block if |
| there are already ``count`` futures in progress. |
| |
| Otherwise, calls to ``submit()`` will always return immediately |
| and callables will be queued internally. |
| |
| logger (~logging.Logger): |
| a logger used for messages from this executor |
| |
| name (str): |
| a name for this executor |
| |
| .. versionchanged:: 2.7.0 |
| Introduced ``name``. |
| |
| .. versionchanged:: 2.11.0 |
| Introduced ``block``. |
| """ |
| self._log = LogWrapper( |
| logger if logger else logging.getLogger("ThrottleExecutor") |
| ) |
| self._block = block |
| self._name = name |
| self._delegate = delegate |
| self._to_submit = deque() |
| self._lock = Lock() |
| self._event = get_event() |
| self._running_count = AtomicInt() |
| self._throttle = count if callable(count) else lambda: count |
| self._last_throttle = self._throttle() |
| self._shutdown = ShutdownHelper() |
| |
| event = self._event |
| self_ref = weakref.ref(self, lambda _: event.set()) |
| |
| metrics.EXEC_INPROGRESS.labels(type="throttle", executor=self._name).inc() |
| metrics.EXEC_TOTAL.labels(type="throttle", executor=self._name).inc() |
| |
| self._thread = Thread( |
| name="ThrottleExecutor-%s" % name, target=_submit_loop, args=(self_ref,) |
| ) |
| self._thread.daemon = True |
| self._thread.start() |
| |
| def submit(self, fn, *args, **kwargs): # pylint: disable=arguments-differ |
| with self._shutdown.ensure_alive(): |
| self._block_until_ready(self._eval_throttle()) |
| |
| out = ThrottleFuture(self) |
| track_future(out, type="throttle", executor=self._name) |
| |
| job = ThrottleJob(out, fn, args, kwargs) |
| with self._lock: |
| self._to_submit.append(job) |
| metrics.THROTTLE_QUEUE.labels(executor=self._name).inc() |
| self._log.debug("Enqueued: %s", job) |
| self._event.set() |
| return out |
| |
| def shutdown(self, wait=True, **_kwargs): |
| if self._shutdown(): |
| self._log.debug("Shutting down") |
| metrics.EXEC_INPROGRESS.labels(type="throttle", executor=self._name).dec() |
| self._delegate.shutdown(wait, **_kwargs) |
| self._event.set() |
| if wait: |
| self._thread.join(MAX_TIMEOUT) |
| |
| def _block_until_ready(self, throttle_val): |
| while self._block and not self._shutdown.is_shutdown: |
| if len(self._to_submit) < throttle_val: |
| return |
| self._log.debug("%s: throttling on submit", self._name) |
| self._event.wait(30.0) |
| |
| def _eval_throttle(self): |
| try: |
| self._last_throttle = self._throttle() |
| except Exception: |
| self._log.exception( |
| "Error evaluating throttle count via %r", self._throttle |
| ) |
| |
| return self._last_throttle |
| |
| def _do_submit(self, job): |
| delegate_future = self._delegate.submit(job.fn, *job.args, **job.kwargs) |
| self._log.debug("Submitted %s yielding %s", job, delegate_future) |
| |
| delegate_future.add_done_callback( |
| partial( |
| self._delegate_future_done, self._log, self._running_count, self._event |
| ) |
| ) |
| job.future._set_delegate(delegate_future) |
| |
| def _do_cancel(self, future): |
| with self._lock: |
| for job in self._to_submit: |
| if job.future is future: |
| self._to_submit.remove(job) |
| self._log.debug("Cancelled %s", job) |
| return True |
| self._log.debug("Could not find for cancel: %s", future) |
| return False |
| |
| @classmethod |
| def _delegate_future_done(cls, log, running_count, event, future): |
| # Whenever an inner future completes, the thread should wake up |
| # in case there's something to be submitted |
| log.debug("Delegate future done: %s", future) |
| running_count.decr() |
| event.set() |
| |
| |
| def _submit_loop_iter(executor): |
| if not executor: |
| return |
| |
| if executor._shutdown.is_shutdown or is_shutdown(): |
| return |
| |
| throttle = executor._eval_throttle() |
| to_submit = [] |
| with executor._lock: |
| while executor._to_submit: |
| if throttle is not None and (executor._running_count.value >= throttle): |
| executor._log.debug("Throttled") |
| break |
| job = executor._to_submit.popleft() |
| executor._log.debug("Will submit: %s", job) |
| to_submit.append(job) |
| |
| # While not actually running yet, we've committed to running it, so... |
| executor._running_count.incr() |
| metrics.THROTTLE_QUEUE.labels(executor=executor._name).dec() |
| |
| executor._log.debug( |
| "Submitting %s, throttling %s", len(to_submit), len(executor._to_submit) |
| ) |
| |
| for job in to_submit: |
| executor._do_submit(job) |
| |
| # Because the throttle count is dynamic, we should wake up at some point in |
| # the future to re-check throttle, even if no other events occurred. |
| # If there's nothing running at all, we should do this sooner, because |
| # there won't be any events from completing futures. |
| return executor._event, 30.0 if executor._running_count.value else 2.0 |
| |
| |
| @executor_loop |
| def _submit_loop(executor_ref): |
| while True: |
| result = _submit_loop_iter(executor_ref()) |
| if not result: |
| break |
| |
| (event, wait_time) = result |
| event.wait(wait_time) |
| event.clear() |
