diff --git a/mafic/node.py b/mafic/node.py index 93030da..f6a28fb 100644 --- a/mafic/node.py +++ b/mafic/node.py @@ -713,7 +713,16 @@ async def close(self) -> None: _log.debug("Connection task cancelled.", extra={"label": self._label}) _log.info("Node %s is now closed.", self._label, extra={"label": self._label}) + self.cleanup() + + def cleanup(self) -> None: + """Cleanup the node, as if it was never connected.""" self._available = False + self._ws = None + self._ws_task = None + self._connect_task = None + self.__session = None + self._ready.clear() async def _ws_listener(self) -> None: """Listen for messages from the websocket.""" diff --git a/mafic/pool.py b/mafic/pool.py index f2bc1c4..791d112 100644 --- a/mafic/pool.py +++ b/mafic/pool.py @@ -187,6 +187,25 @@ async def create_node( resuming_session_id=resuming_session_id, ) + await self.add_node(node) + return node + + async def add_node(self, node: Node[ClientT]) -> None: + """Add an existing node to this pool. + + .. note:: + + You generally do not want this, use :meth:`create_node` instead. + This is used for after running :meth:`remove_node` to re-add the node + if it has been restarted. + + .. versionadded:: 2.7 + + Parameters + ---------- + node: + The node to add. + """ # Add to dictionaries, creating a set or extending it if needed. if node.regions: for region in node.regions: @@ -202,11 +221,10 @@ async def create_node( *self._node_shards.get(shard_id, set()), } - _log.info("Created node, connecting it...", extra={"label": label}) + _log.info("Created node, connecting it...", extra={"label": node.label}) await node.connect() - self._nodes[label] = node - return node + self._nodes[node.label] = node async def remove_node( self, node: Node[ClientT] | str, *, transfer_players: bool = True