Repository: jina-ai/jina
File: jina/clients/base/helper.py (Python)

Code before:
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.websocket = None
        self.response_iter = None

    async def send_message(self, request: 'Request'):
        """Send request in bytes to the server.

        :param request: request object
        :return: send request as bytes awaitable
        """
        # we have to send here `retries` information
        for retry in range(self.max_attempts):
            try:
                res = await self.websocket.send_bytes(request.SerializeToString())
                return res
            except ConnectionResetError:
                if retry == self.max_attempts - 1:
                    self.logger.critical(f'server connection closed already!')
                    raise
                else:
                    if retry == 0:
                        wait_time = random.uniform(0, self.initial_backoff)
                    else:
                        wait_time = random.uniform(0, min(self.initial_backoff*self.backoff_multiplier**(retry-1), self.max_backoff))
                    await asyncio.sleep(wait_time)

    async def send_dry_run(self):
        """Query the dry_run endpoint from Gateway

        :return: send dry_run as bytes awaitable
        """

        try:
            return await self.websocket.send_bytes(b'')
        except ConnectionResetError:
            self.logger.critical(f'server connection closed already!')

    async def send_eoi(self):
        """To confirm end of iteration, we send `bytes(True)` to the server.

        :return: send `bytes(True)` awaitable
        """
        try:
            return await self.websocket.send_bytes(bytes(True))
        except ConnectionResetError:
            # server might be in a `CLOSING` state while sending EOI signal
            # which raises a `ConnectionResetError`, this can be ignored.
            pass

    async def recv_message(self) -> 'DataRequest':

Review comment (suggestion):
```suggestion
        for retry in range(1, self.max_attempts + 1):
            try:
                return await self.websocket.send_bytes(request.SerializeToString())
            except ConnectionResetError:
                if retry == self.max_attempts:
                    self.logger.critical(f'server connection closed already!')
                    raise
                else:
                    wait_time = random.uniform(0, min(self.initial_backoff*self.backoff_multiplier**(retry-1), self.max_backoff))
                    await asyncio.sleep(wait_time)
```

Code after:

    async def send_message(self, request: 'Request'):
        """Send request in bytes to the server.

        :param request: request object
        :return: send request as bytes awaitable
        """
        # we have to send here `retries` information
        for retry in range(1, self.max_attempts + 1):
            try:
                return await self.websocket.send_bytes(request.SerializeToString())
            except ConnectionResetError:
                if retry == self.max_attempts:
                    self.logger.critical(f'server connection closed already!')
                    raise
                else:
                    wait_time = random.uniform(0, min(self.initial_backoff*self.backoff_multiplier**(retry-1), self.max_backoff))
                    await asyncio.sleep(wait_time)

    async def send_dry_run(self):
        """Query the dry_run endpoint from Gateway

        :return: send dry_run as bytes awaitable
        """

        try:
            return await self.websocket.send_bytes(b'')
        except ConnectionResetError:
            self.logger.critical(f'server connection closed already!')

    async def send_eoi(self):
        """To confirm end of iteration, we send `bytes(True)` to the server.

        :return: send `bytes(True)` awaitable
        """
        try:
            return await self.websocket.send_bytes(bytes(True))
        except ConnectionResetError:
            # server might be in a `CLOSING` state while sending EOI signal
            # which raises a `ConnectionResetError`, this can be ignored.
            pass

    async def recv_message(self) -> 'DataRequest':
        """Receive messages in bytes from server and convert to `DataRequest`

        ..note::
            aiohttp allows only one task which can `receive` concurrently.
            we need to make sure we don't create multiple tasks with `recv_message`

        :yield: response objects received from server
        """