Repository: hikari-py/hikari
Pull request #1814 (2024-01-21 09:59:32)

Allow subscribing to generic events


Patch:
diff --git a/changes/1813.bugfix.md b/changes/1813.bugfix.md
new file mode 100644
index 00000000..a61c166f
--- /dev/null
+++ b/changes/1813.bugfix.md
@@ -0,0 +1,1 @@
+Shard rate-limiters are now reset per websocket connection, avoiding a rare issue where a persistent network issue could allow the shard to be rate-limited
diff --git a/changes/1814.feature.md b/changes/1814.feature.md
new file mode 100644
index 00000000..e7eb091d
--- /dev/null
+++ b/changes/1814.feature.md
@@ -0,0 +1,1 @@
+Allow subscribing to generic events
diff --git a/hikari/impl/event_manager_base.py b/hikari/impl/event_manager_base.py
index fb18e4a8..5d3e86db 100644
--- a/hikari/impl/event_manager_base.py
+++ b/hikari/impl/event_manager_base.py
@@ -379,6 +379,10 @@ def _enabled_for_event(self, event_type: typing.Type[base_events.Event], /) -> b
         return False
 
     def _check_event(self, event_type: typing.Type[typing.Any], nested: int) -> None:
+        # Extract the underlying type from generics
+        if (origin_type := typing.get_origin(event_type)) is not None:
+            event_type = origin_type
+
         try:
             is_event = issubclass(event_type, base_events.Event)
         except TypeError:
diff --git a/hikari/impl/rate_limits.py b/hikari/impl/rate_limits.py
index abd25b55..ac4a9285 100644
--- a/hikari/impl/rate_limits.py
+++ b/hikari/impl/rate_limits.py
@@ -86,7 +86,7 @@ class BurstRateLimiter(BaseRateLimiter, abc.ABC):
     complete logic for safely aborting any pending tasks when being shut down.
     """
 
-    __slots__: typing.Sequence[str] = ("name", "throttle_task", "queue", "_closed")
+    __slots__: typing.Sequence[str] = ("name", "throttle_task", "queue")
 
     name: str
     """The name of the rate limiter."""
@@ -101,7 +101,6 @@ def __init__(self, name: str) -> None:
         self.name = name
         self.throttle_task = None
         self.queue = []
-        self._closed = False
 
     @abc.abstractmethod
     async def acquire(self) -> None:
@@ -112,13 +111,7 @@ async def acquire(self) -> None:
         """
 
     def close(self) -> None:
-        """Close the rate limiter, and shut down any pending tasks.
-
-        Once this is invoked, you should not reuse this object.
-        """
-        if self._closed:
-            return
-
+        """Close the rate limiter, and shut down any pending tasks."""
         if self.throttle_task is not None:
             self.throttle_task.cancel()
             self.throttle_task = None
@@ -134,7 +127,6 @@ def close(self) -> None:
             _LOGGER.debug("%s rate limiter closed with %s pending tasks!", self.name, failed_tasks)
         else:
             _LOGGER.debug("%s rate limiter closed", self.name)
-        self._closed = True
 
     @property
     def is_empty(self) -> bool:
diff --git a/hikari/impl/shard.py b/hikari/impl/shard.py
index 48f3b9b4..112966a7 100644
--- a/hikari/impl/shard.py
+++ b/hikari/impl/shard.py
@@ -838,9 +838,13 @@ async def _connect(self) -> typing.Tuple[asyncio.Task[None], ...]:
         )
         poll_events_task = asyncio.create_task(self._poll_events(), name=f"poll events (shard {self._shard_id})")
 
+        # Rate-limits are imposed per websocket connection
+        self._total_rate_limit.close()
+        self._non_priority_rate_limit.close()
+
         # Perform handshake
         if self._seq is None:
-            self._logger.debug("identifying with new session")
+            self._logger.info("identifying with new session")
             await self._send_json(
                 {
                     _OP: _IDENTIFY,
@@ -860,7 +864,7 @@ async def _connect(self) -> typing.Tuple[asyncio.Task[None], ...]:
                 }
             )
         else:
