diff --git a/chia/farmer/farmer.py b/chia/farmer/farmer.py index c06d3c3c2471..33770f279833 100644 --- a/chia/farmer/farmer.py +++ b/chia/farmer/farmer.py @@ -313,7 +313,7 @@ class Farmer: value=ErrorResponse(uint16(PoolErrorCode.REQUEST_FAILED.value), error_message).to_json_dict(), ) - def on_disconnect(self, connection: WSChiaConnection) -> None: + async def on_disconnect(self, connection: WSChiaConnection) -> None: self.log.info(f"peer disconnected {connection.get_peer_logging()}") self.state_changed("close_connection", {}) if connection.connection_type is NodeType.HARVESTER: diff --git a/chia/full_node/full_node.py b/chia/full_node/full_node.py index c15ac2218abb..dfe1bc0b638c 100644 --- a/chia/full_node/full_node.py +++ b/chia/full_node/full_node.py @@ -895,7 +895,7 @@ class FullNode: elif connection.connection_type is NodeType.TIMELORD: await self.send_peak_to_timelords() - def on_disconnect(self, connection: WSChiaConnection) -> None: + async def on_disconnect(self, connection: WSChiaConnection) -> None: self.log.info(f"peer disconnected {connection.get_peer_logging()}") self._state_changed("close_connection") self._state_changed("sync_mode") diff --git a/chia/harvester/harvester.py b/chia/harvester/harvester.py index f1545aa19480..8a539a02fdc1 100644 --- a/chia/harvester/harvester.py +++ b/chia/harvester/harvester.py @@ -173,7 +173,7 @@ class Harvester: if event == PlotRefreshEvents.done: self.plot_sync_sender.sync_done(update_result.removed, update_result.duration) - def on_disconnect(self, connection: WSChiaConnection) -> None: + async def on_disconnect(self, connection: WSChiaConnection) -> None: self.log.info(f"peer disconnected {connection.get_peer_logging()}") self.state_changed("close_connection") self.plot_sync_sender.stop() diff --git a/chia/server/server.py b/chia/server/server.py index 92f5e38a98f9..63b171b30c82 100644 --- a/chia/server/server.py +++ b/chia/server/server.py @@ -522,7 +522,9 @@ class ChiaServer: return False - def connection_closed(self, connection: WSChiaConnection, ban_time: int, closed_connection: bool = False) -> None: + async def connection_closed( + self, connection: WSChiaConnection, ban_time: int, closed_connection: bool = False + ) -> None: # closed_connection is true if the callback is being called with a connection that was previously closed # in this case we still want to do the banning logic and remove the conection from the list # but the other cleanup should already have been done so we skip that @@ -555,7 +557,7 @@ class ChiaServer: connection.cancel_tasks() on_disconnect = getattr(self.node, "on_disconnect", None) if on_disconnect is not None: - on_disconnect(connection) + await on_disconnect(connection) async def validate_broadcast_message_type(self, messages: List[Message], node_type: NodeType) -> None: for message in messages: diff --git a/chia/server/ws_connection.py b/chia/server/ws_connection.py index 889eebcb14dc..a09723bd4b36 100644 --- a/chia/server/ws_connection.py +++ b/chia/server/ws_connection.py @@ -52,7 +52,7 @@ def create_default_last_message_time_dict() -> Dict[ProtocolMessageTypes, float] class ConnectionClosedCallbackProtocol(Protocol): - def __call__( + async def __call__( self, connection: WSChiaConnection, ban_time: int, @@ -277,7 +277,7 @@ class WSChiaConnection: with log_exceptions(self.log, consume=True): self.log.debug(f"Closing already closed connection for {self.peer_info.host}") if self.close_callback is not None: - self.close_callback(self, ban_time, closed_connection=True) + await self.close_callback(self, ban_time, closed_connection=True) self._close_event.set() return None self.closed = True @@ -307,7 +307,7 @@ class WSChiaConnection: finally: with log_exceptions(self.log, consume=True): if self.close_callback is not None: - self.close_callback(self, ban_time, closed_connection=False) + await self.close_callback(self, ban_time, closed_connection=False) self._close_event.set() async def wait_until_closed(self) -> None: diff --git a/chia/wallet/wallet_node.py b/chia/wallet/wallet_node.py index 349c48101664..0418feb19bf9 100644 --- a/chia/wallet/wallet_node.py +++ b/chia/wallet/wallet_node.py @@ -691,7 +691,7 @@ class WalletNode: ) asyncio.create_task(self.wallet_peers.start()) - def on_disconnect(self, peer: WSChiaConnection) -> None: + async def on_disconnect(self, peer: WSChiaConnection) -> None: if self.is_trusted(peer): self.local_node_synced = False self.initialize_wallet_peers()