Repository: jina-ai/jina
File: jina/serve/runtimes/worker/batch_queue.py (Python)

Code before:

        :param request: The request to append to the queue.
        :param http: Flag to determine if the request is served via HTTP for some optims

        :return: The queue that will receive when the request is processed.
        """
        docs = request.docs

        if not self._timer_task or self._timer_finished:
            # If there is no timer (first arrival), or the timer is already consumed, any new push should trigger a new Timer, before
            # this push requests the data lock. The order of accessing the data lock guarantees that this request will be put in the `big_doc`
            # before the `flush` task processes it.
            self._start_timer()
        async with self._data_lock:
            if not self._flush_task:
                self._flush_task = asyncio.create_task(self._await_then_flush(http))

            self._big_doc.extend(docs)
            next_req_idx = len(self._requests)
            num_docs = len(docs)
            self._request_idxs.extend([next_req_idx] * num_docs)
            self._request_lens.append(len(docs))
            self._requests.append(request)
            queue = asyncio.Queue()
            self._requests_completed.append(queue)
            if self._custom_metric is not None:
                if len(self._big_doc) >= self._preferred_batch_size:
                    self._flush_trigger.set()
            else:
                if self._custom_metric(self._big_doc) >= self._preferred_batch_size:
                    self._flush_trigger.set()

        return queue

    async def _await_then_flush(self, http=False) -> None:
        """Process all requests in the queue once flush_trigger event is set.
        :param http: Flag to determine if the request is served via HTTP for some optims
        """

        def _get_docs_groups_completed_request_indexes(
                non_assigned_docs,
                non_assigned_docs_reqs_idx,
                sum_from_previous_mini_batch_in_first_req_idx,
        ):
            """
            This method groups all the `non_assigned_docs` into groups of docs according to the `req_idx` they belong to.
            They are only distributed when we are sure that the request is full.

            :param non_assigned_docs: The documents that have already been processed but have not been assigned to a request result
            :param non_assigned_docs_reqs_idx: The request IDX that are not yet completed (not all of its docs have been processed)
            :param sum_from_previous_mini_batch_in_first_req_idx: The number of docs from previous iteration that belong to the first non_assigned_req_idx. This is useful to make sure we know when a request is completed.

Review comment (suggestion):
```suggestion
            if self._custom_metric is None:
```

Code after:
        """Sleep and then set the event"""
        self._timer_finished = False
        await asyncio.sleep(self._timeout / 1000)
        self._flush_trigger.set()
        self._timer_finished = True

    async def push(self, request: DataRequest, http=False) -> asyncio.Queue:
        """Append request to the the list of requests to be processed.

        This method creates an asyncio Queue for that request and keeps track of it. It returns
        this queue to the caller so that the caller can now when this request has been processed

        :param request: The request to append to the queue.
        :param http: Flag to determine if the request is served via HTTP for some optims

        :return: The queue that will receive when the request is processed.
        """
        docs = request.docs

        if not self._timer_task or self._timer_finished:
            # If there is no timer (first arrival), or the timer is already consumed, any new push should trigger a new Timer, before
            # this push requests the data lock. The order of accessing the data lock guarantees that this request will be put in the `big_doc`
            # before the `flush` task processes it.
            self._start_timer()
        async with self._data_lock:
            if not self._flush_task:
                self._flush_task = asyncio.create_task(self._await_then_flush(http))

            self._big_doc.extend(docs)
            next_req_idx = len(self._requests)
            num_docs = len(docs)
            metric_value = num_docs
            if self._custom_metric is not None:
                metrics = [self._custom_metric(doc) for doc in docs]
                metric_value += sum(metrics)
                self._docs_metrics.extend(metrics)
            self._metric_value += metric_value
            self._request_idxs.extend([next_req_idx] * num_docs)
            self._request_lens.append(num_docs)
            self._requests.append(request)
            queue = asyncio.Queue()
            self._requests_completed.append(queue)
            if self._metric_value >= self._preferred_batch_size:
                self._flush_trigger.set()

        return queue

    async def _await_then_flush(self, http=False) -> None:
        """Process all requests in the queue once flush_trigger event is set.
        :param http: Flag to determine if the request is served via HTTP for some optims
        """