Fix self connection leak (#43)

* Fix self connection leak
* Increase sleep to improve performance
* Add load test for unfinished blocks
* Fix flake8 lint

Author: mariano54

.gitignore

@@ -42,7 +42,7 @@ pip-delete-this-directory.txt
 
 # pyenv
 .python-version
-
+.eggs
 .venv
 
 # mypy

src/full_node.py

@@ -43,6 +43,7 @@
 from src.util.ints import uint32, uint64
 
 log = logging.getLogger(__name__)
+OutboundMessageGenerator = AsyncGenerator[OutboundMessage, None]
 
 
 class FullNode:
@@ -61,7 +62,7 @@ def _set_server(self, server: ChiaServer):
 
     async def _send_tips_to_farmers(
         self, delivery: Delivery = Delivery.BROADCAST
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         Sends all of the current heads to all farmer peers. Also sends the latest
         estimated proof of time rate, so farmer can calulate which proofs are good.
@@ -99,7 +100,7 @@ async def _send_tips_to_farmers(
 
     async def _send_challenges_to_timelords(
         self, delivery: Delivery = Delivery.BROADCAST
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         Sends all of the current heads (as well as Pos infos) to all timelord peers.
         """
@@ -137,7 +138,7 @@ async def _send_challenges_to_timelords(
                 delivery,
             )
 
-    async def _on_connect(self) -> AsyncGenerator[OutboundMessage, None]:
+    async def _on_connect(self) -> OutboundMessageGenerator:
         """
         Whenever we connect to another node, send them our current heads. Also send heads to farmers
         and challenges to timelords.
@@ -162,7 +163,7 @@ async def _on_connect(self) -> AsyncGenerator[OutboundMessage, None]:
         async for msg in self._send_tips_to_farmers(Delivery.RESPOND):
             yield msg
 
-    def _num_needed_peers(self):
+    def _num_needed_peers(self) -> int:
         diff = self.config["target_peer_count"] - len(
             self.server.global_connections.get_full_node_connections()
         )
@@ -177,7 +178,7 @@ def _start_bg_tasks(self):
         introducer_peerinfo = PeerInfo(introducer["host"], introducer["port"])
 
         async def introducer_client():
-            async def on_connect():
+            async def on_connect() -> OutboundMessageGenerator:
                 msg = Message("request_peers", peer_protocol.RequestPeers())
                 yield OutboundMessage(NodeType.INTRODUCER, msg, Delivery.RESPOND)
 
@@ -196,7 +197,7 @@ async def on_connect():
     def _shutdown(self):
         self._shut_down = True
 
-    async def _sync(self):
+    async def _sync(self) -> OutboundMessageGenerator:
         """
         Performs a full sync of the blockchain.
             - Check which are the heaviest tips
@@ -222,13 +223,20 @@ async def _sync(self):
                 Tuple[bytes32, FullBlock]
             ] = await self.store.get_potential_tips_tuples()
             log.info(f"Have collected {len(potential_tips)} potential tips")
-            for header_hash, block in potential_tips:
-                if block.header_block.challenge is None:
-                    raise ValueError(f"Invalid tip block {block.header_hash} received")
-                if block.header_block.challenge.total_weight > highest_weight:
-                    highest_weight = block.header_block.challenge.total_weight
-                    tip_block = block
-                    tip_height = block.header_block.challenge.height
+            for header_hash, potential_tip_block in potential_tips:
+                if potential_tip_block.header_block.challenge is None:
+                    raise ValueError(
+                        f"Invalid tip block {potential_tip_block.header_hash} received"
+                    )
+                if (
+                    potential_tip_block.header_block.challenge.total_weight
+                    > highest_weight
+                ):
+                    highest_weight = (
+                        potential_tip_block.header_block.challenge.total_weight
+                    )
+                    tip_block = potential_tip_block
+                    tip_height = potential_tip_block.header_block.challenge.height
             if highest_weight <= max(
                 [t.weight for t in self.blockchain.get_current_tips()]
             ):
@@ -474,7 +482,9 @@ async def _sync(self):
             for height in range(height_checkpoint, end_height):
                 if self._shut_down:
                     return
-                block = await self.store.get_potential_block(uint32(height))
+                block: Optional[FullBlock] = await self.store.get_potential_block(
+                    uint32(height)
+                )
                 assert block is not None
                 start = time.time()
                 async with self.store.lock:
@@ -498,7 +508,7 @@ async def _sync(self):
         assert max([h.height for h in self.blockchain.get_current_tips()]) == tip_height
         log.info(f"Finished sync up to height {tip_height}")
 
-    async def _finish_sync(self):
+    async def _finish_sync(self) -> OutboundMessageGenerator:
         """
         Finalize sync by setting sync mode to False, clearing all sync information, and adding any final
         blocks that we have finalized recently.
@@ -525,7 +535,7 @@ async def _finish_sync(self):
     @api_request
     async def request_all_header_hashes(
         self, request: peer_protocol.RequestAllHeaderHashes
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         try:
             header_hashes = self.blockchain.get_header_hashes(request.tip_header_hash)
             message = Message(
@@ -538,7 +548,7 @@ async def request_all_header_hashes(
     @api_request
     async def all_header_hashes(
         self, all_header_hashes: peer_protocol.AllHeaderHashes
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         assert len(all_header_hashes.header_hashes) > 0
         async with self.store.lock:
             self.store.set_potential_hashes(all_header_hashes.header_hashes)
@@ -549,7 +559,7 @@ async def all_header_hashes(
     @api_request
     async def request_header_blocks(
         self, request: peer_protocol.RequestHeaderBlocks
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         A peer requests a list of header blocks, by height. Used for syncing or light clients.
         """
@@ -579,7 +589,7 @@ async def request_header_blocks(
     @api_request
     async def header_blocks(
         self, request: peer_protocol.HeaderBlocks
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         Receive header blocks from a peer.
         """
@@ -597,7 +607,7 @@ async def header_blocks(
     @api_request
     async def request_sync_blocks(
         self, request: peer_protocol.RequestSyncBlocks
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         Responsd to a peers request for syncing blocks.
         """
@@ -642,7 +652,7 @@ async def request_sync_blocks(
     @api_request
     async def sync_blocks(
         self, request: peer_protocol.SyncBlocks
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         We have received the blocks that we needed for syncing. Add them to processing queue.
         """
@@ -662,7 +672,7 @@ async def sync_blocks(
     @api_request
     async def request_header_hash(
         self, request: farmer_protocol.RequestHeaderHash
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         Creates a block body and header, with the proof of space, coinbase, and fee targets provided
         by the farmer, and sends the hash of the header data back to the farmer.
@@ -747,7 +757,7 @@ async def request_header_hash(
     @api_request
     async def header_signature(
         self, header_signature: farmer_protocol.HeaderSignature
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         Signature of header hash, by the harvester. This is enough to create an unfinished
         block, which only needs a Proof of Time to be finished. If the signature is valid,
@@ -783,7 +793,7 @@ async def header_signature(
     @api_request
     async def proof_of_time_finished(
         self, request: timelord_protocol.ProofOfTimeFinished
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         A proof of time, received by a peer timelord. We can use this to complete a block,
         and call the block routine (which handles propagation and verification of blocks).
@@ -846,7 +856,7 @@ async def proof_of_time_finished(
     @api_request
     async def new_proof_of_time(
         self, new_proof_of_time: peer_protocol.NewProofOfTime
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         A proof of time, received by a peer full node. If we have the rest of the block,
         we can complete it. Otherwise, we just verify and propagate the proof.
@@ -878,7 +888,7 @@ async def new_proof_of_time(
     @api_request
     async def unfinished_block(
         self, unfinished_block: peer_protocol.UnfinishedBlock
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         """
         We have received an unfinished block, either created by us, or from another peer.
         We can validate it and if it's a good block, propagate it to other peers and
@@ -930,7 +940,7 @@ async def unfinished_block(
             leader: Tuple[uint32, uint64] = self.store.get_unfinished_block_leader()
             if leader is None or unfinished_block.block.height > leader[0]:
                 log.info(
-                    f"This is the first block at height {unfinished_block.block.height}, so propagate."
+                    f"This is the first unfinished block at height {unfinished_block.block.height}, so propagate."
                 )
                 # If this is the first block we see at this height, propagate
                 self.store.set_unfinished_block_leader(
@@ -974,9 +984,7 @@ async def unfinished_block(
         )
 
     @api_request
-    async def block(
-        self, block: peer_protocol.Block
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    async def block(self, block: peer_protocol.Block) -> OutboundMessageGenerator:
         """
         Receive a full block from a peer full node (or ourselves).
         """
@@ -1020,15 +1028,15 @@ async def block(
                 )
                 try:
                     # Performs sync, and catch exceptions so we don't close the connection
-                    async for msg in self._sync():
-                        yield msg
+                    async for ret_msg in self._sync():
+                        yield ret_msg
                 except asyncio.CancelledError:
                     log.warning("Syncing failed, CancelledError")
                 except BaseException as e:
                     log.warning(f"Error {type(e)}{e} with syncing")
                 finally:
-                    async for msg in self._finish_sync():
-                        yield msg
+                    async for ret_msg in self._finish_sync():
+                        yield ret_msg
 
             elif block.block.height >= tip_height - 3:
                 log.info(
@@ -1062,7 +1070,7 @@ async def block(
                     ips_changed = True
             if ips_changed:
                 rate_update = farmer_protocol.ProofOfTimeRate(next_vdf_ips)
-                log.error(f"Sending proof of time rate {next_vdf_ips}")
+                log.info(f"Sending proof of time rate {next_vdf_ips}")
                 yield OutboundMessage(
                     NodeType.FARMER,
                     Message("proof_of_time_rate", rate_update),
@@ -1123,8 +1131,8 @@ async def block(
                 FullBlock
             ] = await self.store.get_disconnected_block_by_prev(block.block.header_hash)
         if next_block is not None:
-            async for msg in self.block(peer_protocol.Block(next_block)):
-                yield msg
+            async for ret_msg in self.block(peer_protocol.Block(next_block)):
+                yield ret_msg
 
         async with self.store.lock:
             # Removes all temporary data for old blocks
@@ -1137,7 +1145,7 @@ async def block(
     @api_request
     async def request_block(
         self, request_block: peer_protocol.RequestBlock
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    ) -> OutboundMessageGenerator:
         block: Optional[FullBlock] = await self.store.get_block(
             request_block.header_hash
         )
@@ -1149,9 +1157,7 @@ async def request_block(
             )
 
     @api_request
-    async def peers(
-        self, request: peer_protocol.Peers
-    ) -> AsyncGenerator[OutboundMessage, None]:
+    async def peers(self, request: peer_protocol.Peers) -> OutboundMessageGenerator:
         conns = self.server.global_connections
         for peer in request.peer_list:
             conns.peers.add(peer)

src/introducer.py

@@ -13,6 +13,7 @@
 import asyncio
 
 import logging
+
 log = logging.getLogger(__name__)
 
 
@@ -31,27 +32,23 @@ async def request_peers(
     ) -> AsyncGenerator[OutboundMessage, None]:
         max_peers = self.config["max_peers_to_send"]
         rawpeers = self.server.global_connections.peers.get_peers(
-            max_peers*2, True, self.config["recent_peer_threshold"]
+            max_peers * 2, True, self.config["recent_peer_threshold"]
         )
 
         peers = []
 
         for peer in rawpeers:
             if peer.get_hash() not in self.vetted:
                 try:
-                    r, w = await asyncio.open_connection(
-                        peer.host, int(peer.port)
-                    )
+                    r, w = await asyncio.open_connection(peer.host, int(peer.port))
                     w.close()
                 except (
                     ConnectionRefusedError,
                     TimeoutError,
                     OSError,
                     asyncio.TimeoutError,
                 ) as e:
-                    log.warning(
-                        f"Could not vet {peer}. {type(e)}{str(e)}"
-                    )
+                    log.warning(f"Could not vet {peer}. {type(e)}{str(e)}")
                     self.vetted[peer.get_hash()] = False
                     continue
 

src/server/connection.py

@@ -47,7 +47,6 @@ def __init__(
         self.on_connect = on_connect
 
         # Connection metrics
-        self.handshake_finished = False
         self.creation_type = time.time()
         self.bytes_read = 0
         self.bytes_written = 0