diff --git a/squeaknode/client/network_controller.py b/squeaknode/client/network_controller.py index 4375cec4..66395531 100644 --- a/squeaknode/client/network_controller.py +++ b/squeaknode/client/network_controller.py @@ -24,10 +24,8 @@ from concurrent.futures import ThreadPoolExecutor from concurrent.futures import wait from typing import Optional -from squeak.core.keys import SqueakPublicKey - -from squeaknode.client.peer_downloader import RangeDownloader -from squeaknode.client.peer_downloader import SingleDownloader +from squeaknode.client.peer_downloader import PeerDownloader +from squeaknode.core.squeak_peer import SqueakPeer from squeaknode.node.squeak_store import SqueakStore logger = logging.getLogger(__name__) @@ -48,7 +46,15 @@ class NetworkController: self.proxy_host = proxy_host self.proxy_port = proxy_port - def download_timeline_async( + def get_downloader(self, peer: SqueakPeer): + return PeerDownloader( + peer, + self.squeak_store, + self.proxy_host, + self.proxy_port, + ) + + def download_timeline( self, interest_block_interval: int, ) -> None: @@ -56,49 +62,37 @@ class NetworkController: min_block = max(0, max_block - interest_block_interval) followed_public_keys = self.squeak_store.get_followed_public_keys() peers = self.squeak_store.get_autoconnect_peers() - for peer in peers: - downloader = RangeDownloader( - peer, - self.squeak_store, - self.proxy_host, - self.proxy_port, - min_block, - max_block, - followed_public_keys, - ) - downloader.download_async() - - def download_pubkey_squeaks_async(self, pubkey: SqueakPublicKey) -> None: - min_block = 0 # TODO - max_block = 999999999999 # TODO - peers = self.squeak_store.get_autoconnect_peers() - for peer in peers: - downloader = RangeDownloader( - peer, - self.squeak_store, - self.proxy_host, - self.proxy_port, - min_block, - max_block, - [pubkey], - ) - downloader.download_async() + downloaders = [ + self.get_downloader(peer) + for peer in peers + ] + with ThreadPoolExecutor(50) as executor: + # submit tasks and collect futures + futures = [ + executor.submit( + downloader.download_interest_range, + min_block, + max_block, + followed_public_keys, + ) + for downloader in downloaders + ] + # wait for all tasks to complete + wait(futures) def download_single_squeak(self, squeak_hash: bytes) -> None: peers = self.squeak_store.get_autoconnect_peers() downloaders = [ - SingleDownloader( - peer, - self.squeak_store, - self.proxy_host, - self.proxy_port, - squeak_hash, - ) for peer in peers + self.get_downloader(peer) + for peer in peers ] with ThreadPoolExecutor(50) as executor: # submit tasks and collect futures - futures = [executor.submit(downloader.download) - for downloader in downloaders] + futures = [ + executor.submit( + downloader.download_single_squeak, + squeak_hash + ) + for downloader in downloaders] # wait for all tasks to complete wait(futures) - logger.info('All downloads are done!') diff --git a/squeaknode/client/peer_downloader.py b/squeaknode/client/peer_downloader.py index 7e874c6b..5d90dc73 100644 --- a/squeaknode/client/peer_downloader.py +++ b/squeaknode/client/peer_downloader.py @@ -20,13 +20,10 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. import logging -import threading from abc import ABC -from abc import abstractmethod from typing import List from typing import Optional -from squeak.core import CSqueak from squeak.core.keys import SqueakPublicKey from squeaknode.client.peer_client import PeerClient @@ -55,59 +52,76 @@ class PeerDownloader(ABC): self.client = PeerClient(peer, proxy_host, proxy_port) self.squeak_store = squeak_store - @abstractmethod - def get_hashes(self) -> List[bytes]: - """Get list of squeak hashes to download. - """ - - @abstractmethod - def is_squeak_wanted(self, squeak: CSqueak) -> bool: - """Return true if squeak is supposed to be downloaded. - """ - - def download_async(self) -> None: - thread = threading.Thread( - target=self.download, - args=(), + def download_interest_range( + self, + min_block: int, + max_block: int, + pubkeys: List[SqueakPublicKey], + ) -> None: + squeak_hashes = self.client.lookup( + min_block, + max_block, + pubkeys, ) - thread.start() - - def download(self) -> None: - squeak_hashes = self.get_hashes() for squeak_hash in squeak_hashes: # Download the squeak if not already owned. - self.get_squeak(squeak_hash) - # Download the secret key if not already unlocked. - self.get_secret_key(squeak_hash) - # Download the offer if not already unlocked. - self.get_offer(squeak_hash) + if not self.squeak_store.get_squeak(squeak_hash): + squeak = self.client.get_squeak(squeak_hash) + if squeak and \ + squeak.nBlockHeight >= min_block and \ + squeak.nBlockHeight <= max_block and \ + squeak.GetPubKey() in pubkeys: + self.squeak_store.save_squeak(squeak) - def get_squeak(self, squeak_hash: bytes) -> None: + # Get the local squeak. + squeak = self.squeak_store.get_squeak(squeak_hash) + + # Download the secret key if not already owned. + if squeak and \ + not self.squeak_store.get_squeak_secret_key(squeak_hash): + secret_key = self.client.get_secret_key(squeak_hash) + if secret_key: + self.squeak_store.save_secret_key(squeak_hash, secret_key) + + # Download offer if the secret key if not already owned. + if squeak and \ + not self.squeak_store.get_squeak_secret_key(squeak_hash): + offer = self.client.get_offer(squeak_hash) + if offer: + self.squeak_store.handle_offer( + squeak, + offer, + self.peer.address, + ) + + def download_single_squeak( + self, + squeak_hash: bytes, + ) -> None: # Download the squeak if not already owned. if self.squeak_store.get_squeak(squeak_hash): - return - squeak = self.client.get_squeak(squeak_hash) - if squeak and self.is_squeak_wanted(squeak): - self.squeak_store.save_squeak(squeak) + raise Exception('Squeak already saved.') - def get_secret_key(self, squeak_hash: bytes) -> None: - # Get the squeak from the database. + squeak = self.client.get_squeak(squeak_hash) + if squeak and \ + get_hash(squeak) == squeak_hash: + self.squeak_store.save_squeak(squeak) + else: + raise Exception('Failed to download squeak.') + + # Get the local squeak. squeak = self.squeak_store.get_squeak(squeak_hash) - if squeak and self.is_squeak_wanted(squeak): - # Download the secret key is not already unlocked. - if self.squeak_store.get_squeak_secret_key(squeak_hash): - return + + # Download the secret key if not already owned. + if squeak and \ + not self.squeak_store.get_squeak_secret_key(squeak_hash): secret_key = self.client.get_secret_key(squeak_hash) if secret_key: self.squeak_store.save_secret_key(squeak_hash, secret_key) - def get_offer(self, squeak_hash: bytes) -> None: - # Get the squeak from the database. - squeak = self.squeak_store.get_squeak(squeak_hash) - if squeak and self.is_squeak_wanted(squeak): - # Download the secret key is not already unlocked. - if self.squeak_store.get_squeak_secret_key(squeak_hash): - return + # Download offer if the secret key if not already owned. + if squeak and \ + not self.squeak_store.get_squeak_secret_key(squeak_hash): offer = self.client.get_offer(squeak_hash) if offer: self.squeak_store.handle_offer( @@ -115,53 +129,3 @@ class PeerDownloader(ABC): offer, self.peer.address, ) - - -class RangeDownloader(PeerDownloader): - - def __init__( - self, - peer: SqueakPeer, - squeak_store: SqueakStore, - proxy_host: Optional[str], - proxy_port: Optional[int], - min_block: int, - max_block: int, - pubkeys: List[SqueakPublicKey], - ): - super().__init__(peer, squeak_store, proxy_host, proxy_port) - self.min_block = min_block - self.max_block = max_block - self.pubkeys = pubkeys - - def get_hashes(self) -> List[bytes]: - return self.client.lookup( - self.min_block, - self.max_block, - self.pubkeys, - ) - - def is_squeak_wanted(self, squeak: CSqueak) -> bool: - return squeak.nBlockHeight >= self.min_block and \ - squeak.nBlockHeight <= self.max_block and \ - squeak.GetPubKey() in self.pubkeys - - -class SingleDownloader(PeerDownloader): - - def __init__( - self, - peer: SqueakPeer, - squeak_store: SqueakStore, - proxy_host: Optional[str], - proxy_port: Optional[int], - squeak_hash: bytes, - ): - super().__init__(peer, squeak_store, proxy_host, proxy_port) - self.squeak_hash = squeak_hash - - def get_hashes(self) -> List[bytes]: - return [self.squeak_hash] - - def is_squeak_wanted(self, squeak: CSqueak) -> bool: - return get_hash(squeak) == self.squeak_hash diff --git a/squeaknode/node/squeak_controller.py b/squeaknode/node/squeak_controller.py index e7e0b255..6ab78f48 100644 --- a/squeaknode/node/squeak_controller.py +++ b/squeaknode/node/squeak_controller.py @@ -261,10 +261,6 @@ class SqueakController: limit: int, last_entry: Optional[SqueakEntry], ) -> List[SqueakEntry]: - # TODO: remove this temporary hack, after converting this to websockets. - # logger.info('Start downloading timeline...') - # self.network_controller.download_timeline() - # logger.info('Finished downloading timeline.') return self.squeak_store.get_timeline_squeak_entries(limit, last_entry) def get_liked_squeak_entries( @@ -294,10 +290,6 @@ class SqueakController: limit: int, last_entry: Optional[SqueakEntry], ) -> List[SqueakEntry]: - # TODO: remove this temporary hack, after converting this to websockets. - logger.info('Start downloading pubkey squeaks...') - self.network_controller.download_pubkey_squeaks_async(public_key) - logger.info('Finished downloading pubkey squeaks.') return self.squeak_store.get_squeak_entries_for_public_key( public_key, limit, diff --git a/squeaknode/node/squeak_download_worker.py b/squeaknode/node/squeak_download_worker.py index a2a04a8d..fc1faeb6 100644 --- a/squeaknode/node/squeak_download_worker.py +++ b/squeaknode/node/squeak_download_worker.py @@ -42,7 +42,7 @@ class SqueakDownloadWorker(PeriodicWorker): self.network_controller = network_controller def work_fn(self): - self.network_controller.download_timeline_async( + self.network_controller.download_timeline( self.interest_block_interval, )