Fix shutdown threadpool in download manager (#1727)

* Fix shutdown threadpool in download manager

* Call event.wait with timeout value so it stops eventually

* Cancel all running downloads when download manager shuts down

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:

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

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

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

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