Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion bot/exts/filtering/_ui/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from bot.exts.filtering._filter_context import Event, FilterContext
from bot.exts.filtering._filter_lists import FilterList
from bot.exts.filtering._utils import FakeContext, normalize_type
from bot.utils.lock import lock_arg
from bot.utils.async_utils import lock_arg
from bot.utils.messages import format_channel, format_user, upload_log

log = get_logger(__name__)
Expand Down
2 changes: 1 addition & 1 deletion bot/exts/filtering/filtering.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,8 @@
from bot.exts.utils.snekbox._io import FileAttachment
from bot.log import get_logger
from bot.pagination import LinePaginator
from bot.utils.async_utils import lock_arg
from bot.utils.channel import is_mod_channel
from bot.utils.lock import lock_arg
from bot.utils.message_cache import MessageCache

log = get_logger(__name__)
Expand Down
2 changes: 1 addition & 1 deletion bot/exts/info/doc/_cog.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from bot.converters import Inventory, PackageName, ValidURL
from bot.log import get_logger
from bot.pagination import LinePaginator
from bot.utils.lock import SharedEvent, lock
from bot.utils.async_utils import SharedEvent, lock
from bot.utils.messages import send_denial, wait_for_deletion

from . import NAMESPACE, PRIORITY_PACKAGES, _batch_parser, doc_cache
Expand Down
2 changes: 1 addition & 1 deletion bot/exts/info/doc/_redis_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from async_rediscache.types.base import RedisObject

from bot.log import get_logger
from bot.utils.lock import lock
from bot.utils.async_utils import lock

from ._doc_item import DocItem

Expand Down
273 changes: 168 additions & 105 deletions bot/exts/moderation/clean.py

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion bot/exts/moderation/silence.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from bot.bot import Bot
from bot.converters import HushDurationConverter
from bot.log import get_logger
from bot.utils.lock import LockedResourceError, lock, lock_arg
from bot.utils.async_utils import LockedResourceError, lock, lock_arg

log = get_logger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion bot/exts/utils/reminders.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,8 @@
from bot.log import get_logger
from bot.pagination import LinePaginator
from bot.utils import time
from bot.utils.async_utils import lock_arg
from bot.utils.checks import has_any_role_check, has_no_roles_check
from bot.utils.lock import lock_arg
from bot.utils.messages import send_denial

log = get_logger(__name__)
Expand Down
2 changes: 1 addition & 1 deletion bot/exts/utils/snekbox/_cog.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
from bot.exts.utils.snekbox._eval import EvalJob, EvalResult
from bot.exts.utils.snekbox._io import FileAttachment
from bot.log import get_logger
from bot.utils.lock import LockedResourceError, lock_arg
from bot.utils.async_utils import LockedResourceError, lock_arg

if TYPE_CHECKING:
from bot.exts.filtering.filtering import Filtering
Expand Down
31 changes: 31 additions & 0 deletions bot/utils/lock.py → bot/utils/async_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,37 @@ async def wait(self) -> None:
await self._event.wait()


class AsyncExecutor:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should there be a public interface for cancelling the tasks? Might be needed if the bot is shutting down.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good point

"""Execute coroutines with a pool limit for concurrent execution."""

def __init__(self, *, limit: int):
self._semaphore = asyncio.Semaphore(limit)
self._running_tasks = set()

def submit[T](self, coro: Awaitable[T]) -> asyncio.Task[T]:
"""Schedule the coroutine on the event loop, and ensure cleanup."""
task = asyncio.create_task(self.execute(coro))
self._running_tasks.add(task)
return task

async def execute[T](self, coro: Awaitable[T]) -> T:
"""Execute the coroutine once there is an available slot in the pool."""
async with self._semaphore:
return await coro

async def gather(self, return_exceptions: bool = False) -> list[Any]:
"""Wait for all submitted coroutines to finish execution."""
result = await asyncio.gather(*self._running_tasks, return_exceptions=return_exceptions)
self._running_tasks.clear()
return result

def cancel_all(self) -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we wait for the tasks to be cancelled or is it fine to fire and forget?

"""Cancel all running tasks."""
for task in self._running_tasks:
task.cancel()
self._running_tasks.clear()


def lock(
namespace: Hashable,
resource_id: ResourceId,
Expand Down
7 changes: 5 additions & 2 deletions tests/bot/exts/moderation/test_clean.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import unittest
from unittest.mock import AsyncMock, MagicMock, patch

import arrow

from bot.exts.moderation.clean import Clean
from tests.helpers import MockBot, MockContext, MockGuild, MockMember, MockMessage, MockRole, MockTextChannel

Expand All @@ -20,7 +22,8 @@ def setUp(self):
self.cog._modlog_cleaned_messages = AsyncMock(return_value=self.log_url)

self.cog._use_cache = MagicMock(return_value=True)
self.cog._delete_found = AsyncMock(return_value=[42, 84])
self.cog._use_api = MagicMock(return_value=False)
self.cog._delete_bulk = AsyncMock(return_value=([42, 84], {}))

@patch("bot.exts.moderation.clean.is_mod_channel")
async def test_clean_deletes_invocation_in_non_mod_channel(self, mod_channel_check):
Expand Down Expand Up @@ -51,7 +54,7 @@ async def test_clean_doesnt_attempt_deletion_when_attempt_delete_invocation_is_f
await self.cog._clean_messages(
self.ctx,
None,
first_limit=MockMessage(),
first_limit=arrow.utcnow().datetime,
attempt_delete_invocation=False,
),
self.log_url,
Expand Down