async def f():
server = Server({"ping": pingpong}, blocked_handlers=["ping"])
server.listen(8881)
comm = await connect(server.address)
await comm.write({"op": "ping"})
msg = await comm.read()
assert "exception" in msg
assert isinstance(msg["exception"], ValueError)
assert "'ping' handler has been explicitly disallowed" in repr(msg["exception"])
await comm.close()
server.stop()
def f():
server = Server({'echo': echo})
server.listen(8887)
data = b'0' * int(200e6) # slightly more than 100MB
conn = rpc(ip='127.0.0.1', port=8887)
result = yield conn.echo(x=data)
assert result == data
d = {'x': data}
result = yield conn.echo(x=d)
assert result == d
server.stop()
def f():
server = Server({'ping': pingpong})
server.listen(8887)
stream = yield connect('127.0.0.1', 8887)
yield write(stream, {'op': 'ping'})
response = yield read(stream)
assert response == b'pong'
yield write(stream, {'op': 'ping', 'close': True})
response = yield read(stream)
assert response == b'pong'
server.stop()
async def f():
async with Server({"ping": pingpong},
blocked_handlers=["ping"]) as server:
await server.listen(8881)
comm = await connect(server.address)
await comm.write({"op": "ping"})
msg = await comm.read()
_, exception, _ = clean_exception(msg["exception"])
assert isinstance(exception, ValueError)
assert "'ping' handler has been explicitly disallowed" in repr(
exception)
await comm.close()
def f():
server = Server({'ping': pingpong})
server.listen(8887)
with rpc(ip='127.0.0.1', port=8887) as remote:
response = yield remote.ping()
assert response == b'pong'
response = yield remote.ping(close=True)
assert response == b'pong'
assert not remote.streams
assert remote.status == 'closed'
server.stop()
async def check_large_packets(listen_arg):
""" tornado has a 100MB cap by default """
server = Server({"echo": echo})
await server.listen(listen_arg)
data = b"0" * int(200e6) # slightly more than 100MB
async with rpc(server.address) as conn:
result = await conn.echo(x=data)
assert result == data
d = {"x": data}
result = await conn.echo(x=d)
assert result == d
server.stop()
def f():
server = Server({'echo': echo})
server.listen(8886)
data = b'0' * int(200e6) # slightly more than 100MB
conn = rpc('127.0.0.1:8886')
result = yield conn.echo(x=data)
assert result == data
d = {'x': data}
result = yield conn.echo(x=d)
assert result == d
conn.close_comms()
server.stop()
def f():
server = Server({'ping': pingpong}, blocked_handlers=['ping'])
server.listen(8881)
comm = yield connect(server.address)
yield comm.write({'op': 'ping'})
msg = yield comm.read()
assert 'exception' in msg
assert isinstance(msg['exception'], ValueError)
assert "'ping' handler has been explicitly disallowed" in repr(
msg['exception'])
comm.close()
server.stop()
def check_rpc_with_many_connections(listen_arg):
@gen.coroutine
def g():
for i in range(10):
yield remote.ping()
server = Server({'ping': pingpong})
server.listen(listen_arg)
remote = rpc(server.address)
yield [g() for i in range(10)]
server.stop()
remote.close_comms()
assert all(comm.closed() for comm in remote.comms)
def check_large_packets(listen_arg):
""" tornado has a 100MB cap by default """
server = Server({'echo': echo})
server.listen(listen_arg)
data = b'0' * int(200e6) # slightly more than 100MB
conn = rpc(server.address)
result = yield conn.echo(x=data)
assert result == data
d = {'x': data}
result = yield conn.echo(x=d)
assert result == d
conn.close_comms()
server.stop()
async def check_rpc_with_many_connections(listen_arg):
async def g():
for i in range(10):
await remote.ping()
server = Server({"ping": pingpong})
await server.listen(listen_arg)
async with rpc(server.address) as remote:
for i in range(10):
await g()
server.stop()
remote.close_comms()
assert all(comm.closed() for comm in remote.comms)
async def test_send_recv_args():
server = Server({"echo": echo})
await server.listen(0)
comm = await connect(server.address)
result = await send_recv(comm, op="echo", x=b"1")
assert result == b"1"
assert not comm.closed()
result = await send_recv(comm, op="echo", x=b"2", reply=False)
assert result is None
assert not comm.closed()
result = await send_recv(comm, op="echo", x=b"3", close=True)
assert result == b"3"
assert comm.closed()
server.stop()
def test_send_recv_args():
server = Server({'echo': echo})
server.listen(0)
comm = yield connect(server.address)
result = yield send_recv(comm, op='echo', x=b'1')
assert result == b'1'
assert not comm.closed()
result = yield send_recv(comm, op='echo', x=b'2', reply=False)
assert result == None
assert not comm.closed()
result = yield send_recv(comm, op='echo', x=b'3', close=True)
assert result == b'3'
assert comm.closed()
server.stop()
def test_send_recv_args():
server = Server({'echo': echo})
server.listen(0)
result = yield send_recv(arg=('127.0.0.1', server.port), op='echo', x=b'1')
assert result == b'1'
result = yield send_recv(addr=('127.0.0.1:%d' % server.port).encode(),
op='echo', x=b'1')
assert result == b'1'
result = yield send_recv(ip=b'127.0.0.1', port=server.port, op='echo',
x=b'1')
assert result == b'1'
result = yield send_recv(ip=b'127.0.0.1', port=server.port, op='echo',
x=b'1', reply=False)
assert result == None
server.stop()
async def test_connection_pool():
async def ping(comm, delay=0.1):
await asyncio.sleep(delay)
return "pong"
servers = [Server({"ping": ping}) for i in range(10)]
for server in servers:
await server.listen(0)
rpc = await ConnectionPool(limit=5)
# Reuse connections
await asyncio.gather(
*[rpc(ip="127.0.0.1", port=s.port).ping() for s in servers[:5]]
)
await asyncio.gather(*[rpc(s.address).ping() for s in servers[:5]])
await asyncio.gather(*[rpc("127.0.0.1:%d" % s.port).ping() for s in servers[:5]])
await asyncio.gather(
*[rpc(ip="127.0.0.1", port=s.port).ping() for s in servers[:5]]
)
assert sum(map(len, rpc.available.values())) == 5
assert sum(map(len, rpc.occupied.values())) == 0
assert rpc.active == 0
assert rpc.open == 5
# Clear out connections to make room for more
await asyncio.gather(
*[rpc(ip="127.0.0.1", port=s.port).ping() for s in servers[5:]]
)
assert rpc.active == 0
assert rpc.open == 5
s = servers[0]
await asyncio.gather(
*[rpc(ip="127.0.0.1", port=s.port).ping(delay=0.1) for i in range(3)]
)
assert len(rpc.available["tcp://127.0.0.1:%d" % s.port]) == 3
# Explicitly clear out connections
rpc.collect()
start = time()
while any(rpc.available.values()):
await asyncio.sleep(0.01)
assert time() < start + 2
await rpc.close()
async def test_connection_pool_close_while_connecting(monkeypatch):
"""
Ensure a closed connection pool guarantees to have no connections left open
even if it is closed mid-connecting
"""
from distributed.comm.registry import backends
from distributed.comm.tcp import TCPBackend, TCPConnector
class SlowConnector(TCPConnector):
async def connect(self, address, deserialize, **connection_args):
await asyncio.sleep(0.1)
return await super().connect(address,
deserialize=deserialize,
**connection_args)
class SlowBackend(TCPBackend):
_connector_class = SlowConnector
monkeypatch.setitem(backends, "tcp", SlowBackend())
server = Server({})
await server.listen("tcp://")
pool = await ConnectionPool(limit=2)
async def connect_to_server():
comm = await pool.connect(server.address)
pool.reuse(server.address, comm)
tasks = [asyncio.create_task(connect_to_server()) for _ in range(30)]
await asyncio.sleep(0)
assert pool._connecting
close_fut = asyncio.create_task(pool.close())
with pytest.raises(
CommClosedError,
match="ConnectionPool not running. Status: Status.closed"):
await asyncio.gather(*tasks)
await close_fut
assert not pool.open
assert not pool._n_connecting
for t in tasks:
t.cancel()
def f():
server = Server({'ping': pingpong})
server.listen(8883)
with rpc('127.0.0.1:8883') as remote:
response = yield remote.ping()
assert response == b'pong'
assert remote.comms
assert remote.address == 'tcp://127.0.0.1:8883'
response = yield remote.ping(close=True)
assert response == b'pong'
assert not remote.comms
assert remote.status == 'closed'
server.stop()
def f():
server = Server({'ping': pingpong})
with pytest.raises(OSError):
server.port
server.listen(8887)
assert server.port == 8887
stream = yield connect('127.0.0.1', 8887)
yield write(stream, {'op': 'ping'})
response = yield read(stream)
assert response == b'pong'
yield write(stream, {'op': 'ping', 'close': True})
response = yield read(stream)
assert response == b'pong'
server.stop()
async def test_send_recv_cancelled():
"""Test that the comm channel is closed on CancelledError"""
async def get_stuck(comm):
await asyncio.Future()
server = Server({"get_stuck": get_stuck})
await server.listen(0)
client_comm = await connect(server.address, deserialize=False)
while not server._comms:
await asyncio.sleep(0.01)
server_comm = next(iter(server._comms))
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(send_recv(client_comm, op="get_stuck"),
timeout=0.1)
assert client_comm.closed()
while not server_comm.closed():
await asyncio.sleep(0.01)
async def test_connection_pool_respects_limit():
limit = 5
async def ping(comm, delay=0.01):
await asyncio.sleep(delay)
return "pong"
async def do_ping(pool, port):
assert pool.open <= limit
await pool(ip="127.0.0.1", port=port).ping()
assert pool.open <= limit
servers = [Server({"ping": ping}) for i in range(10)]
for server in servers:
await server.listen(0)
pool = await ConnectionPool(limit=limit)
await asyncio.gather(*[do_ping(pool, s.port) for s in servers])
def check_rpc(listen_arg, rpc_arg=None):
server = Server({'ping': pingpong})
server.listen(listen_arg)
if rpc_arg is None:
rpc_arg = server.address
with rpc(rpc_arg) as remote:
response = yield remote.ping()
assert response == b'pong'
assert remote.comms
response = yield remote.ping(close=True)
assert response == b'pong'
response = yield remote.ping()
assert response == b'pong'
assert not remote.comms
assert remote.status == 'closed'
server.stop()
def check_rpc(listen_addr, rpc_addr=None, listen_args=None, connection_args=None):
server = Server({"ping": pingpong})
server.listen(listen_addr, listen_args=listen_args)
if rpc_addr is None:
rpc_addr = server.address
with rpc(rpc_addr, connection_args=connection_args) as remote:
response = yield remote.ping()
assert response == b"pong"
assert remote.comms
response = yield remote.ping(close=True)
assert response == b"pong"
response = yield remote.ping()
assert response == b"pong"
assert not remote.comms
assert remote.status == "closed"
server.stop()
async def check_rpc(listen_addr, rpc_addr=None, listen_args={}, connection_args={}):
server = Server({"ping": pingpong})
await server.listen(listen_addr, **listen_args)
if rpc_addr is None:
rpc_addr = server.address
async with rpc(rpc_addr, connection_args=connection_args) as remote:
response = await remote.ping()
assert response == b"pong"
assert remote.comms
response = await remote.ping(close=True)
assert response == b"pong"
response = await remote.ping()
assert response == b"pong"
assert not remote.comms
assert remote.status == "closed"
server.stop()
await asyncio.sleep(0)
def test_connection_pool():
@gen.coroutine
def ping(comm, delay=0.1):
yield gen.sleep(delay)
raise gen.Return('pong')
servers = [Server({'ping': ping}) for i in range(10)]
for server in servers:
server.listen(0)
rpc = ConnectionPool(limit=5)
# Reuse connections
yield [rpc(ip='127.0.0.1', port=s.port).ping() for s in servers[:5]]
yield [rpc(s.address).ping() for s in servers[:5]]
yield [rpc('127.0.0.1:%d' % s.port).ping() for s in servers[:5]]
yield [rpc(ip='127.0.0.1', port=s.port).ping() for s in servers[:5]]
assert sum(map(len, rpc.available.values())) == 5
assert sum(map(len, rpc.occupied.values())) == 0
assert rpc.active == 0
assert rpc.open == 5
# Clear out connections to make room for more
yield [rpc(ip='127.0.0.1', port=s.port).ping() for s in servers[5:]]
assert rpc.active == 0
assert rpc.open == 5
s = servers[0]
yield [rpc(ip='127.0.0.1', port=s.port).ping(delay=0.1) for i in range(3)]
assert len(rpc.available['tcp://127.0.0.1:%d' % s.port]) == 3
# Explicitly clear out connections
rpc.collect()
start = time()
while any(rpc.available.values()):
yield gen.sleep(0.01)
assert time() < start + 2
rpc.close()
def test_send_recv_args():
server = Server({'echo': echo})
server.listen(0)
addr = '127.0.0.1:%d' % server.port
addr2 = server.address
result = yield send_recv(addr=addr, op='echo', x=b'1')
assert result == b'1'
result = yield send_recv(addr=addr, op='echo', x=b'2', reply=False)
assert result == None
result = yield send_recv(addr=addr2, op='echo', x=b'2')
assert result == b'2'
comm = yield connect(addr)
result = yield send_recv(comm, op='echo', x=b'3')
assert result == b'3'
assert not comm.closed()
result = yield send_recv(comm, op='echo', x=b'4', close=True)
assert result == b'4'
assert comm.closed()
server.stop()
async def test_connection_pool_remove():
async def ping(comm, delay=0.01):
await asyncio.sleep(delay)
return "pong"
servers = [Server({"ping": ping}) for i in range(5)]
for server in servers:
await server.listen(0)
rpc = await ConnectionPool(limit=10)
serv = servers.pop()
await asyncio.gather(*[rpc(s.address).ping() for s in servers])
await asyncio.gather(*[rpc(serv.address).ping() for i in range(3)])
await rpc.connect(serv.address)
assert sum(map(len, rpc.available.values())) == 6
assert sum(map(len, rpc.occupied.values())) == 1
assert rpc.active == 1
assert rpc.open == 7
rpc.remove(serv.address)
assert serv.address not in rpc.available
assert serv.address not in rpc.occupied
assert sum(map(len, rpc.available.values())) == 4
assert sum(map(len, rpc.occupied.values())) == 0
assert rpc.active == 0
assert rpc.open == 4
rpc.collect()
# this pattern of calls (esp. `reuse` after `remove`)
# can happen in case of worker failures:
comm = await rpc.connect(serv.address)
rpc.remove(serv.address)
rpc.reuse(serv.address, comm)
await rpc.close()
async def test_connection_pool_tls():
"""
Make sure connection args are supported.
"""
sec = tls_security()
connection_args = sec.get_connection_args("client")
listen_args = sec.get_listen_args("scheduler")
async def ping(comm, delay=0.01):
await asyncio.sleep(delay)
return "pong"
servers = [Server({"ping": ping}) for i in range(10)]
for server in servers:
await server.listen("tls://", **listen_args)
rpc = await ConnectionPool(limit=5, connection_args=connection_args)
await asyncio.gather(*[rpc(s.address).ping() for s in servers[:5]])
await asyncio.gather(*[rpc(s.address).ping() for s in servers[::2]])
await asyncio.gather(*[rpc(s.address).ping() for s in servers])
assert rpc.active == 0
await rpc.close()
def test_server_status_compare_enum_is_quiet():
"""That would be the default in user code"""
server = Server({})
server.status == Status.running
def test_server_status_assign_non_variant_raises():
server = Server({})
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("ignore")
with pytest.raises(AssertionError):
server.status = "I do not exists"
def test_errors(loop):
s = Server({})
try:
s.port
except OSError as e:
assert '.listen' in str(e)
