diff --git a/squeaknode/node/active_download_manager.py b/squeaknode/node/active_download_manager.py index 812d5922..bbdedb98 100644 --- a/squeaknode/node/active_download_manager.py +++ b/squeaknode/node/active_download_manager.py @@ -71,8 +71,11 @@ class ActiveDownload(ABC): def mark_complete(self): self.stopped.set() - def wait_for_complete(self) -> None: - self.stopped.wait() + def cancel(self): + self.stopped.set() + + def wait_for_complete(self, timeout_s: int) -> None: + self.stopped.wait(timeout=timeout_s) def get_result(self) -> DownloadResult: return DownloadResult( @@ -122,12 +125,22 @@ class HashDownload(ActiveDownload): class ActiveDownloadManager: - def __init__(self, broadcast_fn): + def __init__(self): self.downloads: Dict[str, ActiveDownload] = dict() - # self.interests: Dict[str, ActiveDownload] = ExpiringDict( - # max_len=100, max_age_seconds=10) - self.executor = ThreadPoolExecutor(max_workers=10) + self.executor = None + self.broadcast_fn = None + + def start(self, broadcast_fn): + logger.info("Starting Download Manager...") self.broadcast_fn = broadcast_fn + self.executor = ThreadPoolExecutor(max_workers=10) + + def stop(self): + for download in self.downloads.values(): + download.cancel() + logger.info("Stopping Download Manager...") + self.executor.shutdown(wait=True) + logger.info("Stopped Download Manager.") def lookup_counter(self, squeak: CSqueak) -> Optional[ActiveDownload]: for name, interest in self.downloads.items(): @@ -140,7 +153,7 @@ class ActiveDownloadManager: self.downloads[name_key] = download future = self.executor.submit(self.download_task, download) try: - return future.result(DOWNLOAD_TIMEOUT_S) + return future.result() except TimeoutError: return download.get_result() finally: @@ -148,7 +161,7 @@ class ActiveDownloadManager: def download_task(self, download: ActiveDownload) -> DownloadResult: download.initiate_download(self.broadcast_fn) - download.wait_for_complete() + download.wait_for_complete(DOWNLOAD_TIMEOUT_S) return download.get_result() def download_interest(self, limit: int, interest: CInterested) -> DownloadResult: diff --git a/squeaknode/node/squeak_controller.py b/squeaknode/node/squeak_controller.py index 1463cc7b..991386c1 100644 --- a/squeaknode/node/squeak_controller.py +++ b/squeaknode/node/squeak_controller.py @@ -60,7 +60,6 @@ from squeaknode.core.squeak_profile import SqueakProfile from squeaknode.core.squeaks import get_hash from squeaknode.core.update_subscriptions_event import UpdateSubscriptionsEvent from squeaknode.node.active_download_manager import ActiveDownload -from squeaknode.node.active_download_manager import ActiveDownloadManager from squeaknode.node.listener_subscription_client import EventListener from squeaknode.node.received_payments_subscription_client import ReceivedPaymentsSubscriptionClient # from squeaknode.node.temporary_interest_manager import TemporaryInterest @@ -78,6 +77,7 @@ class SqueakController: squeak_core, payment_processor, network_manager, + download_manager, config, ): self.squeak_db = squeak_db @@ -89,8 +89,7 @@ class SqueakController: self.new_secret_key_listener = EventListener() self.new_follow_listener = EventListener() # self.temporary_interest_manager = TemporaryInterestManager() - self.active_download_manager = ActiveDownloadManager( - self.broadcast_msg) + self.active_download_manager = download_manager self.config = config def save_squeak(self, squeak: CSqueak) -> Optional[bytes]: diff --git a/squeaknode/node/squeak_node.py b/squeaknode/node/squeak_node.py index 0aef8043..85c1915a 100644 --- a/squeaknode/node/squeak_node.py +++ b/squeaknode/node/squeak_node.py @@ -35,6 +35,7 @@ from squeaknode.db.db_engine import get_engine from squeaknode.db.squeak_db import SqueakDb from squeaknode.lightning.lnd_lightning_client import LNDLightningClient from squeaknode.network.network_manager import NetworkManager +from squeaknode.node.active_download_manager import ActiveDownloadManager from squeaknode.node.payment_processor import PaymentProcessor from squeaknode.node.peer_connection_worker import PeerConnectionWorker from squeaknode.node.peer_subscription_update_worker import PeerSubscriptionUpdateWorker @@ -63,6 +64,7 @@ class SqueakNode: self.initialize_squeak_core() self.initialize_payment_processor() self.initialize_network_manager() + self.initialize_download_manager() self.initialize_squeak_controller() self.initialize_admin_handler() self.initialize_admin_rpc_server() @@ -92,8 +94,12 @@ class SqueakNode: self.new_secret_key_worker.start_running() self.new_follow_worker.start_running() self.new_bitcoin_block_worker.start_running() + self.download_manager.start( + self.squeak_controller.broadcast_msg, + ) def stop_running(self): + self.download_manager.stop() self.admin_web_server.stop() self.admin_rpc_server.stop() self.network_manager.stop() @@ -165,6 +171,7 @@ class SqueakNode: self.squeak_core, self.payment_processor, self.network_manager, + self.download_manager, self.config, ) @@ -236,3 +243,6 @@ class SqueakNode: self.squeak_controller, self.bitcoin_block_subscription_client, ) + + def initialize_download_manager(self): + self.download_manager = ActiveDownloadManager() diff --git a/tests/node/test_active_download.py b/tests/node/test_active_download.py index 07b62c56..a122d8a9 100644 --- a/tests/node/test_active_download.py +++ b/tests/node/test_active_download.py @@ -105,7 +105,7 @@ def test_download_hash_mark_complete_not_called(download_hash, squeak): def test_download_hash_wait_for_complete(download_hash): download_hash.mark_complete() - download_hash.wait_for_complete() + download_hash.wait_for_complete(50) def test_download_hash_get_result(download_hash, squeak): diff --git a/tests/node/test_squeak_controller.py b/tests/node/test_squeak_controller.py index bf17e883..bc973700 100644 --- a/tests/node/test_squeak_controller.py +++ b/tests/node/test_squeak_controller.py @@ -30,6 +30,7 @@ from squeaknode.core.squeak_core import SqueakCore from squeaknode.core.squeak_peer import SqueakPeer from squeaknode.db.squeak_db import SqueakDb from squeaknode.network.network_manager import NetworkManager +from squeaknode.node.active_download_manager import ActiveDownloadManager from squeaknode.node.payment_processor import PaymentProcessor from squeaknode.node.squeak_controller import SqueakController @@ -99,12 +100,18 @@ def payment_processor(): return mock.Mock(spec=PaymentProcessor) +@pytest.fixture +def download_manager(): + return mock.Mock(spec=ActiveDownloadManager) + + @pytest.fixture def squeak_controller( squeak_db, squeak_core, payment_processor, network_manager, + download_manager, config, ): return SqueakController( @@ -112,6 +119,7 @@ def squeak_controller( squeak_core, payment_processor, network_manager, + download_manager, config, ) @@ -122,6 +130,7 @@ def regtest_squeak_controller( squeak_core, payment_processor, network_manager, + download_manager, regtest_config, ): return SqueakController( @@ -129,6 +138,7 @@ def regtest_squeak_controller( squeak_core, payment_processor, network_manager, + download_manager, regtest_config, )