Skip to content

Commit 805c604

Browse files
committed
stop cleanup task when manager is destroyed
1 parent 977363d commit 805c604

File tree

2 files changed

+10
-3
lines changed

2 files changed

+10
-3
lines changed

botcore/utils/cooldown.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import random
77
import time
88
import typing
9+
import weakref
910
from collections.abc import Awaitable, Hashable, Iterable
1011
from contextlib import suppress
1112
from dataclasses import dataclass
@@ -106,6 +107,7 @@ def __init__(self, *, cooldown_duration: float):
106107
self._periodical_cleanup(random.uniform(0, 10)),
107108
name="CooldownManager cleanup",
108109
)
110+
weakref.finalize(self, self.cleanup_task.cancel)
109111

110112
def set_cooldown(self, channel: Hashable, call_arguments: Iterable[object]) -> None:
111113
"""Set `call_arguments` arguments on cooldown in `channel`."""
@@ -145,11 +147,15 @@ async def _periodical_cleanup(self, initial_delay: float) -> None:
145147
Delete stale items every hour after waiting for `initial_delay`.
146148
147149
The `initial_delay` ensures cleanups are not running for every command at the same time.
150+
A strong reference to self is only kept while cleanup is running.
148151
"""
152+
weak_self = weakref.ref(self)
153+
del self
154+
149155
await asyncio.sleep(initial_delay)
150156
while True:
151157
await asyncio.sleep(60 * 60)
152-
self._delete_stale_items()
158+
weak_self()._delete_stale_items()
153159

154160
def _delete_stale_items(self) -> None:
155161
"""Remove expired items from internal collections."""

tests/botcore/utils/test_cooldown.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
import unittest
2+
from collections.abc import Iterable
23
from unittest.mock import patch
34

4-
from botcore.utils.cooldown import _ArgsTuple, _CommandCooldownManager
5+
from botcore.utils.cooldown import _CommandCooldownManager
56

67

7-
def create_argument_tuple(*args, **kwargs) -> _ArgsTuple:
8+
def create_argument_tuple(*args, **kwargs) -> Iterable[object]:
89
return (*args, *kwargs.items())
910

1011

0 commit comments

Comments
 (0)