Skip to content

Commit

Permalink
Add default instead of calculating window size (#147)
Browse files Browse the repository at this point in the history
(cherry picked from commit praw-dev/prawcore@e2a22e0)
  • Loading branch information
Watchful1 authored and LilSpazJoekp committed Nov 26, 2023
1 parent 44a7f7e commit e69abe8
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 30 deletions.
1 change: 1 addition & 0 deletions asyncprawcore/const.py
Expand Up @@ -7,3 +7,4 @@
AUTHORIZATION_PATH = "/api/v1/authorize"
REVOKE_TOKEN_PATH = "/api/v1/revoke_token"
TIMEOUT = float(os.environ.get("prawcore_timeout", 16))
WINDOW_SIZE = 600
9 changes: 2 additions & 7 deletions asyncprawcore/rate_limit.py
Expand Up @@ -17,13 +17,13 @@ class RateLimiter(object):
"""

def __init__(self) -> None:
def __init__(self, *, window_size: int) -> None:
"""Create an instance of the RateLimit class."""
self.remaining: Optional[float] = None
self.next_request_timestamp: Optional[float] = None
self.reset_timestamp: Optional[float] = None
self.used: Optional[int] = None
self.window_size: Optional[float] = None
self.window_size: int = window_size

async def call(
self,
Expand Down Expand Up @@ -81,11 +81,6 @@ def update(self, response_headers: Mapping[str, str]) -> None:
self.used = int(response_headers["x-ratelimit-used"])
self.reset_timestamp = now + seconds_to_reset

if self.window_size is None:
self.window_size = seconds_to_reset + self.used
elif self.window_size < seconds_to_reset:
self.window_size = seconds_to_reset

if self.remaining <= 0:
self.next_request_timestamp = self.reset_timestamp
return
Expand Down
25 changes: 19 additions & 6 deletions asyncprawcore/sessions.py
Expand Up @@ -2,6 +2,7 @@
import asyncio
import logging
import random
import time
from abc import ABC, abstractmethod
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
Expand All @@ -11,7 +12,7 @@

from .auth import BaseAuthorizer
from .codes import codes
from .const import TIMEOUT
from .const import TIMEOUT, WINDOW_SIZE
from .exceptions import (
BadJSON,
BadRequest,
Expand Down Expand Up @@ -134,7 +135,7 @@ def _log_request(
params: Dict[str, int],
url: str,
):
log.debug(f"Fetching: {method} {url}")
log.debug(f"Fetching: {method} {url} at {time.time()}")
log.debug(f"Data: {data}")
log.debug(f"Params: {params}")

Expand All @@ -148,16 +149,21 @@ def _preprocess_dict(data: Dict[str, Any]) -> Dict[str, str]:
new_data[key] = str(value) if not isinstance(value, str) else value
return new_data

def __init__(self, authorizer: Optional["Authorizer"]) -> None:
def __init__(
self,
authorizer: Optional[BaseAuthorizer],
window_size: int = WINDOW_SIZE,
) -> None:
"""Prepare the connection to Reddit's API.
:param authorizer: An instance of :class:`.Authorizer`.
:param window_size: The size of the rate limit reset window in seconds.
"""
if not isinstance(authorizer, BaseAuthorizer):
raise InvalidInvocation(f"invalid Authorizer: {authorizer}")
self._authorizer = authorizer
self._rate_limiter = RateLimiter()
self._rate_limiter = RateLimiter(window_size=window_size)
self._retry_strategy_class = FiniteRetryStrategy

async def __aenter__(self) -> "Session":
Expand Down Expand Up @@ -221,6 +227,9 @@ async def _make_request(
log.debug(
f"Response: {response.status}"
f" ({response.headers.get('content-length')} bytes)"
f" (rst-{response.headers.get('x-ratelimit-reset')}:"
f"rem-{response.headers.get('x-ratelimit-remaining')}:"
f"used-{response.headers.get('x-ratelimit-used')} ratelimit) at {time.time()}"
)
return response, None
except RequestException as exception:
Expand Down Expand Up @@ -410,10 +419,14 @@ async def request(
)


def session(authorizer: "Authorizer" = None) -> Session:
def session(
authorizer: "Authorizer" = None,
window_size: int = WINDOW_SIZE,
) -> Session:
"""Return a :class:`.Session` instance.
:param authorizer: An instance of :class:`.Authorizer`.
:param window_size: The size of the rate limit reset window in seconds.
"""
return Session(authorizer=authorizer)
return Session(authorizer=authorizer, window_size=window_size)
14 changes: 9 additions & 5 deletions tests/integration/test_sessions.py
Expand Up @@ -35,11 +35,15 @@ async def test_request__accepted(
caplog.set_level(logging.DEBUG)
session = asyncprawcore.Session(script_authorizer)
await session.request("POST", "api/read_all_messages")
assert (
"asyncprawcore",
logging.DEBUG,
"Response: 202 (2 bytes)",
) in caplog.record_tuples
found_message = False
for package, level, message in caplog.record_tuples:
if (
package == "asyncprawcore"
and level == logging.DEBUG
and "Response: 202 (2 bytes)" in message
):
found_message = True
assert found_message, f"'Response: 202 (2 bytes)' in {caplog.record_tuples}"

async def test_request__bad_gateway(
self, readonly_authorizer: asyncprawcore.ReadOnlyAuthorizer
Expand Down
17 changes: 5 additions & 12 deletions tests/unit/test_rate_limit.py
Expand Up @@ -12,7 +12,7 @@
class TestRateLimiter(UnitTest):
@pytest.fixture
def rate_limiter(self):
rate_limiter = RateLimiter()
rate_limiter = RateLimiter(window_size=600)
rate_limiter.next_request_timestamp = 100
return rate_limiter

Expand Down Expand Up @@ -68,29 +68,22 @@ def test_update__compute_delay_with_no_previous_info(self, mock_time, rate_limit
@patch("time.time")
def test_update__compute_delay_with_single_client(self, mock_time, rate_limiter):
rate_limiter.remaining = 61
rate_limiter.window_size = 150
mock_time.return_value = 100
rate_limiter.update(self._headers(50, 100, 60))
assert rate_limiter.remaining == 50
assert rate_limiter.used == 100
assert rate_limiter.next_request_timestamp == 106.66666666666667
assert rate_limiter.next_request_timestamp == 110

@patch("time.time")
def test_update__compute_delay_with_six_clients(self, mock_time, rate_limiter):
rate_limiter.remaining = 66
rate_limiter.window_size = 180
mock_time.return_value = 100
rate_limiter.update(self._headers(60, 100, 72))
assert rate_limiter.remaining == 60
assert rate_limiter.used == 100
assert rate_limiter.next_request_timestamp == 107.5

@patch("time.time")
def test_update__compute_delay_with_window_set(self, mock_time, rate_limiter):
rate_limiter.window_size = 550
mock_time.return_value = 100
rate_limiter.update(self._headers(599, 1, 600))
assert rate_limiter.remaining == 599
assert rate_limiter.used == 1
assert rate_limiter.next_request_timestamp == 101.0
assert rate_limiter.next_request_timestamp == 104.5

@patch("time.time")
def test_update__delay_full_time_with_negative_remaining(
Expand Down

0 comments on commit e69abe8

Please sign in to comment.