Repository: jina-ai/jina
File: jina/serve/networking.py (Python)

Code before:

    @staticmethod
    def __aio_channel_with_tracing_interceptor(
        address,
        credentials=None,
        options=None,
        interceptors=None,
    ) -> grpc.aio.Channel:
        if credentials:
            return grpc.aio.secure_channel(
                address,
                credentials,
                options=options,
                interceptors=interceptors,
            )
        return grpc.aio.insecure_channel(
            address,
            options=options,
            interceptors=interceptors,
        )

    @staticmethod
    def __channel_with_tracing_interceptor(
        address,
        credentials=None,
        options=None,
        interceptor=None,
    ) -> grpc.Channel:
        if credentials:
            channel = grpc.secure_channel(address, credentials, options=options)
        else:
            channel = grpc.insecure_channel(address, options=options)

        if interceptor:
            from opentelemetry.instrumentation.grpc.grpcext import intercept_channel

            return intercept_channel(
                channel,
                interceptor,
            )
        else:
            return channel

    @staticmethod
    def get_grpc_channel(
        address: str,
        options: Optional[list] = None,
        asyncio: bool = False,
        tls: bool = False,
        root_certificates: Optional[str] = None,
        aio_tracing_client_interceptors: Optional[Sequence['ClientInterceptor']] = None,

Review comment (refactor):
I can extract out the warmup logic into a `Warmer` class but the `self._connections` is the inner class `_ConnectionPoolMap` of the `GrpcConnectionPool`. Does the inner class still make sense? 

Code after:
                    for stub in replica.warmup_stubs:
                        if stub.address not in warmed_up_targets:
                            tasks.append(
                                asyncio.create_task(
                                    task_wrapper(replica_warmup_responses, stub)
                                )
                            )

                await asyncio.gather(*tasks, return_exceptions=True)
                for target, response in replica_warmup_responses.items():
                    if response:
                        warmed_up_targets.add(target)

                now = time.time()
                if now > timeout or all(list(replica_warmup_responses.values())):
                    self._logger.debug(f'completed warmup task in {now - start_time}s.')
                    return

                await asyncio.sleep(0.2)
        except Exception as ex:
            self._logger.error(f'error with warmup up task: {ex}')
            return

    def _get_all_replicas(self, deployment):
        replica_set = set()
        replica_set.update(self._connections.get_replicas_all_shards(deployment))
        replica_set.add(
            self._connections.get_replicas(deployment=deployment, head=True)
        )

        return set(filter(None, replica_set))

    @staticmethod
    def __aio_channel_with_tracing_interceptor(
        address,
        credentials=None,
        options=None,
        interceptors=None,
    ) -> grpc.aio.Channel:
        if credentials:
            return grpc.aio.secure_channel(
                address,
                credentials,
                options=options,
                interceptors=interceptors,
            )
        return grpc.aio.insecure_channel(
            address,
            options=options,
            interceptors=interceptors,
        )