From 4cc617a94cd17505dffadaf22f1398002f71d202 Mon Sep 17 00:00:00 2001 From: Jonathan Zernik Date: Sun, 15 Aug 2021 22:15:42 -0700 Subject: [PATCH] Refactor network commands into network manager (#949) * Refactor network commands into network manager * Start and stop peer server and peer client inside network manager * Pass handle connection fn to peer handler instead of self object * Remove old comments --- squeaknode/network/network_manager.py | 120 ++++++++++++++++++++++++++ squeaknode/network/peer_client.py | 2 +- squeaknode/network/peer_handler.py | 34 ++------ squeaknode/network/peer_server.py | 7 +- squeaknode/node/squeak_controller.py | 71 +++------------ squeaknode/node/squeak_node.py | 25 ++---- tests/node/test_squeak_controller.py | 34 ++------ 7 files changed, 154 insertions(+), 139 deletions(-) create mode 100644 squeaknode/network/network_manager.py diff --git a/squeaknode/network/network_manager.py b/squeaknode/network/network_manager.py new file mode 100644 index 00000000..cece9a41 --- /dev/null +++ b/squeaknode/network/network_manager.py @@ -0,0 +1,120 @@ +import logging +import socket +from typing import List + +import squeak.params +from squeak.messages import MsgSerializable + +from squeaknode.core.peer_address import PeerAddress +from squeaknode.core.squeak_peer import SqueakPeer +from squeaknode.network.connection_manager import ConnectionManager +from squeaknode.network.peer import Peer +from squeaknode.network.peer_client import PeerClient +from squeaknode.network.peer_handler import PeerHandler +from squeaknode.network.peer_server import PeerServer + + +MIN_PEERS = 5 +MAX_PEERS = 10 +UPDATE_THREAD_SLEEP_TIME = 10 + + +logger = logging.getLogger(__name__) + + +class NetworkManager(object): + """Interface for doing things involving the network. + """ + + def __init__(self, config): + self.config = config + self.peer_server = None + self.peer_client = None + self.connection_manager = ConnectionManager() + + def start(self, squeak_controller): + peer_handler = PeerHandler( + squeak_controller, + self.handle_connection, + ) + self.peer_server = PeerServer( + peer_handler, + self.config.server.rpc_port, + ) + self.peer_client = PeerClient( + peer_handler, + ) + self.peer_server.start() + + def stop(self): + self.peer_server.stop() + self.connection_manager.stop_all_connections() + + def connect_peer(self, host: str, port: int) -> None: + port = port or squeak.params.params.DEFAULT_PORT + peer_address = PeerAddress(host=host, port=port) + self.peer_client.connect_address(peer_address) + + def connect_peers(self, peers: List[SqueakPeer]) -> None: + connected_peers = self.connection_manager.peers + connected_peer_addrs = [ + peer.address + for peer in connected_peers + ] + for peer in peers: + peer_addr = (peer.address.host, peer.address.port) + if peer_addr not in connected_peer_addrs: + logger.info("Connect to peer: {}".format( + peer, + )) + self.peer_client.connect_address(peer.address) + + def disconnect_peer(self, host: str, port: int) -> None: + peer_address = PeerAddress(host=host, port=port) + self.connection_manager.stop_connection(peer_address) + + def get_connected_peer(self, host, port): + address = (host, port) + return self.connection_manager.get_peer(address) + + def get_connected_peers(self): + return self.connection_manager.peers + + def broadcast_msg(self, msg: MsgSerializable) -> None: + for peer in self.connection_manager.peers: + try: + peer.send_msg(msg) + except Exception: + logger.exception("Failed to send msg to peer: {}".format( + peer, + )) + + def handle_connection(self, squeak_controller, peer_socket, address, outgoing): + """Handles all sending and receiving of messages for the given peer. + + This method blocks until the peer connection has stopped. + """ + if self.connection_manager.has_connection(address): + return + + logger.debug( + 'Setting up controller for peer address {} ...'.format(address)) + with Peer(peer_socket, address, outgoing).open_connection(squeak_controller) as peer: + self.connection_manager.add_peer(peer) + try: + peer.handle_messages(squeak_controller) + except Exception: + logger.exception("Handling messages failed.") + finally: + self.connection_manager.remove_peer(peer) + logger.debug('Stopped controller for peer address {}.'.format(address)) + + def get_address(self): + # TODO: Add return type. + return (self.peer_server.ip, self.peer_server.port) + + def get_remote_address(self, address): + # TODO: Add return type. + hostname, port = address + ip = socket.gethostbyname(hostname) + return (ip, port) diff --git a/squeaknode/network/peer_client.py b/squeaknode/network/peer_client.py index 8b7d7426..d7f549d8 100644 --- a/squeaknode/network/peer_client.py +++ b/squeaknode/network/peer_client.py @@ -15,7 +15,7 @@ class PeerClient(object): """Creates outgoing connections to other peers in the network. """ - def start(self, peer_handler): + def __init__(self, peer_handler): self.peer_handler = peer_handler def make_connection(self, address): diff --git a/squeaknode/network/peer_handler.py b/squeaknode/network/peer_handler.py index cf1ec440..a977a7e2 100644 --- a/squeaknode/network/peer_handler.py +++ b/squeaknode/network/peer_handler.py @@ -1,10 +1,6 @@ import logging import threading -from squeaknode.network.connection_manager import ConnectionManager -from squeaknode.network.peer import Peer -from squeaknode.node.squeak_controller import SqueakController - logger = logging.getLogger(__name__) @@ -15,37 +11,17 @@ class PeerHandler(): def __init__( self, - squeak_controller: SqueakController, - connection_manager: ConnectionManager, + squeak_controller, + handle_connection_fn, ): super().__init__() self.squeak_controller = squeak_controller - self.connection_manager = connection_manager - - def start(self, peer_socket, address, outgoing): - """Handles all sending and receiving of messages for the given peer. - - This method blocks until the peer connection has stopped. - """ - if self.connection_manager.has_connection(address): - return - - logger.debug( - 'Setting up controller for peer address {} ...'.format(address)) - with Peer(peer_socket, address, outgoing).open_connection(self.squeak_controller) as peer: - self.connection_manager.add_peer(peer) - try: - peer.handle_messages(self.squeak_controller) - except Exception: - logger.exception("Handling messages failed.") - finally: - self.connection_manager.remove_peer(peer) - logger.debug('Stopped controller for peer address {}.'.format(address)) + self.handle_connection_fn = handle_connection_fn def handle_connection(self, peer_socket, address, outgoing): threading.Thread( - target=self.start, - args=(peer_socket, address, outgoing,), + target=self.handle_connection_fn, + args=(self.squeak_controller, peer_socket, address, outgoing,), ).start() diff --git a/squeaknode/network/peer_server.py b/squeaknode/network/peer_server.py index 55438d84..36901696 100644 --- a/squeaknode/network/peer_server.py +++ b/squeaknode/network/peer_server.py @@ -17,17 +17,16 @@ class PeerServer(object): """Maintains connections to other peers in the network. """ - def __init__(self, port=None): + def __init__(self, peer_handler, port=None): + self.peer_handler = peer_handler self.ip = socket.gethostbyname('localhost') self.port = port or squeak.params.params.DEFAULT_PORT self.listen_socket = socket.socket() - def start(self, peer_handler): - self.peer_handler = peer_handler + def start(self): logger.info("Starting peer server with port: {}".format( self.port, )) - # Start Listen thread threading.Thread( target=self.accept_connections, diff --git a/squeaknode/node/squeak_controller.py b/squeaknode/node/squeak_controller.py index 818a7353..678f9ff2 100644 --- a/squeaknode/node/squeak_controller.py +++ b/squeaknode/node/squeak_controller.py @@ -1,5 +1,4 @@ import logging -import socket import threading from typing import List from typing import Optional @@ -31,9 +30,6 @@ from squeaknode.core.squeak_peer import SqueakPeer from squeaknode.core.squeak_profile import SqueakProfile from squeaknode.core.util import get_hash from squeaknode.core.util import is_address_valid -from squeaknode.network.connection_manager import ConnectionManager -from squeaknode.network.peer_client import PeerClient -from squeaknode.network.peer_server import PeerServer from squeaknode.node.received_payments_subscription_client import ReceivedPaymentsSubscriptionClient @@ -48,18 +44,14 @@ class SqueakController: squeak_core, squeak_rate_limiter, payment_processor, - peer_server: PeerServer, - peer_client: PeerClient, - connection_manager: ConnectionManager, + network_manager, config, ): self.squeak_db = squeak_db self.squeak_core = squeak_core self.squeak_rate_limiter = squeak_rate_limiter self.payment_processor = payment_processor - self.peer_server = peer_server - self.peer_client = peer_client - self.connection_manager = connection_manager + self.network_manager = network_manager self.config = config def save_squeak( @@ -487,51 +479,27 @@ class SqueakController: ) def connect_peer(self, host: str, port: int) -> None: - # peer = self.squeak_db.get_peer(peer_id) - # if peer is None: - # raise Exception("Peer with id {} not found.".format( - # peer_id, - # )) - # # TODO logger.info("Connect to peer: {}:{}".format( host, port, )) - port = port or squeak.params.params.DEFAULT_PORT - peer_address = PeerAddress(host=host, port=port) - self.peer_client.connect_address(peer_address) + self.network_manager.connect_peer(host, port) def connect_peers(self) -> None: peers = self.squeak_db.get_peers() - connected_peers = self.connection_manager.peers - connected_peer_addrs = [ - peer.address - for peer in connected_peers - ] - for peer in peers: - peer_addr = (peer.address.host, peer.address.port) - if peer_addr not in connected_peer_addrs: - logger.info("Connect to peer: {}".format( - peer, - )) - self.peer_client.connect_address(peer.address) + self.network_manager.connect_peers(peers) def get_address(self): - # TODO: Add return type. - return (self.peer_server.ip, self.peer_server.port) + return self.network_manager.get_address() def get_remote_address(self, address): - # TODO: Add return type. - hostname, port = address - ip = socket.gethostbyname(hostname) - return (ip, port) - - def get_connected_peers(self): - return self.connection_manager.peers + return self.network_manager.get_remote_address(address) def get_connected_peer(self, host, port): - address = (host, port) - return self.connection_manager.get_peer(address) + return self.network_manager.get_connected_peer(host, port) + + def get_connected_peers(self): + return self.network_manager.get_connected_peers() def lookup_squeaks_for_interest( self, @@ -632,26 +600,11 @@ class SqueakController: return ret def broadcast_msg(self, msg: MsgSerializable) -> None: - for peer in self.connection_manager.peers: - try: - peer.send_msg(msg) - except Exception: - logger.exception("Failed to send msg to peer: {}".format( - peer, - )) + self.network_manager.broadcast_msg(msg) def disconnect_peer(self, host: str, port: int) -> None: - # peer = self.squeak_db.get_peer(peer_id) - # if peer is None: - # raise Exception("Peer with id {} not found.".format( - # peer_id, - # )) - # logger.info("Disconnect peer: {}".format( - # peer, - # )) logger.info("Disconnect to peer: {}:{}".format( host, port, )) - peer_address = PeerAddress(host=host, port=port) - self.connection_manager.stop_connection(peer_address) + self.network_manager.disconnect_peer(host, port) diff --git a/squeaknode/node/squeak_node.py b/squeaknode/node/squeak_node.py index 77fc6475..de30a5fb 100644 --- a/squeaknode/node/squeak_node.py +++ b/squeaknode/node/squeak_node.py @@ -14,10 +14,7 @@ from squeaknode.db.db_engine import get_engine from squeaknode.db.db_engine import get_sqlite_connection_string from squeaknode.db.squeak_db import SqueakDb from squeaknode.lightning.lnd_lightning_client import LNDLightningClient -from squeaknode.network.connection_manager import ConnectionManager -from squeaknode.network.peer_client import PeerClient -from squeaknode.network.peer_handler import PeerHandler -from squeaknode.network.peer_server import PeerServer +from squeaknode.network.network_manager import NetworkManager from squeaknode.node.payment_processor import PaymentProcessor from squeaknode.node.peer_connection_worker import PeerConnectionWorker from squeaknode.node.process_received_payments_worker import ProcessReceivedPaymentsWorker @@ -71,29 +68,21 @@ class SqueakNode: self.config.core.subscribe_invoices_retry_s, ) - self.connection_manager = ConnectionManager() - self.peer_server = PeerServer(self.config.server.rpc_port) - self.peer_client = PeerClient() + self.network_manager = NetworkManager(self.config) squeak_controller = SqueakController( squeak_db, squeak_core, squeak_rate_limiter, payment_processor, - self.peer_server, - self.peer_client, - self.connection_manager, + self.network_manager, self.config, ) + self.squeak_controller = squeak_controller admin_handler = load_admin_handler( lightning_client, squeak_controller) - self.peer_handler = PeerHandler( - squeak_controller, - self.connection_manager, - ) - self.admin_rpc_server = load_admin_rpc_server( self.config, admin_handler, self.stopped) @@ -128,8 +117,7 @@ class SqueakNode: start_admin_web_server(self.admin_web_server) # Start peer socket server and peer client - self.peer_server.start(self.peer_handler) - self.peer_client.start(self.peer_handler) + self.network_manager.start(self.squeak_controller) # start peer connection worker if self.config.sync.enabled: @@ -148,8 +136,7 @@ class SqueakNode: self.stopped.set() # TODO: Use explicit stop to stop all components - self.peer_server.stop() - self.connection_manager.stop_all_connections() + self.network_manager.stop() def load_lightning_client(config) -> LNDLightningClient: diff --git a/tests/node/test_squeak_controller.py b/tests/node/test_squeak_controller.py index 85c60cf7..f3f02068 100644 --- a/tests/node/test_squeak_controller.py +++ b/tests/node/test_squeak_controller.py @@ -8,9 +8,7 @@ from squeaknode.core.peer_address import PeerAddress from squeaknode.core.squeak_core import SqueakCore from squeaknode.core.squeak_peer import SqueakPeer from squeaknode.db.squeak_db import SqueakDb -from squeaknode.network.connection_manager import ConnectionManager -from squeaknode.network.peer_client import PeerClient -from squeaknode.network.peer_server import PeerServer +from squeaknode.network.network_manager import NetworkManager from squeaknode.node.payment_processor import PaymentProcessor from squeaknode.node.squeak_controller import SqueakController from squeaknode.node.squeak_rate_limiter import SqueakRateLimiter @@ -39,18 +37,8 @@ def squeak_db(): @pytest.fixture -def peer_server(): - return mock.Mock(spec=PeerServer) - - -@pytest.fixture -def peer_client(): - return mock.Mock(spec=PeerClient) - - -@pytest.fixture -def connection_manager(): - return mock.Mock(spec=ConnectionManager) +def network_manager(): + return mock.Mock(spec=NetworkManager) @pytest.fixture @@ -84,9 +72,7 @@ def squeak_controller( squeak_core, squeak_rate_limiter, payment_processor, - peer_server, - peer_client, - connection_manager, + network_manager, config, ): return SqueakController( @@ -94,9 +80,7 @@ def squeak_controller( squeak_core, squeak_rate_limiter, payment_processor, - peer_server, - peer_client, - connection_manager, + network_manager, config, ) @@ -107,9 +91,7 @@ def regtest_squeak_controller( squeak_core, squeak_rate_limiter, payment_processor, - peer_server, - peer_client, - connection_manager, + network_manager, regtest_config, ): return SqueakController( @@ -117,9 +99,7 @@ def regtest_squeak_controller( squeak_core, squeak_rate_limiter, payment_processor, - peer_server, - peer_client, - connection_manager, + network_manager, regtest_config, )