Skip to content

Commit

Permalink
add rate limiting mechanism
Browse files Browse the repository at this point in the history
  • Loading branch information
schelv committed Jan 11, 2024
1 parent 4406cc0 commit 2583440
Showing 1 changed file with 55 additions and 12 deletions.
67 changes: 55 additions & 12 deletions githubkit/core.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from asyncio import Event, BoundedSemaphore, sleep
from types import TracebackType
from contextvars import ContextVar
from datetime import datetime, timezone, timedelta
Expand Down Expand Up @@ -33,6 +34,7 @@
QueryParamTypes,
)
from .exception import (
RateLimitExceeded,
RequestError,
RequestFailed,
RequestTimeout,
Expand Down Expand Up @@ -135,6 +137,7 @@ def __init__(
follow_redirects: bool = True,
timeout: Optional[Union[float, httpx.Timeout]] = None,
http_cache: bool = True,
max_nr_concurrent_requests: int = 100,
):
auth = auth or UnauthAuthStrategy() # type: ignore
self.auth: A = ( # type: ignore
Expand All @@ -151,6 +154,11 @@ def __init__(
http_cache,
)

self.rate_limit_free = Event()
self.rate_limit_free.set()

self.concurrent_request_semaphore = BoundedSemaphore(max_nr_concurrent_requests)

self.__sync_client: ContextVar[Optional[httpx.Client]] = ContextVar(
"sync_client", default=None
)
Expand Down Expand Up @@ -435,16 +443,51 @@ async def arequest(
cookies: Optional[CookieTypes] = None,
response_model: Type[T] = Any,
error_models: Optional[Dict[str, type]] = None,
retry_attempt_nr: int = 0,
) -> Response[T]:
raw_resp = await self._arequest(
method,
url,
params=params,
content=content,
data=data,
files=files,
json=json,
headers=headers,
cookies=cookies,
)
return self._check(raw_resp, response_model, error_models)
# limit the number of concurrent requests.
async with self.concurrent_request_semaphore:
# only continue if no rate limit active.
await self.rate_limit_free.wait()
try:
raw_resp = await self._arequest(
method,
url,
params=params,
content=content,
data=data,
files=files,
json=json,
headers=headers,
cookies=cookies,
)
return self._check(raw_resp, response_model, error_models)
except RateLimitExceeded as error:

Check warning on line 465 in githubkit/core.py

View check run for this annotation

Codecov / codecov/patch

githubkit/core.py#L465

Added line #L465 was not covered by tests
# block all new requests and try again.
if retry_attempt_nr > 3:
raise error

Check warning on line 468 in githubkit/core.py

View check run for this annotation

Codecov / codecov/patch

githubkit/core.py#L467-L468

Added lines #L467 - L468 were not covered by tests

start_time = datetime.now()
rate_limit_duration = error.retry_after

Check warning on line 471 in githubkit/core.py

View check run for this annotation

Codecov / codecov/patch

githubkit/core.py#L470-L471

Added lines #L470 - L471 were not covered by tests

# print(f"retry request for {url} after {rate_limit_duration} seconds.")
print(f"Started rate limit timer at {start_time} for {rate_limit_duration} seconds.")
self.rate_limit_free.clear()
await sleep(error.retry_after.seconds)
self.rate_limit_free.set()
print(f"rate limit that started at {start_time} stopped at {datetime.now()}")

Check warning on line 478 in githubkit/core.py

View check run for this annotation

Codecov / codecov/patch

githubkit/core.py#L474-L478

Added lines #L474 - L478 were not covered by tests

return await self.arequest(

Check warning on line 480 in githubkit/core.py

View check run for this annotation

Codecov / codecov/patch

githubkit/core.py#L480

Added line #L480 was not covered by tests
method,
url,
params=params,
content=content,
data=data,
files=files,
json=json,
headers=headers,
cookies=cookies,
response_model=response_model,
error_models=error_models,
retry_attempt_nr=retry_attempt_nr + 1,
)

0 comments on commit 2583440

Please sign in to comment.