Repository: hikari-py/hikari
Pull request #1529 (2023-03-01 21:44:57)

Await bucket manager gc task to completion when closing
### Checklist
<!-- Make sure to tick all the following boxes by putting an `x` in between (like this `[x]`) -->
- [x] I have run `nox` and all the pipelines have passed.
- [x] I have made unittests according to the code I have added/modified/deleted.



Patch:
diff --git a/changes/1529.bugfix.md b/changes/1529.bugfix.md
new file mode 100644
index 00000000..c9c0331e
--- /dev/null
+++ b/changes/1529.bugfix.md
@@ -0,0 +1,1 @@
+Await bucket manager gc task to completion when closing
diff --git a/hikari/api/rest.py b/hikari/api/rest.py
index 4175d2c2..c9b07da5 100644
--- a/hikari/api/rest.py
+++ b/hikari/api/rest.py
@@ -265,7 +265,7 @@ async def edit_channel(
             If provided, the auto archive duration Discord's end user client
             should default to when creating threads in this channel.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         default_thread_rate_limit_per_user : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
@@ -300,7 +300,7 @@ async def edit_channel(
             If provided, the new auto archive duration for this thread. This
             only applies to threads.
 
-            This should be either 60, 1440, 4320 or 10080 seconds, as of
+            This should be either 60, 1440, 4320 or 10080 minutes, as of
             writing.
         applied_tags : hikari.undefined.UndefinedOr[typing.Sequence[hikari.channels.ForumTag]]
             If provided, the new tags to apply to the thread. This only applies
@@ -4056,7 +4056,7 @@ async def create_guild_text_channel(
             If provided, the auto archive duration Discord's end user client
             should default to when creating threads in this channel.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         reason : hikari.undefined.UndefinedOr[str]
@@ -4134,7 +4134,7 @@ async def create_guild_news_channel(
             If provided, the auto archive duration Discord's end user client
             should default to when creating threads in this channel.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         reason : hikari.undefined.UndefinedOr[str]
@@ -4220,7 +4220,7 @@ async def create_guild_forum_channel(
             If provided, the auto archive duration Discord's end user client
             should default to when creating threads in this channel.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         default_thread_rate_limit_per_user : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
@@ -4499,7 +4499,7 @@ async def create_message_thread(
         auto_archive_duration : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
             If provided, how long the thread should remain inactive until it's archived.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         rate_limit_per_user : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
@@ -4566,7 +4566,7 @@ async def create_thread(
         auto_archive_duration : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
             If provided, how long the thread should remain inactive until it's archived.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         invitable : undefined.UndefinedOr[bool]
@@ -4734,7 +4734,7 @@ async def create_forum_post(
         auto_archive_duration : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
             If provided, how long the post should remain inactive until it's archived.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         rate_limit_per_user : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
diff --git a/hikari/channels.py b/hikari/channels.py
index 1581662f..baa02f4e 100644
--- a/hikari/channels.py
+++ b/hikari/channels.py
@@ -1037,7 +1037,7 @@ async def edit(
             If provided, the auto archive duration Discord's end user client
             should default to when creating threads in this channel.
 
-            This should be either 60, 1440, 4320 or 10080 seconds, as of
+            This should be either 60, 1440, 4320 or 10080 minutes, as of
             writing.
         flags : hikari.undefined.UndefinedOr[ChannelFlag]
             If provided, the new channel flags to use for the channel. This can
@@ -1050,7 +1050,7 @@ async def edit(
             If provided, the new auto archive duration for this thread. This
             only applies to threads.
 
-            This should be either 60, 1440, 4320 or 10080 seconds, as of
+            This should be either 60, 1440, 4320 or 10080 minutes, as of
             writing.
         locked : hikari.undefined.UndefinedOr[bool]
             If provided, the new locked state for the thread. This only applies
diff --git a/hikari/guilds.py b/hikari/guilds.py
index ddbf3438..34893db9 100644
--- a/hikari/guilds.py
+++ b/hikari/guilds.py
@@ -2113,7 +2113,7 @@ async def create_forum_channel(
             If provided, the auto archive duration Discord's end user client
             should default to when creating threads in this channel.
 
-            This should be either 60, 1440, 4320 or 10080 seconds and, as of
+            This should be either 60, 1440, 4320 or 10080 minutes and, as of
             writing, ignores the parent channel's set default_auto_archive_duration
             when passed as `hikari.undefined.UNDEFINED`.
         default_thread_rate_limit_per_user : hikari.undefined.UndefinedOr[hikari.internal.time.Intervalish]
diff --git a/hikari/impl/buckets.py b/hikari/impl/buckets.py
index 807e5ae8..52ee79bc 100644
--- a/hikari/impl/buckets.py
+++ b/hikari/impl/buckets.py
@@ -385,7 +385,6 @@ class RESTBucketManager:
         "_routes_to_hashes",
         "_real_hashes_to_buckets",
         "_global_ratelimit",
-        "_closed_event",
         "_gc_task",
         "_max_rate_limit",
     )
@@ -393,7 +392,6 @@ class RESTBucketManager:
     def __init__(self, max_rate_limit: float) -> None:
         self._routes_to_hashes: typing.Dict[routes.Route, str] = {}
         self._real_hashes_to_buckets: typing.Dict[str, RESTBucket] = {}
-        self._closed_event: typing.Optional[asyncio.Event] = None
         self._gc_task: typing.Optional[asyncio.Task[None]] = None
         self._max_rate_limit = max_rate_limit
         self._global_ratelimit = rate_limits.ManualRateLimiter()
@@ -405,7 +403,7 @@ def max_rate_limit(self) -> float:
     @property
     def is_alive(self) -> bool:
         """Whether the component is alive."""
-        return self._closed_event is not None
+        return self._gc_task is not None
 
     def start(self, poll_period: float = 20.0, expire_after: float = 10.0) -> None:
         """Start this ratelimiter up.
@@ -426,22 +424,19 @@ def start(self, poll_period: float = 20.0, expire_after: float = 10.0) -> None:
             result. Using `0` will make the bucket get garbage collected as soon
             as the rate limit has reset. Defaults to `10` seconds.
         """
-        if self._closed_event:
+        if self._gc_task:
             raise errors.ComponentStateConflictError("Cannot start an active bucket manager")
 
         # Assert is in running loop
         asyncio.get_running_loop()
 
-        self._closed_event = asyncio.Event()
         self._gc_task = asyncio.create_task(self._gc(poll_period, expire_after))
 
-    def close(self) -> None:
+    async def close(self) -> None:
         """Close the garbage collector and kill any tasks waiting on ratelimits."""
-        if not self._closed_event:
+        if not self._gc_task:
             raise errors.ComponentStateConflictError("Cannot interact with an inactive bucket manager")
 
-        assert self._closed_event is not None
-
         for bucket in self._real_hashes_to_buckets.values():
             bucket.close()
 
@@ -449,49 +444,24 @@ def close(self) -> None:
         self._real_hashes_to_buckets.clear()
         self._routes_to_hashes.clear()
 
-        if self._gc_task is not None:
-            self._gc_task.cancel()
-            self._gc_task = None
-
-        self._closed_event.set()
-        self._closed_event = None
-
-    async def _gc(self, poll_period: float, expire_after: float) -> None:
-        """Run the garbage collector loop.
+        self._gc_task.cancel()
 
-        This is designed to run in the background and manage removing unused
-        route references from the rate-limiter collection to save memory.
+        try:
+            await self._gc_task
+        except asyncio.CancelledError:
+            pass
 
-        This will run forever until `RESTBucketManager.closed_event` is set.
-        This will invoke `RESTBucketManager.do_gc_pass` periodically.
+        self._gc_task = None
 
-        .. warning::
-            You generally have no need to invoke this directly. Use
-            `RESTBucketManager.start` and `RESTBucketManager.close` to control
-            this instead.
-
-        Parameters
-        ----------
-        poll_period : float
-            The period to poll at.
-        expire_after : float
-            Time after which the last `reset_at` was hit for a bucket to
-            remove it. Higher values will retain unneeded ratelimit info for
-            longer, but may produce more effective ratelimiting logic as a
-            result. Using `0` will make the bucket get garbage collected as soon
-            as the rate limit has reset.
-        """
+    async def _gc(self, poll_period: float, expire_after: float) -> None:
         # Prevent filling memory increasingly until we run out by removing dead buckets every 20s
         # Allocations are somewhat cheap if we only do them every so-many seconds, after all.
         _LOGGER.log(ux.TRACE, "rate limit garbage collector started")
 
-        assert self._closed_event is not None
-        while not self._closed_event.is_set():
-            try:
-                await asyncio.wait_for(self._closed_event.wait(), timeout=poll_period)
-            except asyncio.TimeoutError:
-                _LOGGER.log(ux.TRACE, "performing rate limit garbage collection pass")
-                self._purge_stale_buckets(expire_after)
+        while True:
+            await asyncio.sleep(poll_period)
+            _LOGGER.log(ux.TRACE, "performing rate limit garbage collection pass")
+            self._purge_stale_buckets(expire_after)
 
     def _purge_stale_buckets(self, expire_after: float) -> None:
         buckets_to_purge: typing.List[str] = []
@@ -550,7 +520,7 @@ def acquire_bucket(
         typing.AsyncContextManager
             The context manager to use during the duration of the request.
         """
-        if not self._closed_event:
+        if not self._gc_task:
             raise errors.ComponentStateConflictError("Cannot interact with an inactive bucket manager")
 
         authentication_hash = _create_authentication_hash(authentication)
@@ -600,7 +570,7 @@ def update_rate_limits(
         reset_after : float
             The `X-RateLimit-Reset-After` header cast to a `float`.
         """
-        if not self._closed_event:
+        if not self._gc_task:
             raise errors.ComponentStateConflictError("Cannot interact with an inactive bucket manager")
 
         self._routes_to_hashes[compiled_route.route] = bucket_header
diff --git a/hikari/impl/rest.py b/hikari/impl/rest.py
index e21a8042..5e1b42d3 100644
--- a/hikari/impl/rest.py
+++ b/hikari/impl/rest.py
@@ -352,7 +352,7 @@ async def close(self) -> None:
             raise errors.ComponentStateConflictError("Rest app is not running")
 
         await self._client_session.close()
-        self._bucket_manager.close()
+        await self._bucket_manager.close()
 
     @typing.overload
     def acquire(self, token: typing.Optional[rest_api.TokenStrategy] = None) -> RESTClientImpl:
@@ -645,7 +645,7 @@ async def close(self) -> None:
             self._client_session = None
 
         if self._bucket_manager_owner:
-            self._bucket_manager.close()
+            await self._bucket_manager.close()
 
     def start(self) -> None:
         """Start the HTTP client.


Test patch:
diff --git a/tests/hikari/impl/test_buckets.py b/tests/hikari/impl/test_buckets.py
index 5ee3a0bf..fb802c8e 100644
--- a/tests/hikari/impl/test_buckets.py
+++ b/tests/hikari/impl/test_buckets.py
@@ -151,7 +151,7 @@ class TestRESTBucketManager:
     @pytest.fixture()
     def bucket_manager(self):
         manager = buckets.RESTBucketManager(max_rate_limit=float("inf"))
-        manager._closed_event = mock.Mock(wait=mock.AsyncMock(), is_set=mock.Mock(return_value=False))
+        manager._gc_task = object()
 
         return manager
 
@@ -161,32 +161,43 @@ def test_max_rate_limit_property(self, bucket_manager):
         assert bucket_manager.max_rate_limit is bucket_manager._max_rate_limit
 
     @pytest.mark.asyncio()
-    async def test_close_closes_all_buckets(self, bucket_manager):
+    async def test_close(self, bucket_manager):
+        class GcTaskMock:
+            def __init__(self):
+                self._awaited_count = 0
+                self.cancel = mock.Mock()
+
+            def __await__(self):
+                if False:
+                    yield  # Turns it into a generator
+
+                self._awaited_count += 1
+
+            def __call__(self):
+                return self
+
+            def assert_awaited_once(self):
+                assert self._awaited_count == 1
+
         buckets_array = [mock.Mock() for _ in range(30)]
         bucket_manager._real_hashes_to_buckets = {f"blah{i}": b for i, b in enumerate(buckets_array)}
+        bucket_manager._gc_task = gc_task = GcTaskMock()
 
-        bucket_manager.close()
+        await bucket_manager.close()
 
         assert bucket_manager._real_hashes_to_buckets == {}
 
         for i, b in enumerate(buckets_array):
             b.close.assert_called_once(), i
 
-    @pytest.mark.asyncio()
-    async def test_close_sets_closed_event(self, bucket_manager):
-        closed_event = mock.Mock()
-        bucket_manager._closed_event = closed_event
-
-        bucket_manager.close()
-
-        assert bucket_manager._closed_event is None
-        closed_event.set.assert_called_once()
+        assert bucket_manager._gc_task is None
+        gc_task.cancel.assert_called_once_with()
+        gc_task.assert_awaited_once()
 
     @pytest.mark.asyncio()
     async def test_start(self, bucket_manager):
-        bucket_manager._closed_event = None
+        bucket_manager._gc_task = None
 
-        assert bucket_manager._gc_task is None
         bucket_manager.start()
         assert bucket_manager._gc_task is not None
 
@@ -199,41 +210,18 @@ async def test_start(self, bucket_manager):
 
     @pytest.mark.asyncio()
     async def test_start_when_already_started(self, bucket_manager):
-        bucket_manager._closed_event = object()
+        bucket_manager._gc_task = object()
 
         with pytest.raises(errors.ComponentStateConflictError):
             bucket_manager.start()
 
-    @pytest.mark.asyncio()
-    async def test_gc_polls_until_closed_event_set(self, event_loop, bucket_manager):
-        bucket_manager._closed_event = asyncio.Event()
-
-        # Start the gc and initial assertions
-        task = event_loop.create_task(bucket_manager._gc(0.001, float("inf")))
-        assert not task.done()
-
-        # [First poll] event not set => shouldn't complete the task
-        await asyncio.sleep(0.001)
-        assert not task.done()
-
-        # [Second poll] event not set during poll => shouldn't complete the task
-        await asyncio.sleep(0.001)
-        bucket_manager._closed_event.set()
-        assert not task.done()
-
-        # [Third poll] event set => should complete the task
-        await asyncio.sleep(0.001)
-        assert task.done()
-
     @pytest.mark.asyncio()
     async def test_gc_makes_gc_pass(self, bucket_manager):
         class ExitError(Exception):
             ...
 
-        bucket_manager._closed_event.wait = mock.Mock()
-
         with mock.patch.object(buckets.RESTBucketManager, "_purge_stale_buckets") as purge_stale_buckets:
-            with mock.patch.object(asyncio, "wait_for", side_effect=[asyncio.TimeoutError, ExitError]):
+            with mock.patch.object(asyncio, "sleep", side_effect=[None, ExitError]):
                 with pytest.raises(ExitError):
                     await bucket_manager._gc(0.001, 33)
 
@@ -431,7 +419,7 @@ async def test_update_rate_limits_updates_params(self, bucket_manager):
                 bucket_manager.update_rate_limits(route, "auth", "123", 22, 23, 5.32)
                 bucket.update_rate_limit.assert_called_once_with(22, 23, 27 + 5.32)
 
-    @pytest.mark.parametrize(("closed_event", "is_alive"), [(None, False), ("some", True)])
-    def test_is_alive(self, bucket_manager, closed_event, is_alive):
-        bucket_manager._closed_event = closed_event
+    @pytest.mark.parametrize(("gc_task", "is_alive"), [(None, False), ("some", True)])
+    def test_is_alive(self, bucket_manager, gc_task, is_alive):
+        bucket_manager._gc_task = gc_task
         assert bucket_manager.is_alive is is_alive
diff --git a/tests/hikari/impl/test_rest.py b/tests/hikari/impl/test_rest.py
index 732164f1..1a062e5d 100644
--- a/tests/hikari/impl/test_rest.py
+++ b/tests/hikari/impl/test_rest.py
@@ -679,6 +679,7 @@ def test_token_type_property(self, rest_client):
     async def test_close(self, rest_client, client_session_owner, bucket_manager_owner):
         rest_client._close_event = mock_close_event = mock.Mock()
         rest_client._client_session.close = client_close = mock.AsyncMock()
+        rest_client._bucket_manager.close = bucket_close = mock.AsyncMock()
         rest_client._client_session_owner = client_session_owner
         rest_client._bucket_manager_owner = bucket_manager_owner
 
@@ -695,7 +696,7 @@ async def test_close(self, rest_client, client_session_owner, bucket_manager_own
             assert rest_client._client_session is not None
 
         if bucket_manager_owner:
-            rest_client._bucket_manager.close.assert_called_once_with()
+            bucket_close.assert_awaited_once_with()
         else:
             rest_client._bucket_manager.assert_not_called()