Skip to content

Commit

Permalink
format and types
Browse files Browse the repository at this point in the history
  • Loading branch information
riccardobl committed Nov 19, 2024
1 parent 1d8dc88 commit 036285b
Showing 1 changed file with 18 additions and 19 deletions.
37 changes: 18 additions & 19 deletions nwcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,11 @@
from loguru import logger


class RateLimit:
backoff: int = 0
last_attempt_time: int = 0


class MainSubscription:
def __init__(self):
self.requests_sub_id: Optional[str] = None
Expand Down Expand Up @@ -86,8 +91,8 @@ def __init__(self, private_key: Optional[str] = None, relay: Optional[str] = Non

# Subscription
self.sub = None
self.rate_limit = {}
self.rate_limit: Dict[str, RateLimit] = {}

# websocket connection
self.ws = None

Expand Down Expand Up @@ -200,28 +205,22 @@ async def _wait_for_connection(self):
logger.debug("Waiting for connection...")
await asyncio.sleep(1)

async def _ratelimit(self, unit, max_sleep_time = 120):
rate_limit = self.rate_limit.get(unit)
if not rate_limit:
self.rate_limit[unit] = rate_limit = {
"backoff": 0,
"last_attempt_time": 0
}
async def _ratelimit(self, unit: str, max_sleep_time: int = 120) -> None:
limit: Optional[RateLimit] = self.rate_limit.get(unit)
if not limit:
self.rate_limit[unit] = limit = RateLimit()

if time.time() - rate_limit["last_attempt_time"] > max_sleep_time:
if time.time() - limit.last_attempt_time > max_sleep_time:
# reset backoff if action lasted more than max_sleep_time
rate_limit["backoff"] = 0
limit.backoff = 0
else:
# increase backoff
rate_limit["backoff"] = (
min(rate_limit["backoff"] * 2, max_sleep_time)
if rate_limit["backoff"] > 0
else 1
limit.backoff = (
min(limit.backoff * 2, max_sleep_time) if limit.backoff > 0 else 1
)
logger.debug(
"Sleeping for " + str(rate_limit["backoff"]) + " seconds before " + unit)
await asyncio.sleep( rate_limit["backoff"])
rate_limit["last_attempt_time"] = time.time()
logger.debug("Sleeping for " + str(limit.backoff) + " seconds before " + unit)
await asyncio.sleep(limit.backoff)
limit.last_attempt_time = int(time.time())

async def _subscribe(self):
"""
Expand Down

0 comments on commit 036285b

Please sign in to comment.