fix request_blob cancellation

This commit is contained in:
Jack Robison 2019-02-04 16:03:08 -05:00
parent ad30b7df6d
commit 377442d9f0
No known key found for this signature in database
GPG key ID: DF25C68FE0239BB2

View file

@ -29,18 +29,13 @@ class BlobDownloader:
self.ignored: typing.Set['KademliaPeer'] = set() self.ignored: typing.Set['KademliaPeer'] = set()
self.scores: typing.Dict['KademliaPeer', int] = {} self.scores: typing.Dict['KademliaPeer', int] = {}
def request_blob_from_peer(self, blob: 'BlobFile', peer: 'KademliaPeer'): async def request_blob_from_peer(self, blob: 'BlobFile', peer: 'KademliaPeer'):
async def _request_blob():
if blob.get_is_verified(): if blob.get_is_verified():
return return
try:
success, keep_connection = await request_blob( success, keep_connection = await request_blob(
self.loop, blob, peer.address, peer.tcp_port, self.config.peer_connect_timeout, self.loop, blob, peer.address, peer.tcp_port, self.config.peer_connect_timeout,
self.config.blob_download_timeout self.config.blob_download_timeout
) )
finally:
if peer in self.active_connections:
self.active_connections.pop(peer)
if not keep_connection and peer not in self.ignored: if not keep_connection and peer not in self.ignored:
self.ignored.add(peer) self.ignored.add(peer)
log.debug("drop peer %s:%i", peer.address, peer.tcp_port) log.debug("drop peer %s:%i", peer.address, peer.tcp_port)
@ -50,7 +45,6 @@ class BlobDownloader:
self.scores[peer] = self.scores.get(peer, 0) + 2 self.scores[peer] = self.scores.get(peer, 0) + 2
else: else:
self.scores[peer] = self.scores.get(peer, 0) - 1 self.scores[peer] = self.scores.get(peer, 0) - 1
return self.loop.create_task(_request_blob())
async def new_peer_or_finished(self, blob: 'BlobFile'): async def new_peer_or_finished(self, blob: 'BlobFile'):
async def get_and_re_add_peers(): async def get_and_re_add_peers():
@ -77,7 +71,12 @@ class BlobDownloader:
break break
if peer not in self.active_connections and peer not in self.ignored: 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) log.debug("request %s from %s:%i", blob_hash[:8], peer.address, peer.tcp_port)
self.active_connections[peer] = self.request_blob_from_peer(blob, peer) t = self.loop.create_task(self.request_blob_from_peer(blob, peer))
self.active_connections[peer] = t
t.add_done_callback(
lambda _:
None if peer not in self.active_connections else self.active_connections.pop(peer)
)
await self.new_peer_or_finished(blob) await self.new_peer_or_finished(blob)
to_re_add = list(set(filter(lambda peer: peer not in self.ignored, batch))) to_re_add = list(set(filter(lambda peer: peer not in self.ignored, batch)))
to_re_add.sort(key=lambda peer: self.scores.get(peer, 0), reverse=True) to_re_add.sort(key=lambda peer: self.scores.get(peer, 0), reverse=True)