diff --git a/src/utils.py b/src/utils.py index d703da9..8a9a09d 100644 --- a/src/utils.py +++ b/src/utils.py @@ -1,4 +1,5 @@ import json +import asyncio def convert_to_bytes(size: int | str) -> int: try: @@ -23,11 +24,19 @@ def convert_to_bytes(size: int | str) -> int: async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes: size_left = size - while size_left > buffer_size: - size_left -= buffer_size - yield b'\0' * buffer_size - else: - yield b'\0' * size_left + # https://github.com/tiangolo/fastapi/issues/5183 + # https://github.com/encode/starlette/discussions/1776#discussioncomment-3207518 + + try: + while size_left > buffer_size: + size_left -= buffer_size + yield b'\0' * buffer_size + await asyncio.sleep(0) + else: + yield b'\0' * size_left + await asyncio.sleep(0) + except asyncio.CancelledError: + raise GeneratorExit def check_policies(ip: str) -> None: