NodeAgent

async def get(self, node_id: str) -> Optional[NodeAgent]:
    """Gets node agent by given node id

    Args:
        node_id: Id of the node, or None if no agent exists for this node ID
    """
    async with self._store_lock:
        return self._node_agents.get(node_id)

async def get_all(self) -> Dict[str, NodeAgent]:
    """Returns all node agents regardless of their status (ACTIVE, DISCONNECTED, ...).

    Returns:
        Dictionary of node agent objects, by node ID
    """

    async with self._store_lock:
        # a shallow copy is wanted so that
        # - we have a distinct (stable) list of NodeAgents that can be processed in calling func
        # - we use same NodeAgents objects (not a copy)
        return copy.copy(self._node_agents)

async def retrieve(self, node_id: str) -> NodeAgent:
    """Retrieves a node agent for a given node ID.

    Depending if this node is already known to the store this method gets existing agent or.
    registers a new agent.

    Args:
        node_id: ID of receiving node

    Return:
        The node agent to manage tasks that are assigned to it.
    """
    # Lock during all sequence to ensure atomicity
    async with self._store_lock:
        node = self._node_agents.get(node_id)
        if not node:
            node = NodeAgent(
                id=node_id, loop=self._loop, on_forward=self._on_forward
            )
            self._node_agents.update({node_id: node})

    return node

def flush(self, request_id: str, stopped: bool = False) -> None:
    """Flush processed replies

    Args:
        request_id: request ID for which the replies should be flushed
        stopped: the request was stopped during processing
    """
    asyncio.run_coroutine_threadsafe(super().flush(request_id, stopped), self._loop)

def send(self, message: Message, on_reply: Optional[Callable] = None) -> None:
    """Send message to researcher.

    Args:
        message: Message to send to the researcher
    """
    asyncio.run_coroutine_threadsafe(
        self.send_async(message=message, on_reply=on_reply), self._loop
    )

async def change_node_status_after_task(self) -> None:
    """Coroutine to execute each time RPC call is completed"""
    async with self._status_lock:
        self._status = NodeActiveStatus.WAITING

        if self._status_task is None:
            self._status_task = asyncio.create_task(
                self._change_node_status_disconnected()
            )

async def flush(self, request_id: str, stopped: bool = False) -> None:
    """Flushes processed reply

    Args:
        request_id: request ID for which the replies should be flushed
        stopped: the request was stopped during processing
    """
    async with self._replies_lock:
        if stopped and self._replies[request_id]["reply"] is None:
            async with self._stopped_request_ids_lock:
                self._stopped_request_ids.append(request_id)

        self._replies.pop(request_id, None)

def get_task(self) -> Awaitable[Message]:
    """Get tasks assigned by the main thread

    !!! note "Returns coroutine"
        This function return an asyncio coroutine. Please use `await` while calling.

    Returns:
        A coroutine to await for retrieving a list of: a task ; a number of send retries already done ;
            the time of first sending attempt in seconds since epoch
    """
    return self._queue.get()

async def on_reply(self, message: Dict) -> None:
    """Callback to execute each time new reply received from the node"""

    message = Message.from_dict(message)

    # Handle overlay messages to relay to a node
    if isinstance(message, OverlayMessage):
        await self._on_forward(message)
        return

    # Handle RequestReply messages for the researcher
    if not message.request_id:
        logger.error(
            f"Server received a reply from the client {self._id} that does "
            "not contains request id."
        )

    async with self._replies_lock:
        if message.request_id in self._replies:
            if self._replies[message.request_id]["reply"] is None:
                self._replies[message.request_id]["reply"] = message
                self._replies[message.request_id]["callback"](message)
            else:
                # Handle case of multiple replies
                # Avoid conflict with consumption of reply.
                logger.warning(
                    f"Received multiple replies for request {message.request_id}. "
                    "Keep first reply, ignore subsequent replies"
                )
        else:
            async with self._stopped_request_ids_lock:
                if message.request_id in self._stopped_request_ids:
                    logger.warning(
                        "Received a reply from a federated request that has been "
                        f"stopped: {message.request_id}."
                    )
                    self._stopped_request_ids.remove(message.request_id)
                else:
                    logger.warning(
                        f"Received a reply from an unexpected request: {message.request_id}"
                    )

async def send_async(
    self,
    message: Message,
    on_reply: Optional[Callable] = None,
    retry_count: int = 0,
    first_send_time: Optional[float] = None,
) -> None:
    """Async function send message to researcher.

    Args:
        message: Message to send to the researcher
        on_reply: optional callback to execute when receiving message reply
        retry_count: number of retries already done for this message
        first_send_time: time of first send attempt for this message
    """

    async with self._status_lock:
        if self._status == NodeActiveStatus.DISCONNECTED:
            logger.info(f"Node {self._id} is disconnected. Discard message.")
            return

        if self._status == NodeActiveStatus.WAITING:
            logger.info(
                f"Node {self._id} is in WAITING status. Server is "
                "waiting for receiving a request from "
                "this node to convert it as ACTIVE. Node will be updated "
                "as DISCONNECTED soon if no request received."
            )

    # Updates replies
    #
    # Note: as forwarded messages don't have a `request_id` field we don't have to test
    # if this is an OverlayMessage but check whether the field exists
    async with self._replies_lock:
        # update replies only for (1) request-response messages
        # (2) that are not yet registered as pending request
        if (
            hasattr(message, "request_id")
            and message.request_id
            and message.request_id not in self._replies
        ):
            self._replies.update(
                {message.request_id: {"callback": on_reply, "reply": None}}
            )

    if first_send_time is None:
        first_send_time = time.time()
    await self._queue.put([message, retry_count, first_send_time])

async def set_active(self) -> None:
    """Updates node status as active"""

    async with self._status_lock:

        # Inform user that node is online again
        if self._status == NodeActiveStatus.DISCONNECTED:
            logger.info(f"Node {self._id} is back online!")

        self._status = NodeActiveStatus.ACTIVE

        # Cancel status task if there is any running
        if self._status_task:
            self._status_task.cancel()
            self._status_task = None

async def status_async(self) -> NodeActiveStatus:
    """Getter for node status.

    Returns:
        node status
    """
    async with self._status_lock:
        # (deep)copy is not needed as long as this remains a simple value ...
        return self._status

def task_done(self) -> None:
    """Acknowledge completion of a de-queued task"""
    self._queue.task_done()