382 lines
16 KiB
Python
382 lines
16 KiB
Python
|
import asyncio
|
||
|
|
import typing
|
||
|
|
import logging
|
||
|
|
from lbrynet.utils import drain_tasks
|
||
|
|
from lbrynet.dht import constants
|
||
|
|
from lbrynet.dht.error import RemoteException
|
||
|
|
from lbrynet.dht.protocol.distance import Distance
|
||
|
|
|
||
|
|
from typing import TYPE_CHECKING
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from lbrynet.dht.protocol.routing_table import TreeRoutingTable
|
||
|
|
from lbrynet.dht.protocol.protocol import KademliaProtocol
|
||
|
|
from lbrynet.dht.peer import PeerManager, KademliaPeer
|
||
|
|
|
||
|
|
log = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class FindResponse:
|
||
|
|
@property
|
||
|
|
def found(self) -> bool:
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
def get_close_triples(self) -> typing.List[typing.Tuple[bytes, str, int]]:
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
|
||
|
|
class FindNodeResponse(FindResponse):
|
||
|
|
def __init__(self, key: bytes, close_triples: typing.List[typing.Tuple[bytes, str, int]]):
|
||
|
|
self.key = key
|
||
|
|
self.close_triples = close_triples
|
||
|
|
|
||
|
|
@property
|
||
|
|
def found(self) -> bool:
|
||
|
|
return self.key in [triple[0] for triple in self.close_triples]
|
||
|
|
|
||
|
|
def get_close_triples(self) -> typing.List[typing.Tuple[bytes, str, int]]:
|
||
|
|
return self.close_triples
|
||
|
|
|
||
|
|
|
||
|
|
class FindValueResponse(FindResponse):
|
||
|
|
def __init__(self, key: bytes, result_dict: typing.Dict):
|
||
|
|
self.key = key
|
||
|
|
self.token = result_dict[b'token']
|
||
|
|
self.close_triples: typing.List[typing.Tuple[bytes, bytes, int]] = result_dict.get(b'contacts', [])
|
||
|
|
self.found_compact_addresses = result_dict.get(key, [])
|
||
|
|
|
||
|
|
@property
|
||
|
|
def found(self) -> bool:
|
||
|
|
return len(self.found_compact_addresses) > 0
|
||
|
|
|
||
|
|
def get_close_triples(self) -> typing.List[typing.Tuple[bytes, str, int]]:
|
||
|
|
return [(node_id, address.decode(), port) for node_id, address, port in self.close_triples]
|
||
|
|
|
||
|
|
|
||
|
|
def get_shortlist(routing_table: 'TreeRoutingTable', key: bytes,
|
||
|
|
shortlist: typing.Optional[typing.List['KademliaPeer']]) -> typing.List['KademliaPeer']:
|
||
|
|
"""
|
||
|
|
If not provided, initialize the shortlist of peers to probe to the (up to) k closest peers in the routing table
|
||
|
|
|
||
|
|
:param routing_table: a TreeRoutingTable
|
||
|
|
:param key: a 48 byte hash
|
||
|
|
:param shortlist: optional manually provided shortlist, this is done during bootstrapping when there are no
|
||
|
|
peers in the routing table. During bootstrap the shortlist is set to be the seed nodes.
|
||
|
|
"""
|
||
|
|
if len(key) != constants.hash_length:
|
||
|
|
raise ValueError("invalid key length: %i" % len(key))
|
||
|
|
if not shortlist:
|
||
|
|
shortlist = routing_table.find_close_peers(key)
|
||
|
|
distance = Distance(key)
|
||
|
|
shortlist.sort(key=lambda peer: distance(peer.node_id), reverse=True)
|
||
|
|
return shortlist
|
||
|
|
|
||
|
|
|
||
|
|
class IterativeFinder:
|
||
|
|
def __init__(self, loop: asyncio.BaseEventLoop, peer_manager: 'PeerManager',
|
||
|
|
routing_table: 'TreeRoutingTable', protocol: 'KademliaProtocol', key: bytes,
|
||
|
|
bottom_out_limit: typing.Optional[int] = 2, max_results: typing.Optional[int] = constants.k,
|
||
|
|
exclude: typing.Optional[typing.List[typing.Tuple[str, int]]] = None,
|
||
|
|
shortlist: typing.Optional[typing.List['KademliaPeer']] = None):
|
||
|
|
if len(key) != constants.hash_length:
|
||
|
|
raise ValueError("invalid key length: %i" % len(key))
|
||
|
|
self.loop = loop
|
||
|
|
self.peer_manager = peer_manager
|
||
|
|
self.routing_table = routing_table
|
||
|
|
self.protocol = protocol
|
||
|
|
|
||
|
|
self.key = key
|
||
|
|
self.bottom_out_limit = bottom_out_limit
|
||
|
|
self.max_results = max_results
|
||
|
|
self.exclude = exclude or []
|
||
|
|
|
||
|
|
self.shortlist: typing.List['KademliaPeer'] = get_shortlist(routing_table, key, shortlist)
|
||
|
|
self.active: typing.List['KademliaPeer'] = []
|
||
|
|
self.contacted: typing.List[typing.Tuple[str, int]] = []
|
||
|
|
self.distance = Distance(key)
|
||
|
|
|
||
|
|
self.closest_peer: typing.Optional['KademliaPeer'] = None if not self.shortlist else self.shortlist[0]
|
||
|
|
self.prev_closest_peer: typing.Optional['KademliaPeer'] = None
|
||
|
|
|
||
|
|
self.iteration_queue = asyncio.Queue(loop=self.loop)
|
||
|
|
|
||
|
|
self.running_probes: typing.List[asyncio.Task] = []
|
||
|
|
self.lock = asyncio.Lock(loop=self.loop)
|
||
|
|
self.iteration_count = 0
|
||
|
|
self.bottom_out_count = 0
|
||
|
|
self.running = False
|
||
|
|
self.tasks: typing.List[asyncio.Task] = []
|
||
|
|
self.delayed_calls: typing.List[asyncio.Handle] = []
|
||
|
|
self.finished = asyncio.Event(loop=self.loop)
|
||
|
|
|
||
|
|
async def send_probe(self, peer: 'KademliaPeer') -> FindResponse:
|
||
|
|
"""
|
||
|
|
Send the rpc request to the peer and return an object with the FindResponse interface
|
||
|
|
"""
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
def check_result_ready(self, response: FindResponse):
|
||
|
|
"""
|
||
|
|
Called with a lock after adding peers from an rpc result to the shortlist.
|
||
|
|
This method is responsible for putting a result for the generator into the Queue
|
||
|
|
"""
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
def get_initial_result(self) -> typing.List['KademliaPeer']:
|
||
|
|
"""
|
||
|
|
Get an initial or cached result to be put into the Queue. Used for findValue requests where the blob
|
||
|
|
has peers in the local data store of blobs announced to us
|
||
|
|
"""
|
||
|
|
return []
|
||
|
|
|
||
|
|
def _is_closer(self, peer: 'KademliaPeer') -> bool:
|
||
|
|
if not self.closest_peer:
|
||
|
|
return True
|
||
|
|
return self.distance.is_closer(peer.node_id, self.closest_peer.node_id)
|
||
|
|
|
||
|
|
def _update_closest(self):
|
||
|
|
self.shortlist.sort(key=lambda peer: self.distance(peer.node_id), reverse=True)
|
||
|
|
if self.closest_peer and self.closest_peer is not self.shortlist[-1]:
|
||
|
|
if self._is_closer(self.shortlist[-1]):
|
||
|
|
self.prev_closest_peer = self.closest_peer
|
||
|
|
self.closest_peer = self.shortlist[-1]
|
||
|
|
|
||
|
|
async def _handle_probe_result(self, peer: 'KademliaPeer', response: FindResponse):
|
||
|
|
async with self.lock:
|
||
|
|
if peer not in self.shortlist:
|
||
|
|
self.shortlist.append(peer)
|
||
|
|
if peer not in self.active:
|
||
|
|
self.active.append(peer)
|
||
|
|
for contact_triple in response.get_close_triples():
|
||
|
|
addr_tuple = (contact_triple[1], contact_triple[2])
|
||
|
|
if addr_tuple not in self.contacted: # and not self.peer_manager.is_ignored(addr_tuple)
|
||
|
|
found_peer = self.peer_manager.get_kademlia_peer(
|
||
|
|
contact_triple[0], contact_triple[1], contact_triple[2]
|
||
|
|
)
|
||
|
|
if found_peer not in self.shortlist and self.peer_manager.peer_is_good(peer) is not False:
|
||
|
|
self.shortlist.append(found_peer)
|
||
|
|
self._update_closest()
|
||
|
|
self.check_result_ready(response)
|
||
|
|
|
||
|
|
async def _send_probe(self, peer: 'KademliaPeer'):
|
||
|
|
try:
|
||
|
|
response = await self.send_probe(peer)
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
return
|
||
|
|
except asyncio.TimeoutError:
|
||
|
|
if peer in self.active:
|
||
|
|
self.active.remove(peer)
|
||
|
|
return
|
||
|
|
except ValueError as err:
|
||
|
|
log.warning(str(err))
|
||
|
|
if peer in self.active:
|
||
|
|
self.active.remove(peer)
|
||
|
|
return
|
||
|
|
except RemoteException:
|
||
|
|
return
|
||
|
|
return await self._handle_probe_result(peer, response)
|
||
|
|
|
||
|
|
async def _search_round(self):
|
||
|
|
"""
|
||
|
|
Send up to constants.alpha (5) probes to the closest peers in the shortlist
|
||
|
|
"""
|
||
|
|
|
||
|
|
added = 0
|
||
|
|
async with self.lock:
|
||
|
|
self.shortlist.sort(key=lambda p: self.distance(p.node_id), reverse=True)
|
||
|
|
while self.running and len(self.shortlist) and added < constants.alpha:
|
||
|
|
peer = self.shortlist.pop()
|
||
|
|
origin_address = (peer.address, peer.udp_port)
|
||
|
|
if origin_address in self.exclude or self.peer_manager.peer_is_good(peer) is False:
|
||
|
|
continue
|
||
|
|
if peer.node_id == self.protocol.node_id:
|
||
|
|
continue
|
||
|
|
if (peer.address, peer.udp_port) == (self.protocol.external_ip, self.protocol.udp_port):
|
||
|
|
continue
|
||
|
|
if (peer.address, peer.udp_port) not in self.contacted:
|
||
|
|
self.contacted.append((peer.address, peer.udp_port))
|
||
|
|
|
||
|
|
t = self.loop.create_task(self._send_probe(peer))
|
||
|
|
|
||
|
|
def callback(_):
|
||
|
|
if t and t in self.running_probes:
|
||
|
|
self.running_probes.remove(t)
|
||
|
|
if not self.running_probes and self.shortlist:
|
||
|
|
self.tasks.append(self.loop.create_task(self._search_task(0.0)))
|
||
|
|
|
||
|
|
t.add_done_callback(callback)
|
||
|
|
self.running_probes.append(t)
|
||
|
|
added += 1
|
||
|
|
|
||
|
|
async def _search_task(self, delay: typing.Optional[float] = constants.iterative_lookup_delay):
|
||
|
|
try:
|
||
|
|
if self.running:
|
||
|
|
await self._search_round()
|
||
|
|
if self.running:
|
||
|
|
self.delayed_calls.append(self.loop.call_later(delay, self._search))
|
||
|
|
except (asyncio.CancelledError, StopAsyncIteration):
|
||
|
|
if self.running:
|
||
|
|
drain_tasks(self.running_probes)
|
||
|
|
self.running = False
|
||
|
|
|
||
|
|
def _search(self):
|
||
|
|
self.tasks.append(self.loop.create_task(self._search_task()))
|
||
|
|
|
||
|
|
def search(self):
|
||
|
|
if self.running:
|
||
|
|
raise Exception("already running")
|
||
|
|
self.running = True
|
||
|
|
self._search()
|
||
|
|
|
||
|
|
async def next_queue_or_finished(self) -> typing.List['KademliaPeer']:
|
||
|
|
peers = self.loop.create_task(self.iteration_queue.get())
|
||
|
|
finished = self.loop.create_task(self.finished.wait())
|
||
|
|
err = None
|
||
|
|
try:
|
||
|
|
await asyncio.wait([peers, finished], loop=self.loop, return_when='FIRST_COMPLETED')
|
||
|
|
if peers.done():
|
||
|
|
return peers.result()
|
||
|
|
raise StopAsyncIteration()
|
||
|
|
except asyncio.CancelledError as error:
|
||
|
|
err = error
|
||
|
|
finally:
|
||
|
|
if not finished.done() and not finished.cancelled():
|
||
|
|
finished.cancel()
|
||
|
|
if not peers.done() and not peers.cancelled():
|
||
|
|
peers.cancel()
|
||
|
|
if err:
|
||
|
|
raise err
|
||
|
|
|
||
|
|
def __aiter__(self):
|
||
|
|
self.search()
|
||
|
|
return self
|
||
|
|
|
||
|
|
async def __anext__(self) -> typing.List['KademliaPeer']:
|
||
|
|
try:
|
||
|
|
if self.iteration_count == 0:
|
||
|
|
initial_results = self.get_initial_result()
|
||
|
|
if initial_results:
|
||
|
|
self.iteration_queue.put_nowait(initial_results)
|
||
|
|
result = await self.next_queue_or_finished()
|
||
|
|
self.iteration_count += 1
|
||
|
|
return result
|
||
|
|
except (asyncio.CancelledError, StopAsyncIteration):
|
||
|
|
await self.aclose()
|
||
|
|
raise
|
||
|
|
|
||
|
|
def aclose(self):
|
||
|
|
self.running = False
|
||
|
|
|
||
|
|
async def _aclose():
|
||
|
|
async with self.lock:
|
||
|
|
self.running = False
|
||
|
|
if not self.finished.is_set():
|
||
|
|
self.finished.set()
|
||
|
|
drain_tasks(self.tasks)
|
||
|
|
drain_tasks(self.running_probes)
|
||
|
|
while self.delayed_calls:
|
||
|
|
timer = self.delayed_calls.pop()
|
||
|
|
if timer:
|
||
|
|
timer.cancel()
|
||
|
|
|
||
|
|
return asyncio.ensure_future(_aclose(), loop=self.loop)
|
||
|
|
|
||
|
|
|
||
|
|
class IterativeNodeFinder(IterativeFinder):
|
||
|
|
def __init__(self, loop: asyncio.BaseEventLoop, peer_manager: 'PeerManager',
|
||
|
|
routing_table: 'TreeRoutingTable', protocol: 'KademliaProtocol', key: bytes,
|
||
|
|
bottom_out_limit: typing.Optional[int] = 2, max_results: typing.Optional[int] = constants.k,
|
||
|
|
exclude: typing.Optional[typing.List[typing.Tuple[str, int]]] = None,
|
||
|
|
shortlist: typing.Optional[typing.List['KademliaPeer']] = None):
|
||
|
|
super().__init__(loop, peer_manager, routing_table, protocol, key, bottom_out_limit, max_results, exclude,
|
||
|
|
shortlist)
|
||
|
|
self.yielded_peers: typing.Set['KademliaPeer'] = set()
|
||
|
|
|
||
|
|
async def send_probe(self, peer: 'KademliaPeer') -> FindNodeResponse:
|
||
|
|
response = await self.protocol.get_rpc_peer(peer).find_node(self.key)
|
||
|
|
return FindNodeResponse(self.key, response)
|
||
|
|
|
||
|
|
def put_result(self, from_list: typing.List['KademliaPeer']):
|
||
|
|
not_yet_yielded = [peer for peer in from_list if peer not in self.yielded_peers]
|
||
|
|
not_yet_yielded.sort(key=lambda peer: self.distance(peer.node_id))
|
||
|
|
to_yield = not_yet_yielded[:min(constants.k, len(not_yet_yielded))]
|
||
|
|
if to_yield:
|
||
|
|
for peer in to_yield:
|
||
|
|
self.yielded_peers.add(peer)
|
||
|
|
self.iteration_queue.put_nowait(to_yield)
|
||
|
|
|
||
|
|
def check_result_ready(self, response: FindNodeResponse):
|
||
|
|
found = response.found and self.key != self.protocol.node_id
|
||
|
|
|
||
|
|
if found:
|
||
|
|
log.info("found")
|
||
|
|
self.put_result(self.shortlist)
|
||
|
|
if not self.finished.is_set():
|
||
|
|
self.finished.set()
|
||
|
|
return
|
||
|
|
if self.prev_closest_peer and self.closest_peer and not self._is_closer(self.prev_closest_peer):
|
||
|
|
# log.info("improving, %i %i %i %i %i", len(self.shortlist), len(self.active), len(self.contacted),
|
||
|
|
# self.bottom_out_count, self.iteration_count)
|
||
|
|
self.bottom_out_count = 0
|
||
|
|
elif self.prev_closest_peer and self.closest_peer:
|
||
|
|
self.bottom_out_count += 1
|
||
|
|
log.info("bottom out %i %i %i %i", len(self.active), len(self.contacted), len(self.shortlist),
|
||
|
|
self.bottom_out_count)
|
||
|
|
if self.bottom_out_count >= self.bottom_out_limit or self.iteration_count >= self.bottom_out_limit:
|
||
|
|
log.info("limit hit")
|
||
|
|
self.put_result(self.active)
|
||
|
|
if not self.finished.is_set():
|
||
|
|
self.finished.set()
|
||
|
|
return
|
||
|
|
if self.max_results and len(self.active) - len(self.yielded_peers) >= self.max_results:
|
||
|
|
log.info("max results")
|
||
|
|
self.put_result(self.active)
|
||
|
|
if not self.finished.is_set():
|
||
|
|
self.finished.set()
|
||
|
|
return
|
||
|
|
|
||
|
|
|
||
|
|
class IterativeValueFinder(IterativeFinder):
|
||
|
|
def __init__(self, loop: asyncio.BaseEventLoop, peer_manager: 'PeerManager',
|
||
|
|
routing_table: 'TreeRoutingTable', protocol: 'KademliaProtocol', key: bytes,
|
||
|
|
bottom_out_limit: typing.Optional[int] = 2, max_results: typing.Optional[int] = constants.k,
|
||
|
|
exclude: typing.Optional[typing.List[typing.Tuple[str, int]]] = None,
|
||
|
|
shortlist: typing.Optional[typing.List['KademliaPeer']] = None):
|
||
|
|
super().__init__(loop, peer_manager, routing_table, protocol, key, bottom_out_limit, max_results, exclude,
|
||
|
|
shortlist)
|
||
|
|
self.blob_peers: typing.Set['KademliaPeer'] = set()
|
||
|
|
|
||
|
|
async def send_probe(self, peer: 'KademliaPeer') -> FindValueResponse:
|
||
|
|
response = await self.protocol.get_rpc_peer(peer).find_value(self.key)
|
||
|
|
return FindValueResponse(self.key, response)
|
||
|
|
|
||
|
|
def check_result_ready(self, response: FindValueResponse):
|
||
|
|
if response.found:
|
||
|
|
blob_peers = [self.peer_manager.decode_tcp_peer_from_compact_address(compact_addr)
|
||
|
|
for compact_addr in response.found_compact_addresses]
|
||
|
|
to_yield = []
|
||
|
|
self.bottom_out_count = 0
|
||
|
|
for blob_peer in blob_peers:
|
||
|
|
if blob_peer not in self.blob_peers:
|
||
|
|
self.blob_peers.add(blob_peer)
|
||
|
|
to_yield.append(blob_peer)
|
||
|
|
if to_yield:
|
||
|
|
# log.info("found %i new peers for blob", len(to_yield))
|
||
|
|
self.iteration_queue.put_nowait(to_yield)
|
||
|
|
# if self.max_results and len(self.blob_peers) >= self.max_results:
|
||
|
|
# log.info("enough blob peers found")
|
||
|
|
# if not self.finished.is_set():
|
||
|
|
# self.finished.set()
|
||
|
|
return
|
||
|
|
if self.prev_closest_peer and self.closest_peer:
|
||
|
|
self.bottom_out_count += 1
|
||
|
|
if self.bottom_out_count >= self.bottom_out_limit:
|
||
|
|
log.info("blob peer search bottomed out")
|
||
|
|
if not self.finished.is_set():
|
||
|
|
self.finished.set()
|
||
|
|
return
|
||
|
|
|
||
|
|
def get_initial_result(self) -> typing.List['KademliaPeer']:
|
||
|
|
if self.protocol.data_store.has_peers_for_blob(self.key):
|
||
|
|
return self.protocol.data_store.get_peers_for_blob(self.key)
|
||
|
|
return []
|