-            self._logger.debug("resuming session %s", self._session_id)
+            self._logger.info("resuming session %s", self._session_id)
             await self._send_json(
                 {_OP: _RESUME, _D: {"token": self._token, "seq": self._seq, "session_id": self._session_id}}
             )
@@ -924,7 +928,7 @@ async def _keep_alive(self) -> None:
                 backoff.reset()
 
             except errors.GatewayError as ex:
-                self._logger.error("encountered generic gateway error", exc_info=ex)
+                self._logger.error("encountered gateway error", exc_info=ex)
                 raise
 
             except asyncio.CancelledError:


Test patch:
diff --git a/tests/hikari/impl/test_event_manager_base.py b/tests/hikari/impl/test_event_manager_base.py
index 7e16aee6..09d0b79d 100644
--- a/tests/hikari/impl/test_event_manager_base.py
+++ b/tests/hikari/impl/test_event_manager_base.py
@@ -673,7 +673,24 @@ def test__check_event_when_no_intents_required(self, event_manager):
         get_intents.assert_called_once_with(member_events.MemberCreateEvent)
         warn.assert_not_called()
 
-    def test__check_intents_when_intents_correct(self, event_manager):
+    def test__check_event_when_generic_event(self, event_manager):
+        T = typing.TypeVar("T")
+
+        class GenericEvent(typing.Generic[T], base_events.Event):
+            ...
+
+        event_manager._intents = intents.Intents.GUILD_MEMBERS
+
+        with mock.patch.object(
+            base_events, "get_required_intents_for", return_value=intents.Intents.GUILD_MEMBERS
+        ) as get_intents:
+            with mock.patch.object(warnings, "warn") as warn:
+                event_manager._check_event(GenericEvent[int], 0)
+
+        get_intents.assert_called_once_with(GenericEvent)
+        warn.assert_not_called()
+
+    def test__check_event_when_intents_correct(self, event_manager):
         event_manager._intents = intents.Intents.GUILD_EMOJIS | intents.Intents.GUILD_MEMBERS
 
         with mock.patch.object(
@@ -685,7 +702,7 @@ def test__check_intents_when_intents_correct(self, event_manager):
         get_intents.assert_called_once_with(member_events.MemberCreateEvent)
         warn.assert_not_called()
 
-    def test__check_intents_when_intents_incorrect(self, event_manager):
+    def test__check_event_when_intents_incorrect(self, event_manager):
         event_manager._intents = intents.Intents.GUILD_EMOJIS
 
         with mock.patch.object(
diff --git a/tests/hikari/impl/test_rate_limits.py b/tests/hikari/impl/test_rate_limits.py
index 2855e65d..c89430a6 100644
--- a/tests/hikari/impl/test_rate_limits.py
+++ b/tests/hikari/impl/test_rate_limits.py
@@ -65,16 +65,18 @@ def test_is_empty(self, queue, is_empty, mock_burst_limiter):
         mock_burst_limiter.queue = queue
         assert mock_burst_limiter.is_empty is is_empty
 
-    def test_close_removes_all_futures_from_queue(self, mock_burst_limiter):
-        event_loop = asyncio.get_event_loop()
+    @pytest.mark.asyncio()
+    async def test_close_removes_all_futures_from_queue(self, mock_burst_limiter):
+        event_loop = asyncio.get_running_loop()
         mock_burst_limiter.throttle_task = None
         futures = [event_loop.create_future() for _ in range(10)]
         mock_burst_limiter.queue = list(futures)
         mock_burst_limiter.close()
         assert len(mock_burst_limiter.queue) == 0
 
-    def test_close_cancels_all_futures_pending_when_futures_pending(self, mock_burst_limiter):
-        event_loop = asyncio.get_event_loop()
+    @pytest.mark.asyncio()
+    async def test_close_cancels_all_futures_pending_when_futures_pending(self, mock_burst_limiter):
+        event_loop = asyncio.get_running_loop()
         mock_burst_limiter.throttle_task = None
         futures = [event_loop.create_future() for _ in range(10)]
         mock_burst_limiter.queue = list(futures)
@@ -82,21 +84,24 @@ def test_close_cancels_all_futures_pending_when_futures_pending(self, mock_burst
         for i, future in enumerate(futures):
             assert future.cancelled(), f"future {i} was not cancelled"
 
-    def test_close_is_silent_when_no_futures_pending(self, mock_burst_limiter):
+    @pytest.mark.asyncio()
+    async def test_close_is_silent_when_no_futures_pending(self, mock_burst_limiter):
         mock_burst_limiter.throttle_task = None
         mock_burst_limiter.queue = []
         mock_burst_limiter.close()
         assert True, "passed successfully"
 
-    def test_close_cancels_throttle_task_if_running(self, mock_burst_limiter):
-        event_loop = asyncio.get_event_loop()
+    @pytest.mark.asyncio()
+    async def test_close_cancels_throttle_task_if_running(self, mock_burst_limiter):
+        event_loop = asyncio.get_running_loop()
         task = event_loop.create_future()
         mock_burst_limiter.throttle_task = task
         mock_burst_limiter.close()
         assert mock_burst_limiter.throttle_task is None, "task was not overwritten with None"
         assert task.cancelled(), "throttle_task is not cancelled"
 
-    def test_close_when_closed(self, mock_burst_limiter):
+    @pytest.mark.asyncio()
+    async def test_close_when_closed(self, mock_burst_limiter):
         # Double-running shouldn't do anything adverse.
         mock_burst_limiter.close()
         mock_burst_limiter.close()
@@ -105,7 +110,7 @@ def test_close_when_closed(self, mock_burst_limiter):
 class TestManualRateLimiter:
     @pytest.mark.asyncio()
     async def test_acquire_returns_completed_future_if_throttle_task_is_None(self):
-        event_loop = asyncio.get_event_loop()
+        event_loop = asyncio.get_running_loop()
 
         with rate_limits.ManualRateLimiter() as limiter:
             limiter.throttle_task = None
@@ -117,7 +122,7 @@ async def test_acquire_returns_completed_future_if_throttle_task_is_None(self):
 
     @pytest.mark.asyncio()
     async def test_acquire_returns_incomplete_future_if_throttle_task_is_not_None(self):
-        event_loop = asyncio.get_event_loop()
+        event_loop = asyncio.get_running_loop()
 
         with rate_limits.ManualRateLimiter() as limiter:
             limiter.throttle_task = event_loop.create_future()
@@ -129,7 +134,7 @@ async def test_acquire_returns_incomplete_future_if_throttle_task_is_not_None(se
 
     @pytest.mark.asyncio()
     async def test_acquire_places_future_on_queue_if_throttle_task_is_not_None(self):
-        event_loop = asyncio.get_event_loop()
+        event_loop = asyncio.get_running_loop()
 
         with rate_limits.ManualRateLimiter() as limiter:
             limiter.throttle_task = event_loop.create_future()
@@ -163,7 +168,7 @@ async def test_throttle_schedules_throttle(self):
 
     @pytest.mark.asyncio()
     async def test_throttle_chews_queue_completing_futures(self):
-        event_loop = asyncio.get_event_loop()
+        event_loop = asyncio.get_running_loop()
 
         with rate_limits.ManualRateLimiter() as limiter:
             futures = [event_loop.create_future() for _ in range(10)]
@@ -174,7 +179,7 @@ async def test_throttle_chews_queue_completing_futures(self):
 
     @pytest.mark.asyncio()
     async def test_throttle_sleeps_before_popping_queue(self):
-        event_loop = asyncio.get_event_loop()
+        event_loop = asyncio.get_running_loop()
         # GIVEN
         slept_at = float("nan")
         popped_at = []