refactor: apply suggestions

jina/clients/base/http.py

@@ -108,7 +108,9 @@ async def _get_results(
                 HTTPClientlet(url=url, logger=self.logger, **kwargs)
             )
 
-            def _request_handler(request: 'Request') -> 'Tuple[asyncio.Future, None]':
+            def _request_handler(
+                request: 'Request',
+            ) -> 'Tuple[asyncio.Future, Optional[asyncio.Future]]':
                 """
                 For HTTP Client, for each request in the iterator, we `send_message` using
                 http POST request and add it to the list of tasks which is awaited and yielded.

jina/clients/base/websocket.py

@@ -139,7 +139,9 @@ def _handle_end_of_iter():
                 """Send End of iteration signal to the Gateway"""
                 asyncio.create_task(iolet.send_eoi())
 
-            def _request_handler(request: 'Request') -> 'Tuple[asyncio.Future, None]':
+            def _request_handler(
+                request: 'Request',
+            ) -> 'Tuple[asyncio.Future, Optional[asyncio.Future]]':
                 """
                 For each request in the iterator, we send the `Message` using `iolet.send_message()`.
                 For websocket requests from client, for each request in the iterator, we send the request in `bytes`

jina/serve/stream/helper.py

@@ -60,9 +60,6 @@ def __aiter__(self):
         return self
 
     async def __anext__(self):
-        if self._prefetch > 0:
-            while self._request_counter.count > self._prefetch:
-                await asyncio.sleep(0)
         if isinstance(self.iterator, Iterator):
 
             """
@@ -84,5 +81,7 @@ async def __anext__(self):
         elif isinstance(self.iterator, AsyncIterator):
             # we assume that `AsyncIterator` doesn't block the event loop
             request = await self.iterator.__anext__()
-
+        if self._prefetch > 0:
+            while self._request_counter.count >= self._prefetch:
+                await asyncio.sleep(0)
         return request