Repository: jina-ai/serve
Pull request: feat(client): add prefetch to http client

Problem statement:
Http client not streaming requests
The HttpClient is not using streaming correctly together with input generators.

See this test code

```
import time

from jina import Flow, Document, Executor, requests

def slow_gen():
    for j in range(10):
        yield Document()
        time.sleep(1)
        print(j)


class MyExecutor(Executor):

    @requests
    def foo(self, **kwargs):
        print('hi')


for p in ['grpc', 'websocket', 'http']:
    with Flow(prefetch=1, protocol=p).add(uses=MyExecutor) as f:
        f.post('/', slow_gen, request_size=1, on_done=lambda x: print('done'))
```

Hints from issue discussion:

Commit message: feat(client): add prefetch to http client

Patch to review:
diff --git a/jina/clients/base/http.py b/jina/clients/base/http.py
index 018ba4ff32308..0cff3ef93e5f3 100644
--- a/jina/clients/base/http.py
+++ b/jina/clients/base/http.py
@@ -1,6 +1,5 @@
-import asyncio
-from contextlib import nullcontext
-from typing import Callable, Optional
+from contextlib import nullcontext, AsyncExitStack
+from typing import Callable, Dict, Optional, TYPE_CHECKING
 
 from ..base import BaseClient, InputType
 from ..helper import callback_exec
@@ -8,19 +7,70 @@
 from ...importer import ImportExtensions
 from ...logging.profile import ProgressBar
 from ...types.request import Request
+from ...peapods.runtimes.gateway.prefetch import HTTPClientPrefetchCaller
+
+
+if TYPE_CHECKING:
+    from ...logging.logger import JinaLogger
+
+
+class HTTPClientlet:
+    """HTTP Client to be used with prefetcher"""
+
+    def __init__(self, url: str, logger: 'JinaLogger') -> None:
+        """HTTP Client to be used with prefetcher
+
+        :param url: url to send POST request to
+        :param logger: logger
+        """
+        self.url = url
+        self.msg_recv = 0
+        self.msg_sent = 0
+        self.logger = logger
+        self.session = None
+
+    async def send_message(self, request: Dict):
+        """Sends a POST request to the server
+
+        :param request: request as dict
+        :return: send post message
+        """
+        return await self.session.post(url=self.url, json=request).__aenter__()
+
+    async def __aenter__(self):
+        """enter async context
+
+        :return: start self
+        """
+        return await self.start()
+
+    async def start(self):
+        """Create ClientSession and enter context
+
+        :return: self
+        """
+        import aiohttp
+
+        self.session = aiohttp.ClientSession()
+        await self.session.__aenter__()
+        self.logger.debug('aiohttp session started to send requests to gateway')
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close(exc_type, exc_val, exc_tb)
+
+    async def close(self, *args, **kwargs):
+        """Close ClientSession
+
+        :param args: positional args
+        :param kwargs: keyword args"""
+        await self.session.__aexit__(*args, **kwargs)
+        self.logger.debug('aiohttp session successfully closed')
 
 
 class HTTPBaseClient(BaseClient):
     """A MixIn for HTTP Client."""
 
