13388974656

Committed 18 Feb 2025 10:56AM UTC coverage: 93.946% (-0.004%) from 93.95%

Build # 13388974656

Build Type

Pull #1518

github

Committed by

web-flow

Commit Message

Merge 4ce6371ba into 8b28bc770

Pull Request Pull Request #1518: Add custom rate limit for aiohtpp

Run Details

58 of 62 new or added lines in 3 files covered. (93.55%)

8 existing lines in 3 files now uncovered.

9543 of 10158 relevant lines covered (93.95%)

2.82 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.12

/safe_eth/eth/clients/rate_limiter.py

import asyncio
from functools import cache
from logging import getLogger

import aiohttp

logger = getLogger(__name__)


class RateLimiter:
    """
    Class to limit the number of requests per second
    """

    def __init__(self, client, rate):
        self.client = client
        self.rate = rate
        self.available_conns = rate  # Initialize available conns
        self._waiters = []  # List of tasks that are waiting for a connection
        self.loop = asyncio.get_event_loop()
        self._schedule_next_release_connections()  # Schedule first release connections

    async def get(self, *args, **kwargs):
        await self._wait_for_available_conn()
        return self.client.get(*args, **kwargs)

    async def post(self, *args, **kwargs):
        await self._wait_for_available_conn()
        return self.client.post(*args, **kwargs)

    def _wakeup_waiters(self):
        """
        Unblock tasks waiting for connections
        """
        while self.available_conns > 0 and self._waiters:
            future = self._waiters.pop(0)
            future.set_result(None)  # Release await
            self.available_conns -= 1

    async def _wait_for_available_conn(self):
        if self.available_conns < 1:
            future = asyncio.Future()
            self._waiters.append(future)
            await future
        else:
            self.available_conns -= 1

    def _release_available_conns(self):
        """
        Release new connections
        """
        self.available_conns += self.rate - self.available_conns
        self._wakeup_waiters()
        self._schedule_next_release_connections()

    def _schedule_next_release_connections(self):
        """
        Schedule next release connections
        """
        self.loop.call_later(1, self._release_available_conns)


@cache
def get_client_rate_limited(host: str, rate: int) -> "RateLimiter":
    """
    Get a rate limited client by host
    Host parameter is just being used to store in cache different instance by host

    :param host:
    :param rate: number of requests allowed per second
    """
    logger.info(f"Initializing rate limiter for {host} by {rate}/s")
    async_session = aiohttp.ClientSession()
    return RateLimiter(async_session, rate)

1	import asyncio	3✔
2	from functools import cache	3✔
3	from logging import getLogger	3✔
4
5	import aiohttp	3✔
6
7	logger = getLogger(__name__)	3✔
8
9
10	class RateLimiter:	3✔
11	"""
12	Class to limit the number of requests per second
13	"""
14
15	def __init__(self, client, rate):	3✔
16	self.client = client	3✔
17	self.rate = rate	3✔
18	self.available_conns = rate # Initialize available conns	3✔
19	self._waiters = [] # List of tasks that are waiting for a connection	3✔
20	self.loop = asyncio.get_event_loop()	3✔
21	self._schedule_next_release_connections() # Schedule first release connections	3✔
22
23	async def get(self, args, *kwargs):	3✔
24	await self._wait_for_available_conn()	3✔
25	return self.client.get(args, *kwargs)	3✔
26
27	async def post(self, args, *kwargs):	3✔
NEW 28	await self._wait_for_available_conn()	×
NEW 29	return self.client.post(args, *kwargs)	×
30
31	def _wakeup_waiters(self):	3✔
32	"""
33	Unblock tasks waiting for connections
34	"""
35	while self.available_conns > 0 and self._waiters:	3✔
36	future = self._waiters.pop(0)	3✔
37	future.set_result(None) # Release await	3✔
38	self.available_conns -= 1	3✔
39
40	async def _wait_for_available_conn(self):	3✔
41	if self.available_conns < 1:	3✔
42	future = asyncio.Future()	3✔
43	self._waiters.append(future)	3✔
44	await future	3✔
45	else:
46	self.available_conns -= 1	3✔
47
48	def _release_available_conns(self):	3✔
49	"""
50	Release new connections
51	"""
52	self.available_conns += self.rate - self.available_conns	3✔
53	self._wakeup_waiters()	3✔
54	self._schedule_next_release_connections()	3✔
55
56	def _schedule_next_release_connections(self):	3✔
57	"""
58	Schedule next release connections
59	"""
60	self.loop.call_later(1, self._release_available_conns)	3✔
61
62
63	@cache	3✔
64	def get_client_rate_limited(host: str, rate: int) -> "RateLimiter":	3✔
65	"""
66	Get a rate limited client by host
67	Host parameter is just being used to store in cache different instance by host
68
69	:param host:
70	:param rate: number of requests allowed per second
71	"""
72	logger.info(f"Initializing rate limiter for {host} by {rate}/s")	3✔
73	async_session = aiohttp.ClientSession()	3✔
74	return RateLimiter(async_session, rate)	3✔

safe-global / safe-eth-py / 13388974656

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous