Merge pull request #163 from python-astrodynamics/feature/ratelimit-accounting

Count all requests against the rate limit

Author: RazerM

newsfragments/163.fixed.2.rst

@@ -0,0 +1 @@
+Model definition requests are now rate limited like all other requests.

newsfragments/163.fixed.rst

@@ -0,0 +1,2 @@
+Requests sent after a rate limit wait are now counted against the rate limit.
+Previously they were not, so sustained use could exceed Space-Track's limits and trigger server-side rate limit errors.

src/spacetrack/base.py

@@ -754,25 +754,37 @@ def _parse_types(data, predicates):
 
         return data
 
-    def _ratelimited_send_generator(self, request, *, stream=False):
-        """Send a request, handling rate limiting."""
-        minute_limit = self._per_minute_throttle.check(self._per_minute_key, 1)
-        hour_limit = self._per_hour_throttle.check(self._per_hour_key, 1)
-
-        sleep_time = 0
+    def _ratelimit_check_generator(self):
+        """Wait until no throttle is limited, then charge all of them.
 
-        if minute_limit.limited:
-            sleep_time = minute_limit.retry_after.total_seconds()
+        Peeking first means that waiting on one quota does not consume the
+        others.
+        """
+        throttles = [
+            (self._per_minute_throttle, self._per_minute_key),
+            (self._per_hour_throttle, self._per_hour_key),
+        ]
+        if self._additional_throttle is not None:
+            throttles.append((self._additional_throttle, self._additional_key))
 
-        if hour_limit.limited:
-            sleep_time = max(sleep_time, hour_limit.retry_after.total_seconds())
+        while True:
+            limits = [throttle.peek(key) for throttle, key in throttles]
+            if not any(limit.limited for limit in limits):
+                # Another process sharing the store may have used the quota
+                # since the peek, in which case wait and try again.
+                limits = [throttle.check(key, 1) for throttle, key in throttles]
+
+            limited = [limit for limit in limits if limit.limited]
+            if not limited:
+                return
 
-        if self._additional_throttle is not None:
-            additional_limit = self._additional_throttle.check(self._additional_key, 1)
-            sleep_time = max(sleep_time, additional_limit.retry_after.total_seconds())
+            yield RateLimitWait(
+                max(limit.retry_after.total_seconds() for limit in limited)
+            )
 
-        if sleep_time > 0:
-            yield RateLimitWait(sleep_time)
+    def _ratelimited_send_generator(self, request, *, stream=False):
+        """Send a request, handling rate limiting."""
+        yield from self._ratelimit_check_generator()
 
         req_event = NormalRequest(request, stream=stream, follow_redirects=True)
         resp = yield req_event
@@ -794,6 +806,7 @@ def _ratelimited_send_generator(self, request, *, stream=False):
                 yield RateLimitWait(
                     self._per_minute_throttle.rate.period.total_seconds()
                 )
+                yield from self._ratelimit_check_generator()
                 resp = yield req_event
 
         return resp
@@ -881,7 +894,7 @@ def _download_predicate_data_generator(self, class_, controller):
         req = self.client.build_request("GET", url)
         logger.debug(req.url)
 
-        resp = yield NormalRequest(req)
+        resp = yield from self._ratelimited_send_generator(req)
 
         _raise_for_status(resp)
 

tests/test_spacetrack.py

@@ -252,6 +252,86 @@ def test_ratelimit_callback_error(client, httpx2_mock, mock_auth, mock_gp_predic
         client.gp()
 
 
+def test_ratelimit_rechecked_after_wait(
+    client, httpx2_mock, mock_auth, mock_gp_predicates
+):
+    # With one request per window, every query after the model definition
+    # request has to wait, which is only the case if requests sent after a
+    # wait are recorded.
+    client._per_minute_throttle.rate = Quota(
+        period=dt.timedelta(milliseconds=50), count=1
+    )
+
+    url = api_url("basicspacedata/query/class/gp")
+    httpx2_mock.add_response(method="GET", url=url, json={"a": 1}, is_reusable=True)
+
+    waits = []
+    client.callback = waits.append
+
+    for _ in range(3):
+        assert client.gp() == {"a": 1}
+
+    assert len(waits) >= 3
+
+    # Waiting on the per-minute quota must not have charged the per-hour one
+    # more than once per request sent (the model definition plus 3 queries).
+    hour_limit = client._per_hour_throttle.peek(client._per_hour_key)
+    assert hour_limit.limit - hour_limit.remaining == 4
+
+
+def test_ratelimit_per_hour(client, httpx2_mock, mock_auth, mock_gp_predicates):
+    # Trip the per-hour throttle instead of the per-minute one, with a short
+    # period so that the real wait is brief.
+    client._per_hour_throttle.rate = Quota(
+        period=dt.timedelta(milliseconds=50), count=2
+    )
+
+    url = api_url("basicspacedata/query/class/gp")
+    httpx2_mock.add_response(method="GET", url=url, json={"a": 1}, is_reusable=True)
+
+    waits = []
+    client.callback = waits.append
+
+    for _ in range(3):
+        assert client.gp() == {"a": 1}
+
+    assert waits
+
+
+def test_additional_rate_limit(httpx2_mock, mock_auth, mock_gp_predicates):
+    url = api_url("basicspacedata/query/class/gp")
+    httpx2_mock.add_response(method="GET", url=url, json={"a": 1}, is_reusable=True)
+
+    with SpaceTrackClient(
+        "identity",
+        "password",
+        additional_rate_limit=Quota(period=dt.timedelta(milliseconds=50), count=1),
+    ) as client:
+        waits = []
+        client.callback = waits.append
+
+        assert client.gp() == {"a": 1}
+        assert client.gp() == {"a": 1}
+
+    assert waits
+
+
+def test_modeldef_ratelimit_error(client, httpx2_mock, mock_auth):
+    # Change ratelimiter period to speed up test
+    client._per_minute_throttle.rate = Quota(
+        period=dt.timedelta(milliseconds=50), count=30
+    )
+
+    url = api_url("basicspacedata/modeldef/class/gp")
+    httpx2_mock.add_response(
+        method="GET", url=url, status_code=500, text="violated your query rate limit"
+    )
+    httpx2_mock.add_response(method="GET", url=url, json={"data": []})
+
+    assert client.get_predicates("gp") == []
+    assert len(httpx2_mock.get_requests(method="GET", url=url)) == 2
+
+
 def test_predicate_parse_modeldef(client):
     predicates_data = [
         {