Skip to content

feat(loop): add optional overlap support to allow concurrent loop executions #2765

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 9 commits into
base: master
Choose a base branch
from
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,9 @@ These changes are available on the `master` branch, but have not yet been releas
([#2714](https://github.com/Pycord-Development/pycord/pull/2714))
- Added the ability to pass a `datetime.time` object to `format_dt`
([#2747](https://github.com/Pycord-Development/pycord/pull/2747))
- Added the ability to pass an `overlap` parameter to the `loop` decorator and `Loop`
class, allowing concurrent iterations if enabled
([#2765](https://github.com/Pycord-Development/pycord/pull/2765))

### Fixed

Expand Down
26 changes: 24 additions & 2 deletions discord/ext/tasks/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,10 +59,14 @@ def __init__(
relative_delta = discord.utils.compute_timedelta(dt)
self.handle = loop.call_later(relative_delta, future.set_result, True)

def _set_result_safe(self):
if not self.future.done():
self.future.set_result(True)

def recalculate(self, dt: datetime.datetime) -> None:
self.handle.cancel()
relative_delta = discord.utils.compute_timedelta(dt)
self.handle = self.loop.call_later(relative_delta, self.future.set_result, True)
self.handle = loop.call_later(relative_delta, self._set_result_safe)

def wait(self) -> asyncio.Future[Any]:
return self.future
Expand Down Expand Up @@ -91,10 +95,12 @@ def __init__(
count: int | None,
reconnect: bool,
loop: asyncio.AbstractEventLoop,
overlap: bool,
) -> None:
self.coro: LF = coro
self.reconnect: bool = reconnect
self.loop: asyncio.AbstractEventLoop = loop
self.overlap: bool = overlap
self.count: int | None = count
self._current_loop = 0
self._handle: SleepHandle = MISSING
Expand All @@ -115,6 +121,7 @@ def __init__(
self._is_being_cancelled = False
self._has_failed = False
self._stop_next_iteration = False
self._tasks: list[asyncio.Task[Any]] = []

if self.count is not None and self.count <= 0:
raise ValueError("count must be greater than 0 or None.")
Expand Down Expand Up @@ -166,7 +173,11 @@ async def _loop(self, *args: Any, **kwargs: Any) -> None:
self._last_iteration = self._next_iteration
self._next_iteration = self._get_next_sleep_time()
try:
await self.coro(*args, **kwargs)
if self.overlap:
task = asyncio.create_task(self.coro(*args, **kwargs))
self._tasks.append(task)
else:
await self.coro(*args, **kwargs)
self._last_iteration_failed = False
backoff = ExponentialBackoff()
except self._valid_exception:
Expand All @@ -192,6 +203,9 @@ async def _loop(self, *args: Any, **kwargs: Any) -> None:

except asyncio.CancelledError:
self._is_being_cancelled = True
for task in self._tasks:
task.cancel()
await asyncio.gather(*self._tasks, return_exceptions=True)
raise
except Exception as exc:
self._has_failed = True
Expand All @@ -218,6 +232,7 @@ def __get__(self, obj: T, objtype: type[T]) -> Loop[LF]:
count=self.count,
reconnect=self.reconnect,
loop=self.loop,
overlap=self.overlap,
)
copy._injected = obj
copy._before_loop = self._before_loop
Expand Down Expand Up @@ -738,6 +753,7 @@ def loop(
count: int | None = None,
reconnect: bool = True,
loop: asyncio.AbstractEventLoop = MISSING,
overlap: bool = False,
) -> Callable[[LF], Loop[LF]]:
"""A decorator that schedules a task in the background for you with
optional reconnect logic. The decorator returns a :class:`Loop`.
Expand Down Expand Up @@ -774,6 +790,11 @@ def loop(
The loop to use to register the task, if not given
defaults to :func:`asyncio.get_event_loop`.

overlap: :class:`bool`
Whether to allow the next iteration of the loop to run even if the previous one has not completed.

.. versionadded:: 2.7

Raises
------
ValueError
Expand All @@ -793,6 +814,7 @@ def decorator(func: LF) -> Loop[LF]:
time=time,
reconnect=reconnect,
loop=loop,
overlap=overlap,
)

return decorator