diff --git a/squeaknode/network/connection.py b/squeaknode/network/connection.py index dc4873bc..7d685100 100644 --- a/squeaknode/network/connection.py +++ b/squeaknode/network/connection.py @@ -21,7 +21,6 @@ # SOFTWARE. import logging import threading -from contextlib import contextmanager from squeak.messages import msg_addr from squeak.messages import msg_getaddr @@ -69,19 +68,6 @@ class Connection(object): str(self.peer), ) - @contextmanager - def connect(self, connection_manager): - logger.debug("Adding peer.") - connection_manager.add_peer(self.peer) - try: - logger.debug("Yielding peer.") - yield self - except Exception: - logger.exception("Peer connection failed.") - finally: - logger.debug("Removing peer.") - connection_manager.remove_peer(self.peer) - def shutdown(self): logger.debug("Peet shutting down...") self.peer.stop() diff --git a/squeaknode/network/connection_manager.py b/squeaknode/network/connection_manager.py index 1a41407c..5ed27df8 100644 --- a/squeaknode/network/connection_manager.py +++ b/squeaknode/network/connection_manager.py @@ -21,11 +21,13 @@ # SOFTWARE. import logging import threading +from contextlib import contextmanager from typing import Dict from typing import List from typing import Optional from squeaknode.core.peer_address import PeerAddress +from squeaknode.network.connection import Connection from squeaknode.network.peer import Peer from squeaknode.node.listener_subscription_client import EventListener @@ -49,6 +51,19 @@ class ConnectionManager(object): self.single_peer_changed_listener = EventListener() self.accept_connections = True + @contextmanager + def connect(self, peer, squeak_controller): + logger.debug("Adding peer.") + self.add_peer(peer) + try: + logger.debug("Yielding connection.") + yield Connection(peer, squeak_controller) + except Exception: + logger.exception("Peer connection failed.") + finally: + logger.debug("Removing peer.") + self.remove_peer(peer) + @property def peers(self) -> List[Peer]: return list(self._peers.values()) diff --git a/squeaknode/network/peer_handler.py b/squeaknode/network/peer_handler.py index 2477dd73..26267b08 100644 --- a/squeaknode/network/peer_handler.py +++ b/squeaknode/network/peer_handler.py @@ -24,7 +24,6 @@ import socket import threading from squeaknode.core.peer_address import PeerAddress -from squeaknode.network.connection import Connection from squeaknode.network.peer import Peer @@ -100,9 +99,7 @@ class PeerHandler(): def start_connection(self, peer: Peer): """Start a connection """ - with Connection(peer, self.squeak_controller).connect( - self.connection_manager - ) as connection: + with self.connection_manager.connect(peer, self.squeak_controller) as connection: connection.handle_connection()