Skip to content

Commit

Permalink
Change: Improve rate limit algorithm for requests against the NVT API
Browse files Browse the repository at this point in the history
The original algorithm was very simple by just applying a pause after a
specific number of allowed requests. When doing additional costly tasks
after each requests this algorithm results in unnecessary delays and the
sleep might me completely obsolete. Therefore calculate the time since
the last check for a sleep. If this time delta is within the rate limit
we wait for the required time before doing the next request.
  • Loading branch information
bjoernricks committed Nov 3, 2023
1 parent 124bd52 commit c871322
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions pontos/nvd/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import asyncio
import time
from abc import ABC
from datetime import datetime, timezone
from types import TracebackType
Expand Down Expand Up @@ -76,10 +77,6 @@ def convert_camel_case(dct: Dict[str, Any]) -> Dict[str, Any]:
return converted


async def sleep() -> None:
await asyncio.sleep(SLEEP_TIMEOUT)


class NVDApi(ABC):
"""
Abstract base class for querying the NIST NVD API.
Expand Down Expand Up @@ -120,6 +117,7 @@ def __init__(
self._rate_limit = None

self._request_count = 0
self._last_sleep = time.monotonic()

def _request_headers(self) -> Headers:
"""
Expand All @@ -141,7 +139,13 @@ async def _consider_rate_limit(self) -> None:

self._request_count += 1
if self._request_count > self._rate_limit:
await sleep()
time_since_last_sleep = time.monotonic() - self._last_sleep

if time_since_last_sleep < SLEEP_TIMEOUT:
time_to_sleep = SLEEP_TIMEOUT - time_since_last_sleep
await asyncio.sleep(time_to_sleep)

self._last_sleep = time.monotonic()
self._request_count = 0

async def _get(
Expand Down

0 comments on commit c871322

Please sign in to comment.