Repository: hikari-py/hikari
Pull request #2425 (2025-06-18 20:48:21)

Revert "Only lock buckets to a single request if they are unknown"
Reverts hikari-py/hikari#2403


Patch:
diff --git a/changes/2403.optimization.md b/changes/2403.optimization.md
deleted file mode 100644
index e298f48a..00000000
--- a/changes/2403.optimization.md
+++ /dev/null
@@ -1,1 +0,0 @@
-Optimize per bucket locks (this should end up with faster bursts when making concurrent requests to a specific bucket)
diff --git a/hikari/impl/buckets.py b/hikari/impl/buckets.py
index 5f484c65..6280c0ad 100644
--- a/hikari/impl/buckets.py
+++ b/hikari/impl/buckets.py
@@ -265,8 +265,7 @@ def is_unknown(self) -> bool:
 
     def release(self) -> None:
         """Release the lock on the bucket."""
-        if self._lock.locked():
-            self._lock.release()
+        self._lock.release()
 
     @typing_extensions.override
     async def acquire(self) -> None:
@@ -282,17 +281,17 @@ async def acquire(self) -> None:
         hikari.errors.RateLimitTooLongError
             If the rate limit is longer than `max_rate_limit`.
         """
-        if self.is_unknown:
-            await self._lock.acquire()
+        await self._lock.acquire()
 
-            if self.is_unknown:
-                return
+        if self.is_unknown:
+            return
 
         now = time.monotonic()
         retry_after = self.reset_at - now
 
         if self.is_rate_limited(now) and retry_after > self._max_rate_limit:
             # Release lock before we error
+            self._lock.release()
             raise errors.RateLimitTooLongError(
                 route=self._compiled_route,
                 is_global=False,
@@ -308,6 +307,7 @@ async def acquire(self) -> None:
         global_ratelimit = self._global_ratelimit
         if global_ratelimit.reset_at and (global_ratelimit.reset_at - now) > self._max_rate_limit:
             # Release lock before we error
+            self._lock.release()
             raise errors.RateLimitTooLongError(
                 route=self._compiled_route,
                 is_global=True,


Test patch:
diff --git a/tests/hikari/impl/test_buckets.py b/tests/hikari/impl/test_buckets.py
index c367d5e4..bca92087 100644
--- a/tests/hikari/impl/test_buckets.py
+++ b/tests/hikari/impl/test_buckets.py
@@ -43,151 +43,109 @@ def template(self):
     def compiled_route(self, template):
         return routes.CompiledRoute("/foo/bar", template, "1a2b3c")
 
-    @pytest.mark.parametrize(
-        ("name", "expected"), [("spaghetti", False), ("UNKNoWN hash", False), ("UNKNOWN", True), ("UNKNOWN hash", True)]
-    )
-    def test_is_unknown(self, name, compiled_route, expected):
-        bucket = buckets.RESTBucket(name, compiled_route, mock.Mock(), float("inf"))
-
-        assert bucket.is_unknown is expected
-
     @pytest.mark.asyncio
-    async def test_usage_when_unknown(self, compiled_route):
-        bucket = buckets.RESTBucket(buckets.UNKNOWN_HASH, compiled_route, mock.Mock(), float("inf"))
-        bucket._lock = mock.Mock(acquire=mock.AsyncMock(), locked=mock.Mock(return_value=False))
+    async def test_async_context_manager(self, compiled_route):
+        with mock.patch.object(buckets.RESTBucket, "acquire", new=mock.AsyncMock()) as acquire:
+            with mock.patch.object(buckets.RESTBucket, "release") as release:
+                async with buckets.RESTBucket("spaghetti", compiled_route, object(), float("inf")):
+                    acquire.assert_awaited_once_with()
+                    release.assert_not_called()
 
-        async with bucket:
-            bucket._lock.acquire.assert_awaited_once_with()
-            bucket._lock.release.assert_not_called()
+            release.assert_called_once_with()
 
-            bucket._lock.locked.return_value = True
+    @pytest.mark.parametrize("name", ["spaghetti", buckets.UNKNOWN_HASH])
+    def test_is_unknown(self, name, compiled_route):
+        with buckets.RESTBucket(name, compiled_route, object(), float("inf")) as rl:
+            assert rl.is_unknown is (name == buckets.UNKNOWN_HASH)
 
-        bucket._lock.acquire.assert_awaited_once_with()
-        bucket._lock.release.assert_called_once_with()
+    def test_release(self, compiled_route):
+        with buckets.RESTBucket(__name__, compiled_route, object(), float("inf")) as rl:
+            rl._lock = mock.Mock()
 
-    @pytest.mark.asyncio
-    async def test_usage_when_resolved(self, compiled_route):
-        global_ratelimit = mock.Mock(acquire=mock.AsyncMock(), reset_at=None)
-        bucket = buckets.RESTBucket("resolved bucket", compiled_route, global_ratelimit, float("inf"))
-        bucket._lock = mock.Mock(acquire=mock.AsyncMock(), locked=mock.Mock(return_value=False))
+            rl.release()
 
-        async with bucket:
-            bucket._lock.acquire.assert_not_called()
-            bucket._lock.release.assert_not_called()
-
-        bucket._lock.acquire.assert_not_called()
-        bucket._lock.release.assert_not_called()
+            rl._lock.release.assert_called_once_with()
 
     def test_update_rate_limit(self, compiled_route):
-        bucket = buckets.RESTBucket("updating ratelimit test", compiled_route, mock.Mock(), float("inf"))
+        with buckets.RESTBucket(__name__, compiled_route, object(), float("inf")) as rl:
+            rl.remaining = 1
+            rl.limit = 2
+            rl.reset_at = 3
+            rl.period = 2
+
+            with mock.patch.object(hikari_date, "monotonic", return_value=4.20):
+                rl.update_rate_limit(9, 18, 27)
 
-        bucket.remaining = 1
-        bucket.limit = 2
-        bucket.reset_at = 3
-        bucket.period = 2
+            assert rl.remaining == 9
+            assert rl.limit == 18
+            assert rl.reset_at == 27
+            assert rl.period == 27 - 4.20
 
-        with mock.patch.object(hikari_date, "monotonic", return_value=4.20):
-            bucket.update_rate_limit(9, 18, 27)
+    @pytest.mark.asyncio
+    async def test_acquire_when_unknown_bucket(self, compiled_route):
+        with buckets.RESTBucket(buckets.UNKNOWN_HASH, compiled_route, object(), float("inf")) as rl:
+            rl._lock = mock.AsyncMock()
+            with mock.patch.object(rate_limits.WindowedBurstRateLimiter, "acquire") as super_acquire:
+                assert await rl.acquire() is None
 
-        assert bucket.remaining == 9
-        assert bucket.limit == 18
-        assert bucket.reset_at == 27
-        assert bucket.period == 27 - 4.20
+            rl._lock.acquire.assert_awaited_once_with()
+            super_acquire.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_acquire_when_too_long_ratelimit(self, compiled_route):
-        bucket = buckets.RESTBucket("spaghetti", compiled_route, mock.Mock(), 60)
-        bucket.reset_at = time.perf_counter() + 999999999999999999999999999
-        bucket._lock = mock.Mock(acquire=mock.AsyncMock())
+        stack = contextlib.ExitStack()
+        rl = stack.enter_context(buckets.RESTBucket("spaghetti", compiled_route, object(), 60))
+        rl._lock = mock.Mock(acquire=mock.AsyncMock())
+        rl.reset_at = time.perf_counter() + 999999999999999999999999999
+        stack.enter_context(mock.patch.object(buckets.RESTBucket, "is_rate_limited", return_value=True))
+        stack.enter_context(pytest.raises(errors.RateLimitTooLongError))
 
-        with (
-            mock.patch.object(buckets.RESTBucket, "is_rate_limited", return_value=True),
-            pytest.raises(errors.RateLimitTooLongError),
-        ):
-            await bucket.acquire()
+        with stack:
+            await rl.acquire()
 
-        bucket._lock.acquire.assert_not_called()
-        bucket._lock.release.assert_not_called()
+        rl._lock.acquire.assert_awaited_once_with()
+        rl._lock.release.assert_called_once_with()
 
     @pytest.mark.asyncio
     async def test_acquire_when_too_long_global_ratelimit(self, compiled_route):
         global_ratelimit = mock.Mock(reset_at=time.perf_counter() + 999999999999999999999999999)
 
-        bucket = buckets.RESTBucket("spaghetti", compiled_route, global_ratelimit, 1)
-        bucket._lock = mock.Mock(acquire=mock.AsyncMock())
-
-        with (
-            mock.patch.object(rate_limits.WindowedBurstRateLimiter, "acquire") as super_acquire,
-            pytest.raises(errors.RateLimitTooLongError),
-        ):
-            await bucket.acquire()
-
-        bucket._lock.acquire.assert_not_called()
-        bucket._lock.release.assert_not_called()
-        super_acquire.assert_called_once_with()
-        global_ratelimit.acquire.assert_not_called()
-
-    @pytest.mark.asyncio
-    async def test_acquire_when_unknown_bucket(self, compiled_route):
-        global_ratelimit = mock.Mock(acquire=mock.AsyncMock(), reset_at=None)
-
-        bucket = buckets.RESTBucket("UNKNOWN", compiled_route, global_ratelimit, float("inf"))
-        bucket._lock = mock.Mock(acquire=mock.AsyncMock())
-
-        with mock.patch.object(rate_limits.WindowedBurstRateLimiter, "acquire") as super_acquire:
-            await bucket.acquire()
-
-        bucket._lock.acquire.assert_awaited_once_with()
-        super_acquire.assert_not_called()
-        global_ratelimit.acquire.assert_not_called()
-
-    @pytest.mark.asyncio
-    async def test_acquire_when_resolved_while_waiting(self, compiled_route):
-        async def resolve_bucket():
-            nonlocal lock_acquire_called
-
-            bucket.resolve("some real hash")
-            lock_acquire_called += 1
-
-        global_ratelimit = mock.Mock(acquire=mock.AsyncMock(), reset_at=None)
-        bucket = buckets.RESTBucket(buckets.UNKNOWN_HASH, compiled_route, global_ratelimit, float("inf"))
-        bucket._lock = mock.Mock(acquire=resolve_bucket)
-        lock_acquire_called = 0
+        with buckets.RESTBucket("spaghetti", compiled_route, global_ratelimit, 1) as rl:
+            rl._lock = mock.Mock(acquire=mock.AsyncMock())
+            with mock.patch.object(rate_limits.WindowedBurstRateLimiter, "acquire") as super_acquire:
+                with pytest.raises(errors.RateLimitTooLongError):
+                    await rl.acquire()
 
-        with mock.patch.object(rate_limits.WindowedBurstRateLimiter, "acquire") as super_acquire:
-            await bucket.acquire()
-
-        super_acquire.assert_awaited_once_with()
-        assert lock_acquire_called == 1
-        global_ratelimit.acquire.assert_awaited_once_with()
+            rl._lock.acquire.assert_awaited_once_with()
+            super_acquire.assert_awaited_once_with()
+            rl._lock.release.assert_called_once_with()
+            global_ratelimit.acquire.assert_not_called()
 
     @pytest.mark.asyncio
-    async def test_acquire_when_resolved_bucket(self, compiled_route):
+    async def test_acquire(self, compiled_route):
         global_ratelimit = mock.Mock(acquire=mock.AsyncMock(), reset_at=None)
-        bucket = buckets.RESTBucket("spaghetti", compiled_route, global_ratelimit, float("inf"))
-        bucket._lock = mock.Mock()
 
-        with mock.patch.object(rate_limits.WindowedBurstRateLimiter, "acquire") as super_acquire:
-            await bucket.acquire()
+        with buckets.RESTBucket("spaghetti", compiled_route, global_ratelimit, float("inf")) as rl:
+            rl._lock = mock.AsyncMock()
+            with mock.patch.object(rate_limits.WindowedBurstRateLimiter, "acquire") as super_acquire:
+                await rl.acquire()
 
-        super_acquire.assert_awaited_once_with()
-        bucket._lock.acquire.assert_not_called()
-        global_ratelimit.acquire.assert_awaited_once_with()
+            super_acquire.assert_awaited_once_with()
+            rl._lock.acquire.assert_awaited_once_with()
+            global_ratelimit.acquire.assert_awaited_once_with()
 
     def test_resolve_when_not_unknown(self, compiled_route):
-        bucket = buckets.RESTBucket("spaghetti", compiled_route, mock.Mock(), float("inf"))
-
-        with pytest.raises(RuntimeError, match=r"Cannot resolve known bucket"):
-            bucket.resolve("test")
+        with buckets.RESTBucket("spaghetti", compiled_route, object(), float("inf")) as rl:
+            with pytest.raises(RuntimeError, match=r"Cannot resolve known bucket"):
+                rl.resolve("test")
 
-        assert bucket.name == "spaghetti"
+            assert rl.name == "spaghetti"
 
     def test_resolve(self, compiled_route):
-        bucket = buckets.RESTBucket(buckets.UNKNOWN_HASH, compiled_route, mock.Mock(), float("inf"))
-
-        bucket.resolve("test")
+        with buckets.RESTBucket(buckets.UNKNOWN_HASH, compiled_route, object(), float("inf")) as rl:
+            rl.resolve("test")
 
-        assert bucket.name == "test"
+            assert rl.name == "test"
 
 
 class TestRESTBucketManager: