From 5e1788f494d669305009b7c72cf500ccc29f6ad2 Mon Sep 17 00:00:00 2001
From: Lucio Franco <luciofranco14@gmail.com>
Date: Thu, 16 Apr 2020 11:31:58 -0400
Subject: [PATCH] rate: Fix rate limit not resetting (#439)

---
 tower/src/limit/rate/service.rs |  9 +++++----
 tower/tests/limit/rate.rs       | 36 ++++++++++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/tower/src/limit/rate/service.rs b/tower/src/limit/rate/service.rs
index 7e55ec8..7e426c6 100644
--- a/tower/src/limit/rate/service.rs
+++ b/tower/src/limit/rate/service.rs
@@ -67,7 +67,10 @@ where
         match self.state {
             State::Ready { .. } => return Poll::Ready(ready!(self.inner.poll_ready(cx))),
             State::Limited(ref mut sleep) => {
-                ready!(Pin::new(sleep).poll(cx));
+                if let Poll::Pending = Pin::new(sleep).poll(cx) {
+                    tracing::trace!("rate limit exceeded; sleeping.");
+                    return Poll::Pending;
+                }
             }
         }
 
@@ -87,9 +90,7 @@ where
                 // If the period has elapsed, reset it.
                 if now >= until {
                     until = now + self.rate.per();
-                    let rem = self.rate.num();
-
-                    self.state = State::Ready { until, rem }
+                    rem = self.rate.num();
                 }
 
                 if rem > 1 {
diff --git a/tower/tests/limit/rate.rs b/tower/tests/limit/rate.rs
index e608651..946f86f 100644
--- a/tower/tests/limit/rate.rs
+++ b/tower/tests/limit/rate.rs
@@ -9,7 +9,6 @@ async fn reaching_capacity() {
     time::pause();
 
     let rate_limit = RateLimitLayer::new(1, Duration::from_millis(100));
-
     let (mut service, mut handle) = mock::spawn_layer(rate_limit);
 
     assert_ready_ok!(service.poll_ready());
@@ -33,3 +32,38 @@ async fn reaching_capacity() {
 
     assert_eq!(response.await.unwrap(), "done");
 }
+
+#[tokio::test]
+async fn remaining_gets_reset() {
+    // This test checks for the case where the `until` state gets reset
+    // but the `rem` does not. This was a bug found `cd7dd12315706fc0860a35646b1eb7b60c50a5c1`.
+    //
+    // The main premise here is that we can make one request which should initialize the state
+    // as ready. Then we can advance the clock to put us beyond the current period. When we make
+    // subsequent requests the `rem` for the next window is continued from the previous when
+    // it should be totally reset.
+
+    time::pause();
+
+    let rate_limit = RateLimitLayer::new(3, Duration::from_millis(100));
+    let (mut service, mut handle) = mock::spawn_layer(rate_limit);
+
+    assert_ready_ok!(service.poll_ready());
+    let response = service.call("hello");
+    assert_request_eq!(handle, "hello").send_response("world");
+    assert_eq!(response.await.unwrap(), "world");
+
+    time::advance(Duration::from_millis(100)).await;
+
+    assert_ready_ok!(service.poll_ready());
+    let response = service.call("hello");
+    assert_request_eq!(handle, "hello").send_response("world");
+    assert_eq!(response.await.unwrap(), "world");
+
+    assert_ready_ok!(service.poll_ready());
+    let response = service.call("hello");
+    assert_request_eq!(handle, "hello").send_response("world");
+    assert_eq!(response.await.unwrap(), "world");
+
+    assert_ready_ok!(service.poll_ready());
+}