-    async def _get_http_response(self, session, dest_url, req_dict):
-        async with session.post(
-            dest_url,
-            json=req_dict,
-        ) as response:
-            resp_str = await response.json()
-            return response.status, resp_str
-
     async def _get_results(
         self,
         inputs: InputType,
@@ -30,19 +80,6 @@ async def _get_results(
         **kwargs,
     ):
         """
-        :meth:`send_requests()`
-            Traverses through the request iterator
-            Sends each request & awaits :meth:`websocket.send()`
-            Sends & awaits `byte(True)` to acknowledge request iterator is empty
-        Traversal logic:
-            Starts an independent task :meth:`send_requests()`
-            Awaits on each response from :meth:`websocket.recv()` (done in an async loop)
-            This makes sure client makes concurrent invocations
-        Await exit strategy:
-            :meth:`send_requests()` keeps track of num_requests sent
-            Async recv loop keeps track of num_responses received
-            Client exits out of await when num_requests == num_responses
-
         :param inputs: the callable
         :param on_done: the callback for on_done
         :param on_error: the callback for on_error
@@ -54,52 +91,40 @@ async def _get_results(
             import aiohttp
 
         self.inputs = inputs
+        request_iterator = self._get_requests(**kwargs)
 
-        req_iter = self._get_requests(**kwargs)
-        async with aiohttp.ClientSession() as session:
-
+        async with AsyncExitStack() as stack:
             try:
                 cm1 = ProgressBar() if self.show_progress else nullcontext()
                 proto = 'https' if self.args.https else 'http'
                 url = f'{proto}://{self.args.host}:{self.args.port}/post'
 
-                with cm1 as p_bar:
-                    all_responses = []
-                    for req in req_iter:
-                        # fix the mismatch between pydantic model and Protobuf model
-                        req_dict = req.dict()
-                        req_dict['exec_endpoint'] = req_dict['header']['exec_endpoint']
-                        req_dict['data'] = req_dict['data'].get('docs', None)
-
-                        all_responses.append(
-                            asyncio.create_task(
-                                self._get_http_response(
-                                    session,
-                                    url,
-                                    req_dict,
-                                )
-                            )
-                        )
-
-                    for resp in asyncio.as_completed(all_responses):
-                        r_status, r_str = await resp
-                        if r_status == 404:
-                            raise BadClient(f'no such endpoint {url}')
-                        elif r_status < 200 or r_status > 300:
-                            raise ValueError(r_str)
-
-                        resp = Request(r_str)
-                        resp = resp.as_typed_request(resp.request_type).as_response()
-                        callback_exec(
-                            response=resp,
-                            on_error=on_error,
-                            on_done=on_done,
-                            on_always=on_always,
-                            continue_on_error=self.continue_on_error,
-                            logger=self.logger,
-                        )
-                        if self.show_progress:
-                            p_bar.update()
-                        yield resp
+                p_bar = stack.enter_context(cm1)
+                iolet = await stack.enter_async_context(
+                    HTTPClientlet(url=url, logger=self.logger)
+                )
+
+                prefetch_caller = HTTPClientPrefetchCaller(self.args, iolet=iolet)
+                async for response in prefetch_caller.send(request_iterator):
+                    r_status = response.status
+                    r_str = await response.json()
+                    if r_status == 404:
+                        raise BadClient(f'no such endpoint {url}')
+                    elif r_status < 200 or r_status > 300:
+                        raise ValueError(r_str)
+
+                    resp = Request(r_str)
+                    resp = resp.as_typed_request(resp.request_type).as_response()
+                    callback_exec(
+                        response=resp,
+                        on_error=on_error,
+                        on_done=on_done,
+                        on_always=on_always,
+                        continue_on_error=self.continue_on_error,
+                        logger=self.logger,
+                    )
+                    if self.show_progress:
+                        p_bar.update()
+                    yield resp
             except aiohttp.client_exceptions.ClientConnectorError:
                 self.logger.warning(f'Client got disconnected from the HTTP server')
diff --git a/jina/parsers/__init__.py b/jina/parsers/__init__.py
index 7b7e6fd83ab1a..2bedfe0adb356 100644
--- a/jina/parsers/__init__.py
+++ b/jina/parsers/__init__.py
@@ -125,12 +125,16 @@ def set_client_cli_parser(parser=None):
 
         parser = set_base_parser()
 
-    from .peapods.runtimes.remote import mixin_client_gateway_parser
+    from .peapods.runtimes.remote import (
+        mixin_client_gateway_parser,
+        mixin_prefetch_parser,
+    )
     from .client import mixin_client_features_parser, mixin_comm_protocol_parser
 
     mixin_client_gateway_parser(parser)
     mixin_client_features_parser(parser)
     mixin_comm_protocol_parser(parser)
+    mixin_prefetch_parser(parser)
 
     return parser
 
diff --git a/jina/peapods/runtimes/gateway/grpc/__init__.py b/jina/peapods/runtimes/gateway/grpc/__init__.py
index e3b085cbd6281..22d3f542fa205 100644
--- a/jina/peapods/runtimes/gateway/grpc/__init__.py
+++ b/jina/peapods/runtimes/gateway/grpc/__init__.py
@@ -8,21 +8,11 @@
 from .....proto import jina_pb2_grpc
 from ....zmq import AsyncZmqlet
 from ...zmq.asyncio import AsyncNewLoopRuntime
-from ..prefetch import PrefetchCaller
+from ..prefetch import GrpcPrefetchCaller, ZmqPrefetchCaller
 
 __all__ = ['GRPCRuntime']
 
 
-class GRPCPrefetchCall(jina_pb2_grpc.JinaRPCServicer):
-    """JinaRPCServicer """
-
-    def __init__(self, args, iolet):
-        super().__init__()
-        self._servicer = PrefetchCaller(args, iolet=iolet)
-        self.Call = self._servicer.send
-        self.close = self._servicer.close
-
-
 class GRPCRuntime(AsyncNewLoopRuntime):
     """Runtime for gRPC."""
 
@@ -49,10 +39,10 @@ async def async_setup(self):
                 message_callback=None,
                 logger=self.logger,
             )
-            self._prefetcher = GRPCPrefetchCall(self.args, iolet=self._grpclet)
+            self._prefetcher = GrpcPrefetchCaller(self.args, iolet=self._grpclet)
         else:
             self.zmqlet = AsyncZmqlet(self.args, logger=self.logger)
-            self._prefetcher = GRPCPrefetchCall(self.args, iolet=self.zmqlet)
+            self._prefetcher = ZmqPrefetchCaller(self.args, iolet=self.zmqlet)
 
         jina_pb2_grpc.add_JinaRPCServicer_to_server(self._prefetcher, self.server)
         bind_addr = f'{__default_host__}:{self.args.port_expose}'
diff --git a/jina/peapods/runtimes/gateway/http/app.py b/jina/peapods/runtimes/gateway/http/app.py
index cdaee9d2f9cb4..6186f2681ac97 100644
--- a/jina/peapods/runtimes/gateway/http/app.py
+++ b/jina/peapods/runtimes/gateway/http/app.py
@@ -5,7 +5,6 @@
 
 from google.protobuf.json_format import MessageToDict
 
-from ..prefetch import PrefetchCaller
 from ....grpc import Grpclet
 from ....zmq import AsyncZmqlet
 from ..... import __version__
@@ -60,14 +59,19 @@ def get_fastapi_app(args: 'argparse.Namespace', logger: 'JinaLogger'):
         )
 
     if args.grpc_data_requests:
+        from ..prefetch import GrpcPrefetchCaller
+
         iolet = Grpclet(
             args=args,
             message_callback=None,
             logger=logger,
         )
+        servicer = GrpcPrefetchCaller(args, iolet)
     else:
+        from ..prefetch import ZmqPrefetchCaller
+
         iolet = AsyncZmqlet(args, logger)
-    servicer = PrefetchCaller(args, iolet)
+        servicer = ZmqPrefetchCaller(args, iolet)
 
     @app.on_event('shutdown')
     async def _shutdown():
@@ -203,14 +207,14 @@ async def _render_custom_swagger_html(req: Request) -> HTMLResponse:
 
         app.add_route(docs_url, _render_custom_swagger_html, include_in_schema=False)
 
-    async def _get_singleton_result(req_iter) -> Dict:
+    async def _get_singleton_result(request_iterator) -> Dict:
         """
         Streams results from AsyncPrefetchCall as a dict
 
-        :param req_iter: request iterator, with length of 1
+        :param request_iterator: request iterator, with length of 1
         :return: the first result from the request iterator
         """
-        async for k in servicer.send(request_iterator=req_iter):
+        async for k in servicer.send(request_iterator=request_iterator):
             return MessageToDict(
                 k, including_default_value_fields=True, use_integers_for_enums=True
             )  # DO NOT customize other serialization here. Scheme is handled by Pydantic in `models.py`
diff --git a/jina/peapods/runtimes/gateway/prefetch.py b/jina/peapods/runtimes/gateway/prefetch.py
index 82d865f03e214..36f0fa2fa3358 100644
--- a/jina/peapods/runtimes/gateway/prefetch.py
+++ b/jina/peapods/runtimes/gateway/prefetch.py
@@ -1,76 +1,66 @@
 import argparse
 import asyncio
 from asyncio import Future
-from typing import AsyncGenerator, Dict, Union
+from typing import AsyncGenerator, Dict, List, Union, TYPE_CHECKING
 
 from ...grpc import Grpclet
 from ....helper import typename, get_or_reuse_loop
 from ....logging.logger import JinaLogger
 from ....types.message import Message
 
-__all__ = ['PrefetchCaller']
+__all__ = ['ZmqPrefetchCaller', 'GrpcPrefetchCaller', 'HTTPClientPrefetchCaller']
 
-if False:
+if TYPE_CHECKING:
     from ...zmq import AsyncZmqlet
+    from ....clients.base.http import HTTPClientlet
+    from ....types.request import Request, Response
 
 
-class PrefetchCaller:
-    """An async zmq request sender to be used in the Gateway"""
+class BasePrefetchCaller:
+    """An base async request handler"""
 
     def __init__(
-        self, args: argparse.Namespace, iolet: Union['AsyncZmqlet', 'Grpclet']
+        self,
+        args: argparse.Namespace,
+        iolet: Union['AsyncZmqlet', 'Grpclet', 'HTTPClientlet'],
     ):
         """
         :param args: args from CLI
         :param iolet: One of AsyncZmqlet or Grpclet. Used for sending/receiving data to/from the Flow
         """
         self.args = args
-        self.name = args.name or self.__class__.__name__
-        self.logger = JinaLogger(self.name, **vars(args))
-        self._message_buffer: Dict[str, Future[Message]] = dict()
         self.iolet = iolet
+        self.receive_task = self._create_receive_task()
+        self.logger = JinaLogger(self.__class__.__name__, **vars(args))
 
-        if isinstance(iolet, Grpclet):
-            self.iolet.callback = self._unwrap_request
-            self._receive_task = get_or_reuse_loop().create_task(self.iolet.start())
-        else:
-            self._receive_task = get_or_reuse_loop().create_task(self._receive())
+    def _create_receive_task(self) -> asyncio.Task:
+        """Start a receive task to be running in the background
 
-    async def _unwrap_request(self, msg):
-        return self._process_message(msg.request)
+        .. # noqa: DAR202
+        :return: asyncio Task
+        """
+        raise NotImplementedError
 
-    async def _receive(self):
-        try:
-            while True:
-                message = await self.iolet.recv_message(callback=lambda x: x.response)
-                # during shutdown the socket will return None
-                if message is None:
-                    break
+    async def receive(self):
+        """Implement `receive` logic for prefetcher
 
-                self._process_message(message)
-        except asyncio.CancelledError:
-            raise
-        finally:
-            for future in self._message_buffer.values():
-                future.cancel(
-                    'PrefetchCaller closed, all outstanding requests canceled'
-                )
-            self._message_buffer.clear()
+        .. # noqa: DAR202
+        """
+        raise NotImplementedError
 
-    def _process_message(self, message):
-        if message.request_id in self._message_buffer:
-            future = self._message_buffer.pop(message.request_id)
-            future.set_result(message)
-        else:
-            self.logger.warning(
-                f'Discarding unexpected message with request id {message.request_id}'
-            )
+    def handle_request(self, request: 'Request', fetch_to: List):
+        """Handle each request in the iterator
+
+        :param request: current request in the iterator
+        :param fetch_to: list to add the task to
+        """
+        raise NotImplementedError
 
     async def close(self):
         """
         Stop receiving messages
         """
-        self._receive_task.cancel()
+        self.receive_task.cancel()
 
     async def send(self, request_iterator, *args) -> AsyncGenerator[None, Message]:
         """
@@ -84,7 +74,7 @@ async def send(self, request_iterator, *args) -> AsyncGenerator[None, Message]:
         self.iolet: Union['AsyncZmqlet', 'Grpclet']
         self.logger: JinaLogger
 
-        if self._receive_task.done():
+        if self.receive_task.done():
             raise RuntimeError(
                 'PrefetchCaller receive task not running, can not send messages'
             )
@@ -100,23 +90,15 @@ async def prefetch_req(num_req, fetch_to):
             for _ in range(num_req):
                 try:
                     if hasattr(request_iterator, '__anext__'):
-                        next_request = await request_iterator.__anext__()
+                        request = await request_iterator.__anext__()
                     elif hasattr(request_iterator, '__next__'):
-                        next_request = next(request_iterator)
+                        request = next(request_iterator)
                     else:
                         raise TypeError(
                             f'{typename(request_iterator)} does not have `__anext__` or `__next__`'
                         )
 
-                    future = get_or_reuse_loop().create_future()
-                    self._message_buffer[next_request.request_id] = future
-                    asyncio.create_task(
-                        self.iolet.send_message(
-                            Message(None, next_request, 'gateway', **vars(self.args))
-                        )
-                    )
-
-                    fetch_to.append(future)
+                    self.handle_request(request=request, fetch_to=fetch_to)
                 except (StopIteration, StopAsyncIteration):
                     return True
             return False
@@ -157,3 +139,132 @@ async def prefetch_req(num_req, fetch_to):
                 # this list dries, clear it and feed it with on_recv_task
                 prefetch_task.clear()
                 prefetch_task = [j for j in onrecv_task]
+
+
+class ZmqPrefetchCaller(BasePrefetchCaller):
+    """An async zmq request handler used in the Gateway"""
+
+    def __init__(
+        self,
+        args: argparse.Namespace,
+        iolet: Union['AsyncZmqlet', 'Grpclet', 'HTTPClientlet'],
+    ):
+        super().__init__(args, iolet)
+        self.request_buffer: Dict[str, Future[Message]] = dict()
+        self.Call = self.send  # Used in
+
+    def _create_receive_task(self):
+        """Start a receive task that starts the GRPC server & awaits termination.
+
+        :return: asyncio Task
+        """
+        return get_or_reuse_loop().create_task(self.receive())
+
+    def convert_to_message(self, request: 'Request'):
+        """Convert a `Request` to a `Message` to be sent from gateway
+
+        :param request: request from iterator
+        :return: Message from request
+        """
+        return Message(None, request, 'gateway', **vars(self.args))
+
+    def handle_request(self, request: 'Request', fetch_to: List):
+        """
+        For ZMQ & GRPC data requests, for each request in the iterator, we send the `Message` using
+        `iolet.send_message()` and add {<request-id>: <an-empty-future>} to the message buffer.
+        This empty future is used to track the `result` of this request during `receive`
+
+        :param request: current request in the iterator
+        :param fetch_to: list to add the task to
+        """
+        future = get_or_reuse_loop().create_future()
+        self.request_buffer[request.request_id] = future
+        asyncio.create_task(
+            self.iolet.send_message(self.convert_to_message(request=request))
+        )
+        fetch_to.append(future)
+
+    async def receive(self):
+        """Await messages back from Executors and process them in the message buffer"""
+        try:
+            while True:
+                response = await self.iolet.recv_message(callback=lambda x: x.response)
+                # during shutdown the socket will return None
+                if response is None:
+                    break
+
+                self.handle_response(response)
+        except asyncio.CancelledError:
+            raise
+        finally:
+            for future in self.request_buffer.values():
+                future.cancel(
+                    f'{self.__class__.__name__} closed, all outstanding requests canceled'
+                )
+            self.request_buffer.clear()
+
+    def handle_response(self, response: 'Response'):
+        """
+        Set result of each response received from Executors in the request buffer
+
+        :param response: message received during `iolet.recv_message`
+        """
+        if response.request_id in self.request_buffer:
+            future = self.request_buffer.pop(response.request_id)
+            future.set_result(response)
+        else:
+            self.logger.warning(
+                f'Discarding unexpected response with request id {response.request_id}'
+            )
+
+
+class GrpcPrefetchCaller(ZmqPrefetchCaller):
+    """An async grpc request handler used in the Gateway"""
+
+    def __init__(self, args: argparse.Namespace, iolet: 'Grpclet'):
+        super().__init__(args, iolet)
+        self.iolet.callback = lambda response: self.handle_response(response.request)
+
+    def _create_receive_task(self) -> 'asyncio.Task':
+        """Start a receive task that starts the GRPC server & awaits termination.
+
+        :return: asyncio Task
+        """
+        return get_or_reuse_loop().create_task(self.iolet.start())
+
+
+class HTTPClientPrefetchCaller(BasePrefetchCaller):
+    """An async HTTP request handler used in the HTTP Client"""
+
+    def convert_to_message(self, request: 'Request', **kwargs):
+        """Convert request to dict for POST request
+
+        :param request: request from client
+        :param kwargs: keyword args
+        :return: request as dict
+        """
+        req_dict = request.dict()
+        req_dict['exec_endpoint'] = req_dict['header']['exec_endpoint']
+        req_dict['data'] = req_dict['data'].get('docs', None)
+        return req_dict
+
+    def handle_request(self, request: 'Request', fetch_to: List):
+        """
+        For HTTP Client, for each request in the iterator, we send the message (http POST request)
+        and add it to the list of tasks which is awaited.
+
+        :param request: current request in the iterator
+        :param fetch_to: list to add the task to
+        """
+        fetch_to.append(
+            asyncio.create_task(
+                self.iolet.send_message(self.convert_to_message(request=request))
+            )
+        )
+
+    def _create_receive_task(self) -> 'asyncio.Task':
+        """For HTTP Client, there's no task needed for receiving. Sleep like there's no tomorrow!
+
+        :return: asyncio Task
+        """
+        return get_or_reuse_loop().create_task(asyncio.sleep(1e9))
diff --git a/jina/peapods/runtimes/gateway/websocket/app.py b/jina/peapods/runtimes/gateway/websocket/app.py
index 4762d37af4d23..e7074592ec1bc 100644
--- a/jina/peapods/runtimes/gateway/websocket/app.py
+++ b/jina/peapods/runtimes/gateway/websocket/app.py
@@ -2,7 +2,6 @@
 import inspect
 from typing import List
 
-from ..prefetch import PrefetchCaller
 from ....grpc import Grpclet
 from ....zmq import AsyncZmqlet
 from .....importer import ImportExtensions
@@ -38,14 +37,19 @@ def disconnect(self, websocket: WebSocket):
     app = FastAPI()
 
     if args.grpc_data_requests:
+        from ..prefetch import GrpcPrefetchCaller
+
         iolet = Grpclet(
             args=args,
             message_callback=None,
             logger=logger,
         )
+        servicer = GrpcPrefetchCaller(args, iolet)
     else:
+        from ..prefetch import ZmqPrefetchCaller
+
         iolet = AsyncZmqlet(args, logger)
-    servicer = PrefetchCaller(args, iolet)
+        servicer = ZmqPrefetchCaller(args, iolet)
 
     @app.on_event('shutdown')
     async def _shutdown():
diff --git a/jina/peapods/zmq/__init__.py b/jina/peapods/zmq/__init__.py
index 209db65a3dc37..1d4deea710f92 100644
--- a/jina/peapods/zmq/__init__.py
+++ b/jina/peapods/zmq/__init__.py
@@ -107,11 +107,11 @@ def _register_pollin(self):
             self.poller.register(self.in_connect_sock, zmq.POLLIN)
 
     def pause_pollin(self):
-        """Remove :attr:`in_sock` from the poller """
+        """Remove :attr:`in_sock` from the poller"""
         self.poller.unregister(self.in_sock)
 
     def resume_pollin(self):
-        """Put :attr:`in_sock` back to the poller """
+        """Put :attr:`in_sock` back to the poller"""
         self.poller.register(self.in_sock)
 
     @staticmethod
@@ -155,7 +155,7 @@ def _pull(self, interval: int = 1):
             return self.in_connect_sock
 
     def _close_sockets(self):
-        """Close input, output and control sockets of this `Zmqlet`. """
+        """Close input, output and control sockets of this `Zmqlet`."""
         for k in self.opened_socks:
             k.close()
 
@@ -277,7 +277,7 @@ def close(self, *args, **kwargs):
             self.print_stats()
 
     def print_stats(self):
-        """Print out the network stats of of itself """
+        """Print out the network stats of of itself"""
         self.logger.debug(
             f'#sent: {self.msg_sent} '
             f'#recv: {self.msg_recv} '
@@ -382,7 +382,7 @@ def _send_control_to_router(self, command, raise_exception=False):
         )
 
     def _send_idle_to_router(self):
-        """Tell the upstream router this dealer is idle """
+        """Tell the upstream router this dealer is idle"""
         self._send_control_to_router('IDLE')
 
     def _send_cancel_to_router(self, raise_exception=False):
@@ -567,12 +567,12 @@ def close(self, flush: bool = True, *args, **kwargs):
                     self.logger.error(f'failed to stop. {e!r}')
 
     def pause_pollin(self):
-        """Remove :attr:`in_sock` from the poller """
+        """Remove :attr:`in_sock` from the poller"""
         self.in_sock.stop_on_recv()
         self.is_polling_paused = True
 
     def resume_pollin(self):
-        """Put :attr:`in_sock` back to the poller """
+        """Put :attr:`in_sock` back to the poller"""
         if self.is_polling_paused:
             self.in_sock.on_recv(self._in_sock_callback)
             self.is_polling_paused = False


Code review comments:
diff_hunk: @@ -107,11 +107,11 @@ def _register_pollin(self):
             self.poller.register(self.in_connect_sock, zmq.POLLIN)
 
     def pause_pollin(self):
-        """Remove :attr:`in_sock` from the poller """
+        """Remove :attr:`in_sock` from the poller"""
path: jina/peapods/zmq/__init__.py
text: @user1:
Lets remove the changes in this file from the PR? I think they are only linting and this PR should not touch this file it appears

diff_hunk: @@ -100,23 +89,15 @@ async def prefetch_req(num_req, fetch_to):
             for _ in range(num_req):
                 try:
                     if hasattr(request_iterator, '__anext__'):
-                        next_request = await request_iterator.__anext__()
+                        request = await request_iterator.__anext__()
                     elif hasattr(request_iterator, '__next__'):
-                        next_request = next(request_iterator)
+                        request = next(request_iterator)
                     else:
                         raise TypeError(
                             f'{typename(request_iterator)} does not have `__anext__` or `__next__`'
                         )
 
-                    future = get_or_reuse_loop().create_future()
-                    self._message_buffer[next_request.request_id] = future
-                    asyncio.create_task(
-                        self.iolet.send_message(
-                            Message(None, next_request, 'gateway', **vars(self.args))
-                        )
-                    )
-
-                    fetch_to.append(future)
+                    self.handle_request(request=request, fetch_to=fetch_to)
path: jina/peapods/runtimes/prefetch/base.py
text: @user1:
lets return the task in the `handle_request` and append it to the list here? Otherwise we have to expose the fetch_to list to the child class

diff_hunk: @@ -157,3 +139,132 @@ async def prefetch_req(num_req, fetch_to):
                 # this list dries, clear it and feed it with on_recv_task
                 prefetch_task.clear()
                 prefetch_task = [j for j in onrecv_task]
+
+
+class ZmqPrefetchCaller(BasePrefetchCaller):
path: jina/peapods/runtimes/gateway/prefetch.py
text: @user1:
Whats the reasoning to further split the BasePrefetchCaller into this ZmqPrefetchCaller? I think that is slightly misleading, because the ZmqPrefetchCaller actually can use either gRPC or ZMQ. 
So for me it would make more sense to have one abstract (BasePrefetchCaller) Prefetcher which can handle all Flow internal traffic (either zmq or GRPC), but has no functionality to communicate those results back to the outside world.

Then we can have concrete Http/Grpc/Websocket Prefetchers implementing the missing communication logic to the outside.
It might even be better to use compositon over inheritance here? So some generic BasePrefetchCaller is just used by some http/grpc/websocket server wrapping code? I think functionally speaking we mix two different things here: Comunication with the `Flow` and the `Client`