Server

def broadcast(self, message: Message) -> None:
    """Broadcast message to all known and reachable nodes

    Args:
        message: Message to broadcast

    Raises:
        FedbiomedCommunicationError: bad argument type
        FedbiomedCommunicationError: server is not started
    """
    if not isinstance(message, Message):
        raise FedbiomedCommunicationError(
            f"{ErrorNumbers.FB628}: bad argument type for message, expected `Message`, got `{type(message)}`")

    if not self._is_started.is_set():
        raise FedbiomedCommunicationError(
            f"{ErrorNumbers.FB628}: Can not broadcast given message. "
            "Communication client is not initialized.")

    self._run_threadsafe(super().broadcast(message))

def get_all_nodes(self) -> List[NodeAgent]:
    """Returns all known nodes

    Returns:
        A list of node agents

    Raises:
        FedbiomedCommunicationError: server is not started
    """
    if not self._is_started.is_set():
        raise FedbiomedCommunicationError(
            f"{ErrorNumbers.FB628}: Error while getting all nodes "
            "connected:  Communication client is not initialized.")

    return self._run_threadsafe(super().get_all_nodes())

def get_node(self, node_id) -> Optional[NodeAgent]:
    """Returns given node

    Args:
        node_id: ID of node to retrieve

    Returns:
        A node agent

    Raises:
        FedbiomedCommunicationError: server is not started
    """
    if not self._is_started.is_set():
        raise FedbiomedCommunicationError(
            f"{ErrorNumbers.FB628}: Error while getting node '{node_id}':"
            "Communication client is not initialized.")

    return self._run_threadsafe(super().get_node(node_id))

def is_alive(self) -> bool:
    """Checks if the thread running gRPC server still alive

    Returns:
        gRPC server running status

    Raises:
        FedbiomedCommunicationError: server is not started
    """
    if not self._is_started.is_set():
        raise FedbiomedCommunicationError(
            f"{ErrorNumbers.FB628}: Can not check if thread is alive."
            "Communication client is not initialized.")

    # TODO: more tests about gRPC server and task status ?
    return False if not isinstance(self._thread, threading.Thread) else self._thread.is_alive()

def send(self, message: Message, node_id: str) -> None:
    """Send message to a specific node.

    Args:
        message: Message to send
        node_id: Destination node unique ID

    Raises:
        FedbiomedCommunicationError: bad argument type
        FedbiomedCommunicationError: server is not started
    """
    if not isinstance(message, Message):
        raise FedbiomedCommunicationError(
            f"{ErrorNumbers.FB628}: bad argument type for message, expected `Message`, got `{type(message)}`")

    if not self._is_started.is_set():
        raise FedbiomedCommunicationError(
            f"{ErrorNumbers.FB628.value}: Can not send message. "
            "Communication client is not initialized.")

    self._run_threadsafe(super().send(message, node_id))

def start(self) -> None:
    """Starts async GrpcServer """

    self._thread = threading.Thread(target=self._run, daemon=True)
    self._thread.start()

    # FIXME: This implementation assumes that nodes will be able connect and server complete setup with this delay
    logger.info("Starting researcher service...")


    logger.info(f'Waiting {GRPC_SERVER_SETUP_TIMEOUT}s for nodes to connect...')
    time.sleep(GRPC_SERVER_SETUP_TIMEOUT)

    sleep_ = 0
    while len(self.get_all_nodes()) == 0:

        if sleep_ == 0:
            logger.info(f"No nodes found, server will wait "
                        f"{MAX_GRPC_SERVER_SETUP_TIMEOUT - GRPC_SERVER_SETUP_TIMEOUT} "
                        "more seconds until a node creates connection.")

        if sleep_ > MAX_GRPC_SERVER_SETUP_TIMEOUT - GRPC_SERVER_SETUP_TIMEOUT:
            if len(self.get_all_nodes()) == 0:
                logger.warning("Server has not received connection from any remote nodes in "
                               f"MAX_GRPC_SERVER_SETUP_TIMEOUT: {MAX_GRPC_SERVER_SETUP_TIMEOUT} "
                               "This may effect the request created right after the server initialization. "
                               "However, server will keep running in the background so you can retry the "
                               "operations for sending requests to remote nodes until one receives.")
            break

        time.sleep(1)
        sleep_ += 1

