diff --git a/magneticod/magneticod/__main__.py b/magneticod/magneticod/__main__.py index 8295536..0d412ae 100644 --- a/magneticod/magneticod/__main__.py +++ b/magneticod/magneticod/__main__.py @@ -51,7 +51,7 @@ def main(): node = dht.SybilNode(arguments.node_addr, complete_info_hashes, arguments.max_metadata_size) loop.run_until_complete(node.launch(loop)) - loop.create_task(watch_q(database, node._metadata_q)) + watch_q_task = loop.create_task(watch_q(database, node._metadata_q)) try: loop.run_forever() @@ -59,7 +59,9 @@ def main(): logging.critical("Keyboard interrupt received! Exiting gracefully...") finally: database.close() + watch_q_task.cancel() loop.run_until_complete(node.shutdown()) + loop.run_until_complete(watch_q_task) return 0 diff --git a/magneticod/magneticod/dht.py b/magneticod/magneticod/dht.py index bdd1f54..a39ef4d 100644 --- a/magneticod/magneticod/dht.py +++ b/magneticod/magneticod/dht.py @@ -105,15 +105,11 @@ class SybilNode: self.__n_max_neighbours = self.__n_max_neighbours * 101 // 100 async def shutdown(self) -> None: - peers = [peer for peer in itertools.chain.from_iterable(self.__peers.values())] - for peer in peers: - peer.cancel() - for peer in peers: - await peer - for task in self._tasks: - task.cancel() - for task in self._tasks: - await task + futures = [peer for peer in itertools.chain.from_iterable(self.__peers.values())] + futures.extend(self._tasks) + for future in futures: + future.cancel() + await asyncio.wait(futures) self._transport.close() def __on_FIND_NODE_response(self, message: bencode.KRPCDict) -> None: