test_blob_writers_concurrency

This commit is contained in:
Jack Robison 2019-04-23 15:36:06 -04:00
parent dab49e812f
commit 67088b2ca2
No known key found for this signature in database
GPG key ID: DF25C68FE0239BB2
2 changed files with 67 additions and 7 deletions

View file

@ -143,7 +143,8 @@ class AbstractBlob:
def close(self):
while self.writers:
peer, writer = self.writers.popitem()
writer.finished.cancel()
if writer and writer.finished and not writer.finished.done() and not self.loop.is_closed():
writer.finished.cancel()
while self.readers:
reader = self.readers.pop()
if reader:
@ -206,6 +207,12 @@ class AbstractBlob:
writer = HashBlobWriter(self.blob_hash, self.get_length, fut)
self.writers[(peer_address, peer_port)] = writer
def remove_writer(_):
if (peer_address, peer_port) in self.writers:
del self.writers[(peer_address, peer_port)]
fut.add_done_callback(remove_writer)
def writer_finished_callback(finished: asyncio.Future):
try:
err = finished.exception()
@ -215,16 +222,13 @@ class AbstractBlob:
while self.writers:
_, other = self.writers.popitem()
if other is not writer:
other.finished.cancel()
other.close_handle()
self.save_verified_blob(verified_bytes)
return
except (InvalidBlobHashError, InvalidDataError) as error:
log.warning("writer error downloading %s: %s", self.blob_hash[:8], str(error))
except (DownloadCancelledError, asyncio.CancelledError, asyncio.TimeoutError):
pass
finally:
if (peer_address, peer_port) in self.writers:
self.writers.pop((peer_address, peer_port))
fut.add_done_callback(writer_finished_callback)
return writer

View file

@ -109,7 +109,7 @@ class TestBlobExchange(BlobExchangeTestBase):
await self._add_blob_to_server(blob_hash, mock_blob_bytes)
second_client_blob = self.client_blob_manager.get_blob(blob_hash)
second_client_blob = second_client_blob_manager.get_blob(blob_hash)
# download the blob
await asyncio.gather(
@ -122,6 +122,62 @@ class TestBlobExchange(BlobExchangeTestBase):
await second_client_blob.verified.wait()
self.assertEqual(second_client_blob.get_is_verified(), True)
async def test_blob_writers_concurrency(self):
blob_hash = "7f5ab2def99f0ddd008da71db3a3772135f4002b19b7605840ed1034c8955431bd7079549e65e6b2a3b9c17c773073ed"
mock_blob_bytes = b'1' * ((2 * 2 ** 20) - 1)
blob = self.server_blob_manager.get_blob(blob_hash)
write_blob = blob._write_blob
write_called_count = 0
def wrap_write_blob(blob_bytes):
nonlocal write_called_count
write_called_count += 1
write_blob(blob_bytes)
blob._write_blob = wrap_write_blob
writer1 = blob.get_blob_writer(peer_port=1)
writer2 = blob.get_blob_writer(peer_port=2)
reader1_ctx_before_write = blob.reader_context()
with self.assertRaises(OSError):
blob.get_blob_writer(peer_port=2)
with self.assertRaises(OSError):
with blob.reader_context():
pass
blob.set_length(len(mock_blob_bytes))
results = {}
def check_finished_callback(writer, num):
def inner(writer_future: asyncio.Future):
results[num] = writer_future.result()
writer.finished.add_done_callback(inner)
check_finished_callback(writer1, 1)
check_finished_callback(writer2, 2)
def write_task(writer):
async def _inner():
writer.write(mock_blob_bytes)
return self.loop.create_task(_inner())
await asyncio.gather(write_task(writer1), write_task(writer2), loop=self.loop)
self.assertDictEqual({1: mock_blob_bytes, 2: mock_blob_bytes}, results)
self.assertEqual(1, write_called_count)
self.assertTrue(blob.get_is_verified())
self.assertDictEqual({}, blob.writers)
with reader1_ctx_before_write as f:
self.assertEqual(mock_blob_bytes, f.read())
with blob.reader_context() as f:
self.assertEqual(mock_blob_bytes, f.read())
with blob.reader_context() as f:
blob.close()
with self.assertRaises(ValueError):
f.read()
self.assertListEqual([], blob.readers)
async def test_host_different_blobs_to_multiple_peers_at_once(self):
blob_hash = "7f5ab2def99f0ddd008da71db3a3772135f4002b19b7605840ed1034c8955431bd7079549e65e6b2a3b9c17c773073ed"
mock_blob_bytes = b'1' * ((2 * 2 ** 20) - 1)