async def test_locked_comm_drop_in_replacement(loop): a = await MyServer({}) await a.listen(0) read_event = asyncio.Event() read_event.set() read_queue = asyncio.Queue() original_pool = a.rpc a.rpc = _LockedCommPool(original_pool, read_event=read_event, read_queue=read_queue) b = await MyServer({}) await b.listen(0) # Event is set, the pool works like an ordinary pool res = await a.rpc(b.address).ping() assert await read_queue.get() == (b.address, "pong") assert res == "pong" assert b.counter == 1 read_event.clear() # Can also be used without a lock to intercept network traffic a.rpc = _LockedCommPool(original_pool, read_queue=read_queue) a.rpc.remove(b.address) res = await a.rpc(b.address).ping() assert await read_queue.get() == (b.address, "pong")
async def test_locked_comm_intercept_write(loop): a = await MyServer({}) await a.listen(0) b = await MyServer({}) await b.listen(0) write_event = asyncio.Event() write_queue = asyncio.Queue() a.rpc = _LockedCommPool(a.rpc, write_event=write_event, write_queue=write_queue) async def ping_pong(): return await a.rpc(b.address).ping() fut = asyncio.create_task(ping_pong()) with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(asyncio.shield(fut), 0.01) # Write was blocked. The remote hasn't received the message, yet assert b.counter == 0 assert await write_queue.get() == (b.address, { "op": "ping", "reply": True }) write_event.set() assert await fut == "pong"
async def test_locked_comm_intercept_read(loop): a = await MyServer({}) await a.listen(0) b = await MyServer({}) await b.listen(0) read_event = asyncio.Event() read_queue = asyncio.Queue() a.rpc = _LockedCommPool(a.rpc, read_event=read_event, read_queue=read_queue) async def ping_pong(): return await a.rpc(b.address).ping() fut = asyncio.create_task(ping_pong()) # We didn't block the write but merely the read. The remove should have # received the message and responded already while not b.counter: await asyncio.sleep(0.001) with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(asyncio.shield(fut), 0.01) assert await read_queue.get() == (b.address, "pong") read_event.set() assert await fut == "pong"
async def test_worker_stream_died_during_comm(c, s, a, b): write_queue = asyncio.Queue() write_event = asyncio.Event() b.rpc = _LockedCommPool( b.rpc, write_queue=write_queue, write_event=write_event, ) fut = c.submit(inc, 1, workers=[a.address], allow_other_workers=True) await fut # Actually no worker has the data; the scheduler is supposed to reschedule res = c.submit(inc, fut, workers=[b.address]) await write_queue.get() await a.close() write_event.set() await res assert any("receive-dep-failed" in msg for msg in b.log)