async def Feedback(
        self,
        request: ProtoBufMessage,
        unused_context: grpc.aio.ServicerContext
) -> None:
    """Executed for Feedback request received from the nodes

    Args:
        request: Feedback message
        unused_context: Request service context
    """

    # Get the type of Feedback | log or scalar
    one_of = request.WhichOneof("feedback_type")
    feedback = FeedbackMessage.from_proto(request)

    # Execute on_message assigned by the researcher.requests modules
    self._on_message(feedback.get_param(one_of), MessageType.convert(one_of))

    return Empty()

async def GetTaskUnary(
        self,
        request: ProtoBufMessage,
        context: grpc.aio.ServicerContext
) -> None:
    """Gets unary RPC request and return stream of response

    Args:
        request: RPC request
        context: RPC peer context
    """
    task_request = TaskRequest.from_proto(request).get_dict()
    logger.debug(f"Node: {task_request.get('node')} polling for the tasks")

    node_agent = await self._agent_store.retrieve(node_id=task_request["node"])

    # Update node active status as active
    await node_agent.set_active()

    task = None
    try:
        while True:
            task, retry_count, first_send_time = await node_agent.get_task()

            # Choice: mark task as de-queued as soon only if really sent
            node_agent.task_done()

            # discard if message too old
            if first_send_time + MAX_SEND_DURATION > time.time():
                break
            else:
                task = None
                logger.warning(f"Message to send is older than {MAX_SEND_DURATION} seconds. Discard message.")

        task_bytes = Serializer.dumps(task.get_dict())

        chunk_range = range(0, len(task_bytes), MAX_MESSAGE_BYTES_LENGTH)
        for start, iter_ in zip(chunk_range, range(1, len(chunk_range) + 1)):
            stop = start + MAX_MESSAGE_BYTES_LENGTH

            try:
                yield TaskResponse(
                    size=len(chunk_range),
                    iteration=iter_,
                    bytes_=task_bytes[start:stop]
                ).to_proto()
            except GeneratorExit:
                # schedule resend if task sending could not be completed
                # => retry send as long as (1) send not successful
                # (2) max retries not reached
                # => else discard message
                #
                # Note: if node is disconnected then back online, message is retried after reconnection.
                # This is not fully coherent with upper layers (Requests) that may trigger an application
                # level failure in the while, but it is mitigated by the MAX_SEND_DURATION
                if retry_count < MAX_SEND_RETRIES:
                    await node_agent.send_async(
                        message=task, on_reply=None, retry_count=retry_count + 1, first_send_time=first_send_time
                    )
                else:
                    logger.warning(f"Message cannot be sent after {MAX_SEND_RETRIES} retries. Discard message.")
                await node_agent.change_node_status_after_task()
                # need return here to avoid RuntimeError
                return

    except asyncio.CancelledError:
        if task is not None and retry_count is not None and first_send_time is not None:
            # schedule resend if task was pulled from queue
            if retry_count < MAX_SEND_RETRIES:
                await node_agent.send_async(
                    message=task, on_reply=None, retry_count=retry_count + 1, first_send_time=first_send_time
                )
            else:
                logger.warning(f"Message cannot be sent after {MAX_SEND_RETRIES} retries. Discard message.")
    finally:
        await node_agent.change_node_status_after_task()

async def ReplyTask(
        self,
        request_iterator: Iterable[ProtoBufMessage],
        unused_context: grpc.aio.ServicerContext
) -> None:
    """Gets stream replies from the nodes

    Args:
        request_iterator: Iterator for streaming
        unused_context: Request service context
    """

    reply = bytes()
    async for answer in request_iterator:
        reply += answer.bytes_
        if answer.size != answer.iteration:
            continue
        else:
            # Deserialize message
            message = Serializer.loads(reply)

            # Replies are handled by node agent callbacks
            node = await self._agent_store.get(message["node_id"])
            await node.on_reply(message)

            reply = bytes()

    return Empty()