diff --git a/squeaknode/network/connection.py b/squeaknode/network/connection.py index fdb15c3a..dc4873bc 100644 --- a/squeaknode/network/connection.py +++ b/squeaknode/network/connection.py @@ -81,15 +81,20 @@ class Connection(object): finally: logger.debug("Removing peer.") connection_manager.remove_peer(self.peer) - self.peer.stop() - # TODO: Set a stop event here. def shutdown(self): + logger.debug("Peet shutting down...") self.peer.stop() + self.ping_timer.cancel() + self.pong_timer.cancel() def handle_connection(self): - self.initial_sync() - self.handle_msgs() + try: + self.initial_sync() + self.handle_msgs() + finally: + self.shutdown() + # self._stopped.set() def initial_sync(self): self.send_ping() @@ -344,6 +349,12 @@ class PingTimer: self.peer_name) self.timer.start() + def cancel(self): + logger.debug("Cancelling ping timer.") + with self._lock: + if self.timer: + self.timer.cancel() + def send_ping(self): logger.debug("Sending ping triggered by timer.") self.send_fn() @@ -398,6 +409,12 @@ class PongTimer: # Start a new ping timer. self.start_ping_timer() + def cancel(self): + logger.debug("Cancelling pong timer.") + with self._lock: + if self.timer: + self.timer.cancel() + def shutdown(self): logger.debug("Shutdown connection triggered by pong timer.") self.shutdown_fn() diff --git a/squeaknode/network/peer_handler.py b/squeaknode/network/peer_handler.py index 02b114c7..2477dd73 100644 --- a/squeaknode/network/peer_handler.py +++ b/squeaknode/network/peer_handler.py @@ -100,18 +100,10 @@ class PeerHandler(): def start_connection(self, peer: Peer): """Start a connection """ - logger.debug( - 'Setting up connection for peer {}'.format(peer)) - try: - with Connection(peer, self.squeak_controller).connect( - self.connection_manager - ) as connection: - connection.handle_connection() - finally: - peer.stop() - logger.debug( - 'Stopped connection for peer {}.'.format(peer), - ) + with Connection(peer, self.squeak_controller).connect( + self.connection_manager + ) as connection: + connection.handle_connection() class HandshakeTimer: