From d4145c8727a168eecfafc514ae1ae36357c6a6c2 Mon Sep 17 00:00:00 2001
From: Tim Vilgot Mikael Fredenberg <vilgot.fredenberg@proton.me>
Date: Sun, 25 Jun 2023 15:55:36 +0200
Subject: [PATCH] perf(gateway): reuse ratelimiter's cleanup instant (#2212)

Short and simple micro-optimization that removes a syscall from CommandRatelimiter::poll_ready + some minor cleanup.
---
 twilight-gateway/src/future.rs      |  7 ++++---
 twilight-gateway/src/ratelimiter.rs | 30 ++++++++++++++++++-----------
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/twilight-gateway/src/future.rs b/twilight-gateway/src/future.rs
index 29097c93d57..8e85c015f06 100644
--- a/twilight-gateway/src/future.rs
+++ b/twilight-gateway/src/future.rs
@@ -124,9 +124,10 @@ impl Future for NextMessageFuture<'_> {
             return Poll::Ready(NextMessageFutureOutput::SendHeartbeat);
         }
 
-        let ratelimited = this.ratelimiter.as_mut().map_or(false, |ratelimiter| {
-            ratelimiter.poll_available(cx).is_pending()
-        });
+        let ratelimited = this
+            .ratelimiter
+            .as_mut()
+            .map_or(false, |ratelimiter| ratelimiter.poll_ready(cx).is_pending());
 
         // Must poll to register waker.
         if !ratelimited
diff --git a/twilight-gateway/src/ratelimiter.rs b/twilight-gateway/src/ratelimiter.rs
index 1507a9d2b4a..bb88d21424e 100644
--- a/twilight-gateway/src/ratelimiter.rs
+++ b/twilight-gateway/src/ratelimiter.rs
@@ -74,15 +74,22 @@ impl CommandRatelimiter {
         })
     }
 
-    /// Returns when a ratelimit permit becomes available.
+    /// Waits for a permit to become available.
     pub(crate) async fn acquire(&mut self) {
-        poll_fn(|cx| self.poll_available(cx)).await;
+        poll_fn(|cx| self.poll_ready(cx)).await;
 
         self.instants.push(Instant::now() + PERIOD);
     }
 
-    /// Polls for the next time a permit is available.
-    pub(crate) fn poll_available(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+    /// Polls for readiness.
+    ///
+    /// # Return value
+    ///
+    /// The function returns:
+    ///
+    /// * `Poll::Pending` if the ratelimiter is full
+    /// * `Poll::Ready` if the ratelimiter is ready for a new permit.
+    pub(crate) fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<()> {
         if self.instants.len() != self.instants.capacity() {
             return Poll::Ready(());
         }
@@ -92,16 +99,17 @@ impl CommandRatelimiter {
         }
 
         let new_deadline = self.instants[0];
-        if new_deadline > Instant::now() {
-            tracing::trace!(?new_deadline, old_deadline = ?self.delay.deadline());
+        let now = Instant::now();
+        if new_deadline > now {
+            tracing::debug!(duration = ?(new_deadline - now), "ratelimited");
             self.delay.as_mut().reset(new_deadline);
-
-            // Register waker.
             _ = self.delay.as_mut().poll(cx);
 
             Poll::Pending
         } else {
-            let used_permits = (self.max() - self.available()).into();
+            let elapsed_permits = self.instants.partition_point(|&elapsed| elapsed <= now);
+            let used_permits = self.instants.len() - elapsed_permits;
+
             self.instants.rotate_right(used_permits);
             self.instants.truncate(used_permits);
 
@@ -242,11 +250,11 @@ mod tests {
         // Spuriously poll after registering the waker but before the timer has
         // fired.
         poll_fn(|cx| {
-            if ratelimiter.poll_available(cx).is_ready() {
+            if ratelimiter.poll_ready(cx).is_ready() {
                 return Poll::Ready(());
             };
             let deadline = ratelimiter.delay.deadline();
-            assert!(ratelimiter.poll_available(cx).is_pending());
+            assert!(ratelimiter.poll_ready(cx).is_pending());
             assert_eq!(deadline, ratelimiter.delay.deadline(), "deadline was reset");
             Poll::Pending
         })