Async connection close handlers (#17017)

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:

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")

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()

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:

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:

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()