mirror of
https://github.com/yzernik/squeaknode.git
synced 2026-08-18 13:09:08 +02:00
Simplify peer downloader (#2033)
* Remove subclasses for peer downloader and simplify download timeline method * Use thread pool executor for peer timeline download * Rename download timeline method * Save secret key and offer when downloading single squeak * Remove old download pubkey squeaks method from network controller * Fix get pubkey squeak entries method.
This commit is contained in:
Jonathan Zernik
GitHub
parent
bff7bbbc4b
commit
67da1c5a30
4 changed files with 95 additions and 145 deletions
|
|
@ -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!')
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue