def test_is_rate_limited_when_rate_limit_not_expired_only_returns_False(
self, remaining):
with rate_limits.WindowedBurstRateLimiter(__name__, 403, 27) as rl:
now = 420
rl.reset_at = now + 69
rl.remaining = remaining
assert rl.is_rate_limited(now) is (remaining <= 0)
async def test_throttle_consumes_queue(self, event_loop):
with rate_limits.WindowedBurstRateLimiter(__name__, 0.01, 1) as rl:
rl.queue = [event_loop.create_future() for _ in range(15)]
old_queue = list(rl.queue)
await rl.throttle()
assert len(rl.queue) == 0
for i, future in enumerate(old_queue):
assert future.done(), f"future {i} was incomplete!"
def test_is_rate_limited_when_rate_limit_expired_resets_self(self):
with rate_limits.WindowedBurstRateLimiter(__name__, 403, 27) as rl:
now = 180
rl.reset_at = 80
rl.remaining = 4
assert not rl.is_rate_limited(now)
assert rl.reset_at == now + 403
assert rl.remaining == 27
async def test_throttle_when_limited_sleeps_then_bursts_repeatedly(
self, event_loop):
window = 5
sleep_count = 0
futures = [event_loop.create_future() for _ in range(20)]
async def mock_sleep(t):
nonlocal sleep_count
for i, future in enumerate(futures):
if i >= (window * sleep_count):
assert not future.done(
), f"future {i} was complete, expected it to be incomplete!"
else:
assert future.done(
), f"future {i} was incomplete, expected it to be completed!"
sleep_count += 1
stack = contextlib.ExitStack()
rl = stack.enter_context(
rate_limits.WindowedBurstRateLimiter(__name__, 0, window))
stack.enter_context(mock.patch.object(asyncio, "sleep",
new=mock_sleep))
with stack:
rl.queue = list(futures)
rl.reset_at = time.perf_counter()
await rl.throttle()
# die if we take too long...
await asyncio.wait(futures, timeout=3)
assert sleep_count == 4
assert len(rl.queue) == 0
for i, future in enumerate(futures):
assert future.done(), f"future {i} was incomplete!"
async def test_throttle_resets_throttle_task(self, event_loop):
with rate_limits.WindowedBurstRateLimiter(__name__, 0.01, 1) as rl:
rl.queue = [event_loop.create_future() for _ in range(15)]
rl.throttle_task = None
await rl.throttle()
assert rl.throttle_task is None
async def _run_test_throttle_logic_on_loop(event_loop):
limit = 2
period = 1.5
total_requests = int(period * limit * 2)
max_distance_within_window = 0.05
completion_times = []
logger = logging.getLogger(__name__)
def create_task(i):
logger.info("making task %s", i)
future = event_loop.create_future()
future.add_done_callback(
lambda _: completion_times.append(time.perf_counter()))
return future
with rate_limits.WindowedBurstRateLimiter(__name__, period,
limit) as rl:
futures = [create_task(i) for i in range(total_requests)]
rl.queue = list(futures)
rl.reset_at = time.perf_counter()
logger.info("throttling back")
await rl.throttle()
# die if we take too long...
logger.info("waiting for stuff to finish")
await asyncio.wait(futures, timeout=period * limit + period)
assert (
len(completion_times) == total_requests
), f"expected {total_requests} completions but got {len(completion_times)}"
# E203 - Whitespace before ":". Black reformats it
windows = [
completion_times[i:i + limit]
for i in range(0, total_requests, limit)
] # noqa: E203
for i, window in enumerate(windows):
logger.info("window %s %s", i, window)
mode = statistics.mode(window)
for j, element in enumerate(window):
assert math.isclose(
element, mode, abs_tol=max_distance_within_window
), (f"not close! windows[{i}][{j}], future {i * len(window) + j}, "
f"val {element}, mode {mode}, max diff {max_distance_within_window}"
)
assert len(windows) >= 3, "not enough windows to sample correctly"
assert len(
windows[0]
) > 1, "not enough datapoints per window to sample correctly"
for i in range(1, len(windows)):
previous_last = windows[i - 1][-1]
next_first = windows[i][0]
logger.info(
"intra-window index=%s value=%s versus index=%s value=%s",
i - 1, previous_last, i, next_first)
assert math.isclose(
next_first - previous_last,
period,
abs_tol=max_distance_within_window
), (f"distance between windows is not acceptable! {i - 1}={previous_last} {i}={next_first}, "
f"max diff = {max_distance_within_window}")
def __init__(
self,
*,
compression: typing.Optional[str] = shard.GatewayCompression.
PAYLOAD_ZLIB_STREAM,
initial_activity: typing.Optional[presences.Activity] = None,
initial_idle_since: typing.Optional[datetime.datetime] = None,
initial_is_afk: bool = False,
initial_status: presences.Status = presences.Status.ONLINE,
intents: intents_.Intents,
large_threshold: int = 250,
shard_id: int = 0,
shard_count: int = 1,
event_consumer: typing.Callable[
[shard.GatewayShard, str, data_binding.JSONObject], None],
http_settings: config.HTTPSettings,
proxy_settings: config.ProxySettings,
data_format: str = shard.GatewayDataFormat.JSON,
token: str,
url: str,
) -> None:
if data_format != shard.GatewayDataFormat.JSON:
raise NotImplementedError(
f"Unsupported gateway data format: {data_format}")
query = {"v": _VERSION, "encoding": str(data_format)}
if compression is not None:
if compression == shard.GatewayCompression.PAYLOAD_ZLIB_STREAM:
query["compress"] = "zlib-stream"
else:
raise NotImplementedError(
f"Unsupported compression format {compression}")
scheme, netloc, path, params, _, _ = urllib.parse.urlparse(
url, allow_fragments=True)
new_query = urllib.parse.urlencode(query)
self._activity = initial_activity
self._closing = asyncio.Event()
self._closed = asyncio.Event()
self._chunking_rate_limit = rate_limits.WindowedBurstRateLimiter(
f"shard {shard_id} chunking rate limit",
*_CHUNKING_RATELIMIT,
)
self._event_consumer = event_consumer
self._handshake_completed = asyncio.Event()
self._heartbeat_latency = float("nan")
self._http_settings = http_settings
self._idle_since = initial_idle_since
self._intents = intents
self._is_afk = initial_is_afk
self._large_threshold = large_threshold
self._last_heartbeat_ack_received = float("nan")
self._last_heartbeat_sent = float("nan")
self._logger = logging.getLogger(f"hikari.gateway.{shard_id}")
self._proxy_settings = proxy_settings
self._run_task: typing.Optional[asyncio.Task[None]] = None
self._seq: typing.Optional[int] = None
self._session_id: typing.Optional[str] = None
self._shard_count = shard_count
self._shard_id = shard_id
self._status = initial_status
self._token = token
self._total_rate_limit = rate_limits.WindowedBurstRateLimiter(
f"shard {shard_id} total rate limit",
*_TOTAL_RATELIMIT,
)
self._url = urllib.parse.urlunparse(
(scheme, netloc, path, params, new_query, ""))
self._user_id: typing.Optional[snowflakes.Snowflake] = None
self._ws: typing.Optional[_GatewayTransport] = None
