lbry-sdk/lbrynet/blob_exchange/downloader.py

import asyncio
import typing
import logging
from lbrynet.utils import drain_tasks
from lbrynet.blob_exchange.client import request_blob
if typing.TYPE_CHECKING:
    from lbrynet.conf import Config
    from lbrynet.dht.node import Node
    from lbrynet.dht.peer import KademliaPeer
    from lbrynet.blob.blob_manager import BlobFileManager
    from lbrynet.blob.blob_file import BlobFile

log = logging.getLogger(__name__)


class BlobDownloader:
    BAN_TIME = 10.0  # fixme: when connection manager gets implemented, move it out from here
    def __init__(self, loop: asyncio.BaseEventLoop, config: 'Config', blob_manager: 'BlobFileManager',
                 peer_queue: asyncio.Queue):
        self.loop = loop
        self.config = config
        self.blob_manager = blob_manager
        self.peer_queue = peer_queue
        self.active_connections: typing.Dict['KademliaPeer', asyncio.Task] = {}  # active request_blob calls
        self.ignored: typing.Dict['KademliaPeer', int] = {}
        self.scores: typing.Dict['KademliaPeer', int] = {}
        self.connections: typing.Dict['KademliaPeer', asyncio.Transport] = {}
        self.time_since_last_blob = loop.time()

    def should_race_continue(self, blob: 'BlobFile'):
        if len(self.active_connections) >= self.config.max_connections_per_download:
            return False
        # if a peer won 3 or more blob races and is active as a downloader, stop the race so bandwidth improves
        # the safe net side is that any failure will reset the peer score, triggering the race back
        # TODO: this is a good idea for low bandwidth, but doesnt play nice on high bandwidth
        # for peer, task in self.active_connections.items():
        #   if self.scores.get(peer, 0) >= 0 and self.rounds_won.get(peer, 0) >= 3 and not task.done():
        #       return False
        return not (blob.get_is_verified() or blob.file_exists)

    async def request_blob_from_peer(self, blob: 'BlobFile', peer: 'KademliaPeer'):
        if blob.get_is_verified():
            return
        self.scores[peer] = self.scores.get(peer, 0) - 1  # starts losing score, to account for cancelled ones
        transport = self.connections.get(peer)
        start = self.loop.time()
        bytes_received, transport = await request_blob(
            self.loop, blob, peer.address, peer.tcp_port, self.config.peer_connect_timeout,
            self.config.blob_download_timeout, connected_transport=transport
        )
        if bytes_received == blob.get_length():
            self.time_since_last_blob = self.loop.time()
        if not transport and peer not in self.ignored:
            self.ignored[peer] = self.loop.time()
            log.debug("drop peer %s:%i", peer.address, peer.tcp_port)
            if peer in self.connections:
                del self.connections[peer]
        elif transport:
            log.debug("keep peer %s:%i", peer.address, peer.tcp_port)
            self.connections[peer] = transport
            rough_speed = (bytes_received / (self.loop.time() - start)) if bytes_received else 0
            self.scores[peer] = rough_speed

    async def new_peer_or_finished(self, blob: 'BlobFile'):
        async def get_and_re_add_peers():
            try:
                new_peers = await asyncio.wait_for(self.peer_queue.get(), timeout=1.0)
                self.peer_queue.put_nowait(new_peers)
            except asyncio.TimeoutError:
                pass
        tasks = [self.loop.create_task(get_and_re_add_peers()), self.loop.create_task(blob.verified.wait())]
        active_tasks = list(self.active_connections.values())
        try:
            await asyncio.wait(tasks + active_tasks, loop=self.loop, return_when='FIRST_COMPLETED')
        finally:
            drain_tasks(tasks)

    def cleanup_active(self):
        to_remove = [peer for (peer, task) in self.active_connections.items() if task.done()]
        for peer in to_remove:
            del self.active_connections[peer]

    def clearbanned(self):
        now = self.loop.time()
        if now - self.time_since_last_blob > 60.0:
            return
        forgiven = [banned_peer for banned_peer, when in self.ignored.items() if now - when > self.BAN_TIME]
        self.peer_queue.put_nowait(forgiven)
        for banned_peer in forgiven:
            self.ignored.pop(banned_peer)

    async def download_blob(self, blob_hash: str, length: typing.Optional[int] = None) -> 'BlobFile':
        blob = self.blob_manager.get_blob(blob_hash, length)
        if blob.get_is_verified():
            return blob
        try:
            while not blob.get_is_verified():
                batch: typing.List['KademliaPeer'] = []
                while not self.peer_queue.empty():
                    batch.extend(self.peer_queue.get_nowait())
                batch.sort(key=lambda peer: self.scores.get(peer, 0), reverse=True)
                log.debug(
                    "running, %d peers, %d ignored, %d active",
                    len(batch), len(self.ignored), len(self.active_connections)
                )
                for peer in batch:
                    if not self.should_race_continue(blob):
                        break
                    if peer not in self.active_connections and peer not in self.ignored:
                        log.debug("request %s from %s:%i", blob_hash[:8], peer.address, peer.tcp_port)
                        t = self.loop.create_task(self.request_blob_from_peer(blob, peer))
                        self.active_connections[peer] = t
                await self.new_peer_or_finished(blob)
                self.cleanup_active()
                if batch:
                    self.peer_queue.put_nowait(set(batch).difference(self.ignored))
                else:
                    self.clearbanned()
            while self.active_connections:
                peer, task = self.active_connections.popitem()
                if task and not task.done():
                    task.cancel()
            blob.close()
            log.debug("downloaded %s", blob_hash[:8])
            return blob
        except asyncio.CancelledError:
            while self.active_connections:
                peer, task = self.active_connections.popitem()
                if task and not task.done():
                    task.cancel()
            raise
        except (OSError, Exception) as e:
            log.exception(e)
            raise e

    def close(self):
        self.scores.clear()
        self.ignored.clear()
        for transport in self.connections.values():
            transport.close()


async def download_blob(loop, config: 'Config', blob_manager: 'BlobFileManager', node: 'Node',
                        blob_hash: str) -> 'BlobFile':
    search_queue = asyncio.Queue(loop=loop, maxsize=config.max_connections_per_download)
    search_queue.put_nowait(blob_hash)
    peer_queue, accumulate_task = node.accumulate_peers(search_queue)
    downloader = BlobDownloader(loop, config, blob_manager, peer_queue)
    try:
        return await downloader.download_blob(blob_hash)
    finally:
        if accumulate_task and not accumulate_task.done():
            accumulate_task.cancel()
        downloader.close()