add
This commit is contained in:
macbook_max
@@ -0,0 +1,443 @@
|
||||
"""Async I/O backend support utilities."""
|
||||
|
||||
import logging
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from contextlib import contextmanager
|
||||
from queue import Empty
|
||||
from time import sleep
|
||||
from weakref import WeakKeyDictionary
|
||||
|
||||
from kombu.utils.compat import detect_environment
|
||||
|
||||
from celery import states
|
||||
from celery.exceptions import TimeoutError
|
||||
from celery.utils.log import get_logger
|
||||
from celery.utils.threads import THREAD_TIMEOUT_MAX
|
||||
|
||||
E_CELERY_RESTART_REQUIRED = "Celery must be restarted because a shutdown signal was detected."
|
||||
|
||||
E_RETRY_LIMIT_EXCEEDED = """
|
||||
Retry limit exceeded while trying to reconnect to the Celery result store
|
||||
backend. The Celery application must be restarted.
|
||||
"""
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
__all__ = (
|
||||
'AsyncBackendMixin', 'BaseResultConsumer', 'Drainer',
|
||||
'register_drainer',
|
||||
)
|
||||
|
||||
|
||||
class EventletAdaptedEvent:
|
||||
"""
|
||||
An adapted eventlet event, designed to match the API of `threading.Event` and
|
||||
`gevent.event.Event`.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
import eventlet
|
||||
self.evt = eventlet.Event()
|
||||
|
||||
def is_set(self):
|
||||
return self.evt.ready()
|
||||
|
||||
def set(self):
|
||||
return self.evt.send()
|
||||
|
||||
def wait(self, timeout=None):
|
||||
return self.evt.wait(timeout)
|
||||
|
||||
|
||||
drainers = {}
|
||||
|
||||
|
||||
def register_drainer(name):
|
||||
"""Decorator used to register a new result drainer type."""
|
||||
def _inner(cls):
|
||||
drainers[name] = cls
|
||||
return cls
|
||||
return _inner
|
||||
|
||||
|
||||
@register_drainer('default')
|
||||
class Drainer:
|
||||
"""Result draining service."""
|
||||
|
||||
def __init__(self, result_consumer):
|
||||
self.result_consumer = result_consumer
|
||||
|
||||
def start(self):
|
||||
pass
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
def drain_events_until(self, p, timeout=None, interval=1, on_interval=None, wait=None):
|
||||
wait = wait or self.result_consumer.drain_events
|
||||
time_start = time.monotonic()
|
||||
|
||||
while 1:
|
||||
# Total time spent may exceed a single call to wait()
|
||||
if timeout and time.monotonic() - time_start >= timeout:
|
||||
raise socket.timeout()
|
||||
try:
|
||||
yield self.wait_for(p, wait, timeout=interval)
|
||||
except socket.timeout:
|
||||
pass
|
||||
except OSError:
|
||||
# Recoverable connection error (e.g. broker restart).
|
||||
# drain_events handles reconnection internally; if an
|
||||
# OSError still leaks through, we log, sleep for one
|
||||
# interval, and continue rather than spinning hot.
|
||||
logging.warning(
|
||||
'Drainer: connection error during drain_events, '
|
||||
'will retry on next loop iteration.',
|
||||
exc_info=True,
|
||||
)
|
||||
time.sleep(interval)
|
||||
|
||||
if on_interval:
|
||||
on_interval()
|
||||
if p.ready: # got event on the wanted channel.
|
||||
break
|
||||
|
||||
def wait_for(self, p, wait, timeout=None):
|
||||
wait(timeout=timeout)
|
||||
|
||||
def _event(self):
|
||||
return threading.Event()
|
||||
|
||||
|
||||
class greenletDrainer(Drainer):
|
||||
spawn = None
|
||||
_exc = None
|
||||
_g = None
|
||||
_drain_complete_event = None # event, sended (and recreated) after every drain_events iteration
|
||||
|
||||
def _send_drain_complete_event(self):
|
||||
self._drain_complete_event.set()
|
||||
self._drain_complete_event = self._event()
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
self._started = self._event()
|
||||
self._stopped = self._event()
|
||||
self._shutdown = self._event()
|
||||
self._drain_complete_event = self._event()
|
||||
|
||||
def run(self):
|
||||
self._started.set()
|
||||
|
||||
try:
|
||||
while not self._stopped.is_set():
|
||||
try:
|
||||
self.result_consumer.drain_events(timeout=1)
|
||||
self._send_drain_complete_event()
|
||||
except socket.timeout:
|
||||
pass
|
||||
except OSError:
|
||||
# Recoverable connection errors (e.g. broker restart)
|
||||
# are handled inside drain_events via reconnection.
|
||||
# If something still leaks through, we log, back off
|
||||
# briefly, and retry instead of spinning hot.
|
||||
logging.warning(
|
||||
'Drainer: connection error during drain_events, '
|
||||
'will retry on next loop iteration.',
|
||||
exc_info=True,
|
||||
)
|
||||
time.sleep(1)
|
||||
except Exception as e:
|
||||
self._exc = e
|
||||
raise
|
||||
finally:
|
||||
self._send_drain_complete_event()
|
||||
try:
|
||||
self._shutdown.set()
|
||||
except RuntimeError as e:
|
||||
logging.error(f"Failed to set shutdown event: {e}")
|
||||
|
||||
def start(self):
|
||||
self._ensure_not_shut_down()
|
||||
|
||||
if not self._started.is_set():
|
||||
self._g = self.spawn(self.run)
|
||||
self._started.wait()
|
||||
|
||||
def stop(self):
|
||||
self._stopped.set()
|
||||
self._shutdown.wait(THREAD_TIMEOUT_MAX)
|
||||
|
||||
def wait_for(self, p, wait, timeout=None):
|
||||
self.start()
|
||||
if not p.ready:
|
||||
self._drain_complete_event.wait(timeout=timeout)
|
||||
|
||||
self._ensure_not_shut_down()
|
||||
|
||||
def _ensure_not_shut_down(self):
|
||||
"""Currently used to ensure the drainer has not run to completion.
|
||||
|
||||
Raises if the shutdown event has been signaled (either due to an exception
|
||||
or stop() being called).
|
||||
|
||||
The _shutdown event acts as synchronization to ensure _exc is properly
|
||||
set before it is read from, avoiding need for locks.
|
||||
"""
|
||||
if self._shutdown.is_set():
|
||||
if self._exc is not None:
|
||||
raise self._exc
|
||||
else:
|
||||
raise Exception(E_CELERY_RESTART_REQUIRED)
|
||||
|
||||
|
||||
@register_drainer('eventlet')
|
||||
class eventletDrainer(greenletDrainer):
|
||||
|
||||
def spawn(self, func):
|
||||
from eventlet import sleep, spawn
|
||||
g = spawn(func)
|
||||
sleep(0)
|
||||
return g
|
||||
|
||||
def _event(self):
|
||||
return EventletAdaptedEvent()
|
||||
|
||||
|
||||
@register_drainer('gevent')
|
||||
class geventDrainer(greenletDrainer):
|
||||
|
||||
def spawn(self, func):
|
||||
import gevent
|
||||
g = gevent.spawn(func)
|
||||
gevent.sleep(0)
|
||||
return g
|
||||
|
||||
def _event(self):
|
||||
from gevent.event import Event
|
||||
return Event()
|
||||
|
||||
|
||||
class AsyncBackendMixin:
|
||||
"""Mixin for backends that enables the async API."""
|
||||
|
||||
def _collect_into(self, result, bucket):
|
||||
self.result_consumer.buckets[result] = bucket
|
||||
|
||||
def iter_native(self, result, no_ack=True, **kwargs):
|
||||
self._ensure_not_eager()
|
||||
|
||||
results = result.results
|
||||
if not results:
|
||||
raise StopIteration()
|
||||
|
||||
# we tell the result consumer to put consumed results
|
||||
# into these buckets.
|
||||
bucket = deque()
|
||||
for node in results:
|
||||
if not hasattr(node, '_cache'):
|
||||
bucket.append(node)
|
||||
elif node._cache:
|
||||
bucket.append(node)
|
||||
else:
|
||||
self._collect_into(node, bucket)
|
||||
|
||||
for _ in self._wait_for_pending(result, no_ack=no_ack, **kwargs):
|
||||
while bucket:
|
||||
node = bucket.popleft()
|
||||
if not hasattr(node, '_cache'):
|
||||
yield node.id, node.children
|
||||
else:
|
||||
yield node.id, node._cache
|
||||
while bucket:
|
||||
node = bucket.popleft()
|
||||
yield node.id, node._cache
|
||||
|
||||
def add_pending_result(self, result, weak=False, start_drainer=True):
|
||||
if start_drainer:
|
||||
self.result_consumer.drainer.start()
|
||||
try:
|
||||
self._maybe_resolve_from_buffer(result)
|
||||
except Empty:
|
||||
self._add_pending_result(result.id, result, weak=weak)
|
||||
return result
|
||||
|
||||
def _maybe_resolve_from_buffer(self, result):
|
||||
result._maybe_set_cache(self._pending_messages.take(result.id))
|
||||
|
||||
def _add_pending_result(self, task_id, result, weak=False):
|
||||
concrete, weak_ = self._pending_results
|
||||
if task_id not in weak_ and result.id not in concrete:
|
||||
(weak_ if weak else concrete)[task_id] = result
|
||||
self.result_consumer.consume_from(task_id)
|
||||
|
||||
def add_pending_results(self, results, weak=False):
|
||||
self.result_consumer.drainer.start()
|
||||
return [self.add_pending_result(result, weak=weak, start_drainer=False)
|
||||
for result in results]
|
||||
|
||||
def remove_pending_result(self, result):
|
||||
self._remove_pending_result(result.id)
|
||||
self.on_result_fulfilled(result)
|
||||
return result
|
||||
|
||||
def _remove_pending_result(self, task_id):
|
||||
for mapping in self._pending_results:
|
||||
mapping.pop(task_id, None)
|
||||
|
||||
def on_result_fulfilled(self, result):
|
||||
self.result_consumer.cancel_for(result.id)
|
||||
|
||||
def wait_for_pending(self, result,
|
||||
callback=None, propagate=True, **kwargs):
|
||||
self._ensure_not_eager()
|
||||
for _ in self._wait_for_pending(result, **kwargs):
|
||||
pass
|
||||
return result.maybe_throw(callback=callback, propagate=propagate)
|
||||
|
||||
def _wait_for_pending(self, result,
|
||||
timeout=None, on_interval=None, on_message=None,
|
||||
**kwargs):
|
||||
return self.result_consumer._wait_for_pending(
|
||||
result, timeout=timeout,
|
||||
on_interval=on_interval, on_message=on_message,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@property
|
||||
def is_async(self):
|
||||
return True
|
||||
|
||||
|
||||
class BaseResultConsumer:
|
||||
"""Manager responsible for consuming result messages."""
|
||||
|
||||
#: Tuple of transport-layer exceptions that signal a lost connection.
|
||||
#: Subclasses should override this with the appropriate exception types
|
||||
#: so that :meth:`reconnect_on_error` can catch and recover from them.
|
||||
_connection_errors = ()
|
||||
|
||||
def __init__(self, backend, app, accept,
|
||||
pending_results, pending_messages):
|
||||
self.backend = backend
|
||||
self.app = app
|
||||
self.accept = accept
|
||||
self._pending_results = pending_results
|
||||
self._pending_messages = pending_messages
|
||||
self.on_message = None
|
||||
self.buckets = WeakKeyDictionary()
|
||||
self.drainer = drainers[detect_environment()](self)
|
||||
|
||||
def start(self, initial_task_id, **kwargs):
|
||||
raise NotImplementedError()
|
||||
|
||||
@contextmanager
|
||||
def reconnect_on_error(self):
|
||||
"""Context manager that catches connection errors and reconnects.
|
||||
|
||||
Wraps a block of code so that any :attr:`_connection_errors` raised
|
||||
inside it trigger a call to :meth:`_reconnect`. If reconnection
|
||||
itself raises a connection error the consumer is considered
|
||||
unrecoverable and a :exc:`RuntimeError` is raised to signal that
|
||||
the Celery application must be restarted.
|
||||
"""
|
||||
try:
|
||||
yield
|
||||
except self._connection_errors:
|
||||
try:
|
||||
self._reconnect()
|
||||
except self._connection_errors as exc:
|
||||
logger.critical(E_RETRY_LIMIT_EXCEEDED)
|
||||
raise RuntimeError(E_RETRY_LIMIT_EXCEEDED) from exc
|
||||
|
||||
def _reconnect(self):
|
||||
"""Re-establish the backend connection.
|
||||
|
||||
Subclasses must override this method to perform the transport-specific
|
||||
reconnection logic that should be executed when a connection error is
|
||||
caught by :meth:`reconnect_on_error`.
|
||||
"""
|
||||
pass
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
def drain_events(self, timeout=None):
|
||||
raise NotImplementedError()
|
||||
|
||||
def consume_from(self, task_id):
|
||||
raise NotImplementedError()
|
||||
|
||||
def cancel_for(self, task_id):
|
||||
raise NotImplementedError()
|
||||
|
||||
def _after_fork(self):
|
||||
self.buckets.clear()
|
||||
self.buckets = WeakKeyDictionary()
|
||||
self.on_message = None
|
||||
self.on_after_fork()
|
||||
|
||||
def on_after_fork(self):
|
||||
pass
|
||||
|
||||
def drain_events_until(self, p, timeout=None, on_interval=None):
|
||||
return self.drainer.drain_events_until(
|
||||
p, timeout=timeout, on_interval=on_interval)
|
||||
|
||||
def _wait_for_pending(self, result,
|
||||
timeout=None, on_interval=None, on_message=None,
|
||||
**kwargs):
|
||||
self.on_wait_for_pending(result, timeout=timeout, **kwargs)
|
||||
prev_on_m, self.on_message = self.on_message, on_message
|
||||
try:
|
||||
for _ in self.drain_events_until(
|
||||
result.on_ready, timeout=timeout,
|
||||
on_interval=on_interval):
|
||||
yield
|
||||
sleep(0)
|
||||
except socket.timeout:
|
||||
raise TimeoutError('The operation timed out.')
|
||||
finally:
|
||||
self.on_message = prev_on_m
|
||||
|
||||
def on_wait_for_pending(self, result, timeout=None, **kwargs):
|
||||
pass
|
||||
|
||||
def on_out_of_band_result(self, message):
|
||||
self.on_state_change(message.payload, message)
|
||||
|
||||
def _get_pending_result(self, task_id):
|
||||
for mapping in self._pending_results:
|
||||
try:
|
||||
return mapping[task_id]
|
||||
except KeyError:
|
||||
pass
|
||||
raise KeyError(task_id)
|
||||
|
||||
def on_state_change(self, meta, message):
|
||||
if self.on_message:
|
||||
self.on_message(meta)
|
||||
if meta['status'] in states.READY_STATES:
|
||||
task_id = meta['task_id']
|
||||
try:
|
||||
result = self._get_pending_result(task_id)
|
||||
except KeyError:
|
||||
# send to buffer in case we received this result
|
||||
# before it was added to _pending_results.
|
||||
self._pending_messages.put(task_id, meta)
|
||||
else:
|
||||
result._maybe_set_cache(meta)
|
||||
buckets = self.buckets
|
||||
try:
|
||||
# remove bucket for this result, since it's fulfilled
|
||||
bucket = buckets.pop(result)
|
||||
except KeyError:
|
||||
pass
|
||||
else:
|
||||
# send to waiter via bucket
|
||||
bucket.append(result)
|
||||
sleep(0)
|
||||
Reference in New Issue
Block a user