From 9e83da6b1aa1d185a03d10978a02e10d668fc569 Mon Sep 17 00:00:00 2001
From: Kristian Krsnik <git@krsnik.at>
Date: Fri, 12 Apr 2024 18:02:14 +0200
Subject: [PATCH] added await

---
 src/utils.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

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: