400 lines
12 KiB
Python
400 lines
12 KiB
Python
"""Event loop implementation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import errno
|
|
import threading
|
|
from contextlib import contextmanager
|
|
from copy import copy
|
|
from queue import Empty
|
|
from time import sleep
|
|
from types import GeneratorType as generator
|
|
|
|
from vine import Thenable, promise
|
|
|
|
from kombu.log import get_logger
|
|
from kombu.utils.compat import fileno
|
|
from kombu.utils.eventio import ERR, READ, WRITE, poll
|
|
from kombu.utils.objects import cached_property
|
|
|
|
from .timer import Timer
|
|
|
|
__all__ = ('Hub', 'get_event_loop', 'set_event_loop')
|
|
logger = get_logger(__name__)
|
|
|
|
_current_loop: Hub | None = None
|
|
|
|
W_UNKNOWN_EVENT = """\
|
|
Received unknown event %r for fd %r, please contact support!\
|
|
"""
|
|
|
|
|
|
class Stop(BaseException):
|
|
"""Stops the event loop."""
|
|
|
|
|
|
def _raise_stop_error():
|
|
raise Stop()
|
|
|
|
|
|
@contextmanager
|
|
def _dummy_context(*args, **kwargs):
|
|
yield
|
|
|
|
|
|
def get_event_loop() -> Hub | None:
|
|
"""Get current event loop object."""
|
|
return _current_loop
|
|
|
|
|
|
def set_event_loop(loop: Hub | None) -> Hub | None:
|
|
"""Set the current event loop object."""
|
|
global _current_loop
|
|
_current_loop = loop
|
|
return loop
|
|
|
|
|
|
class Hub:
|
|
"""Event loop object.
|
|
|
|
Arguments:
|
|
---------
|
|
timer (kombu.asynchronous.Timer): Specify custom timer instance.
|
|
"""
|
|
|
|
#: Flag set if reading from an fd will not block.
|
|
READ = READ
|
|
|
|
#: Flag set if writing to an fd will not block.
|
|
WRITE = WRITE
|
|
|
|
#: Flag set on error, and the fd should be read from asap.
|
|
ERR = ERR
|
|
|
|
#: List of callbacks to be called when the loop is exiting,
|
|
#: applied with the hub instance as sole argument.
|
|
on_close = None
|
|
|
|
def __init__(self, timer=None):
|
|
self.timer = timer if timer is not None else Timer()
|
|
|
|
self.readers = {}
|
|
self.writers = {}
|
|
self.on_tick = set()
|
|
self.on_close = set()
|
|
self._ready = set()
|
|
self._ready_lock = threading.Lock()
|
|
|
|
self._running = False
|
|
self._loop = None
|
|
|
|
# The eventloop (in celery.worker.loops)
|
|
# will merge fds in this set and then instead of calling
|
|
# the callback for each ready fd it will call the
|
|
# :attr:`consolidate_callback` with the list of ready_fds
|
|
# as an argument. This API is internal and is only
|
|
# used by the multiprocessing pool to find inqueues
|
|
# that are ready to write.
|
|
self.consolidate = set()
|
|
self.consolidate_callback = None
|
|
|
|
self.propagate_errors = ()
|
|
|
|
self._create_poller()
|
|
|
|
@property
|
|
def poller(self):
|
|
if not self._poller:
|
|
self._create_poller()
|
|
return self._poller
|
|
|
|
@poller.setter
|
|
def poller(self, value):
|
|
self._poller = value
|
|
|
|
def reset(self):
|
|
self.close()
|
|
self._create_poller()
|
|
|
|
def _create_poller(self):
|
|
self._poller = poll()
|
|
self._register_fd = self._poller.register
|
|
self._unregister_fd = self._poller.unregister
|
|
|
|
def _close_poller(self):
|
|
if self._poller is not None:
|
|
self._poller.close()
|
|
self._poller = None
|
|
self._register_fd = None
|
|
self._unregister_fd = None
|
|
|
|
def stop(self):
|
|
self.call_soon(_raise_stop_error)
|
|
|
|
def __repr__(self):
|
|
return '<Hub@{:#x}: R:{} W:{}>'.format(
|
|
id(self), len(self.readers), len(self.writers),
|
|
)
|
|
|
|
def fire_timers(self, min_delay=1, max_delay=10, max_timers=10,
|
|
propagate=()):
|
|
timer = self.timer
|
|
delay = None
|
|
if timer and timer._queue:
|
|
for i in range(max_timers):
|
|
delay, entry = next(self.scheduler)
|
|
if entry is None:
|
|
break
|
|
try:
|
|
entry()
|
|
except propagate:
|
|
raise
|
|
except (MemoryError, AssertionError):
|
|
raise
|
|
except OSError as exc:
|
|
if exc.errno == errno.ENOMEM:
|
|
raise
|
|
logger.error('Error in timer: %r', exc, exc_info=1)
|
|
except Exception as exc:
|
|
logger.error('Error in timer: %r', exc, exc_info=1)
|
|
return min(delay or min_delay, max_delay)
|
|
|
|
def _remove_from_loop(self, fd):
|
|
try:
|
|
self._unregister(fd)
|
|
finally:
|
|
self._discard(fd)
|
|
|
|
def add(self, fd, callback, flags, args=(), consolidate=False):
|
|
fd = fileno(fd)
|
|
try:
|
|
self.poller.register(fd, flags)
|
|
except ValueError:
|
|
self._remove_from_loop(fd)
|
|
raise
|
|
else:
|
|
dest = self.readers if flags & READ else self.writers
|
|
if consolidate:
|
|
self.consolidate.add(fd)
|
|
dest[fd] = None
|
|
else:
|
|
dest[fd] = callback, args
|
|
|
|
def remove(self, fd):
|
|
fd = fileno(fd)
|
|
self._remove_from_loop(fd)
|
|
|
|
def run_forever(self):
|
|
self._running = True
|
|
try:
|
|
while 1:
|
|
try:
|
|
self.run_once()
|
|
except Stop:
|
|
break
|
|
finally:
|
|
self._running = False
|
|
|
|
def run_once(self):
|
|
try:
|
|
next(self.loop)
|
|
except StopIteration:
|
|
self._loop = None
|
|
|
|
def call_soon(self, callback, *args):
|
|
if not isinstance(callback, Thenable):
|
|
callback = promise(callback, args)
|
|
with self._ready_lock:
|
|
self._ready.add(callback)
|
|
return callback
|
|
|
|
def call_later(self, delay, callback, *args):
|
|
return self.timer.call_after(delay, callback, args)
|
|
|
|
def call_at(self, when, callback, *args):
|
|
return self.timer.call_at(when, callback, args)
|
|
|
|
def call_repeatedly(self, delay, callback, *args):
|
|
return self.timer.call_repeatedly(delay, callback, args)
|
|
|
|
def add_reader(self, fds, callback, *args):
|
|
return self.add(fds, callback, READ | ERR, args)
|
|
|
|
def add_writer(self, fds, callback, *args):
|
|
return self.add(fds, callback, WRITE, args)
|
|
|
|
def remove_reader(self, fd):
|
|
writable = fd in self.writers
|
|
on_write = self.writers.get(fd)
|
|
try:
|
|
self._remove_from_loop(fd)
|
|
finally:
|
|
if writable:
|
|
cb, args = on_write
|
|
self.add(fd, cb, WRITE, args)
|
|
|
|
def remove_writer(self, fd):
|
|
readable = fd in self.readers
|
|
on_read = self.readers.get(fd)
|
|
try:
|
|
self._remove_from_loop(fd)
|
|
finally:
|
|
if readable:
|
|
cb, args = on_read
|
|
self.add(fd, cb, READ | ERR, args)
|
|
|
|
def _unregister(self, fd):
|
|
try:
|
|
self.poller.unregister(fd)
|
|
except (AttributeError, KeyError, OSError):
|
|
pass
|
|
|
|
def _pop_ready(self):
|
|
with self._ready_lock:
|
|
ready = self._ready
|
|
self._ready = set()
|
|
return ready
|
|
|
|
def close(self, *args):
|
|
[self._unregister(fd) for fd in self.readers]
|
|
self.readers.clear()
|
|
[self._unregister(fd) for fd in self.writers]
|
|
self.writers.clear()
|
|
self.consolidate.clear()
|
|
self._close_poller()
|
|
for callback in self.on_close:
|
|
callback(self)
|
|
|
|
# Complete remaining todo before Hub close
|
|
# Eg: Acknowledge message
|
|
# To avoid infinite loop where one of the callables adds items
|
|
# to self._ready (via call_soon or otherwise).
|
|
# we create new list with current self._ready
|
|
todos = self._pop_ready()
|
|
for item in todos:
|
|
item()
|
|
|
|
def _discard(self, fd):
|
|
fd = fileno(fd)
|
|
self.readers.pop(fd, None)
|
|
self.writers.pop(fd, None)
|
|
self.consolidate.discard(fd)
|
|
|
|
def on_callback_error(self, callback, exc):
|
|
logger.error(
|
|
'Callback %r raised exception: %r', callback, exc, exc_info=1,
|
|
)
|
|
|
|
def create_loop(self,
|
|
generator=generator, sleep=sleep, min=min, next=next,
|
|
Empty=Empty, StopIteration=StopIteration,
|
|
KeyError=KeyError, READ=READ, WRITE=WRITE, ERR=ERR):
|
|
readers, writers = self.readers, self.writers
|
|
poll = self.poller.poll
|
|
fire_timers = self.fire_timers
|
|
hub_remove = self.remove
|
|
scheduled = self.timer._queue
|
|
consolidate = self.consolidate
|
|
consolidate_callback = self.consolidate_callback
|
|
propagate = self.propagate_errors
|
|
|
|
while 1:
|
|
todo = self._pop_ready()
|
|
|
|
for item in todo:
|
|
if item:
|
|
item()
|
|
|
|
poll_timeout = fire_timers(propagate=propagate) if scheduled else 1
|
|
|
|
for tick_callback in copy(self.on_tick):
|
|
tick_callback()
|
|
|
|
# print('[[[HUB]]]: %s' % (self.repr_active(),))
|
|
if readers or writers:
|
|
to_consolidate = []
|
|
try:
|
|
events = poll(poll_timeout)
|
|
# print('[EVENTS]: %s' % (self.repr_events(events),))
|
|
except ValueError: # Issue celery/#882
|
|
return
|
|
|
|
for fd, event in events or ():
|
|
general_error = False
|
|
if fd in consolidate and \
|
|
writers.get(fd) is None:
|
|
to_consolidate.append(fd)
|
|
continue
|
|
cb = cbargs = None
|
|
|
|
if event & READ:
|
|
try:
|
|
cb, cbargs = readers[fd]
|
|
except KeyError:
|
|
self.remove_reader(fd)
|
|
continue
|
|
elif event & WRITE:
|
|
try:
|
|
cb, cbargs = writers[fd]
|
|
except KeyError:
|
|
self.remove_writer(fd)
|
|
continue
|
|
elif event & ERR:
|
|
general_error = True
|
|
else:
|
|
logger.info(W_UNKNOWN_EVENT, event, fd)
|
|
general_error = True
|
|
|
|
if general_error:
|
|
try:
|
|
cb, cbargs = (readers.get(fd) or
|
|
writers.get(fd))
|
|
except TypeError:
|
|
pass
|
|
|
|
if cb is None:
|
|
self.remove(fd)
|
|
continue
|
|
|
|
if isinstance(cb, generator):
|
|
try:
|
|
next(cb)
|
|
except OSError as exc:
|
|
if exc.errno != errno.EBADF:
|
|
raise
|
|
hub_remove(fd)
|
|
except StopIteration:
|
|
pass
|
|
except Exception:
|
|
hub_remove(fd)
|
|
raise
|
|
else:
|
|
try:
|
|
cb(*cbargs)
|
|
except Empty:
|
|
pass
|
|
if to_consolidate:
|
|
consolidate_callback(to_consolidate)
|
|
else:
|
|
# no sockets yet, startup is probably not done.
|
|
sleep(min(poll_timeout, 0.1))
|
|
yield
|
|
|
|
def repr_active(self):
|
|
from .debug import repr_active
|
|
return repr_active(self)
|
|
|
|
def repr_events(self, events):
|
|
from .debug import repr_events
|
|
return repr_events(self, events or [])
|
|
|
|
@cached_property
|
|
def scheduler(self):
|
|
return iter(self.timer)
|
|
|
|
@property
|
|
def loop(self):
|
|
if self._loop is None:
|
|
self._loop = self.create_loop()
|
|
return self._loop
|