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

@@ -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)

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):

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()
 
 

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,

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)

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:

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,
     )