From 1fb353b5eb3b0dbfd9b2f8190d9e112589230cb5 Mon Sep 17 00:00:00 2001 From: Jonathan Zernik Date: Wed, 11 Aug 2021 15:39:27 -0700 Subject: [PATCH] Fix remove peer on disconnect (#912) --- squeaknode/network/connection.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/squeaknode/network/connection.py b/squeaknode/network/connection.py index 15369169..5dea05d4 100644 --- a/squeaknode/network/connection.py +++ b/squeaknode/network/connection.py @@ -86,12 +86,14 @@ class Connection(): def open_connection(self): logger.debug( 'Starting handshake connection with peer ... {}'.format(self.peer)) - self.connection_manager.add_peer(self.peer) - self.handshake() - logger.debug('Peer connection added... {}'.format(self.peer)) - yield self - self.connection_manager.remove_peer(self.peer) - logger.debug('Peer connection removed... {}'.format(self.peer)) + try: + self.connection_manager.add_peer(self.peer) + self.handshake() + logger.debug('Peer connection added... {}'.format(self.peer)) + yield self + finally: + self.connection_manager.remove_peer(self.peer) + logger.debug('Peer connection removed... {}'.format(self.peer)) # def __enter__(self): # logger.debug(