def tls_cluster_context(
worker_kwargs=None, scheduler_kwargs=None, security=None, **kwargs
):
security = security or tls_only_security()
worker_kwargs = assoc(worker_kwargs or {}, "security", security)
scheduler_kwargs = assoc(scheduler_kwargs or {}, "security", security)
with cluster(
worker_kwargs=worker_kwargs, scheduler_kwargs=scheduler_kwargs, **kwargs
) as (s, workers):
yield s, workers
async def run_traffic_jam(nsends, nbytes):
# This test eats `nsends * nbytes` bytes in RAM
np = pytest.importorskip("numpy")
from distributed.protocol import to_serialize
data = bytes(np.random.randint(0, 255, size=(nbytes, )).astype("u1").data)
async with EchoServer() as e:
comm = await connect(e.address)
b = BatchedSend(interval=0.01)
b.start(comm)
msg = {"x": to_serialize(data)}
for i in range(nsends):
b.send(assoc(msg, "i", i))
if np.random.random() > 0.5:
await asyncio.sleep(0.001)
results = []
count = 0
while len(results) < nsends:
# If this times out then I think it's a backpressure issue
# Somehow we're able to flood the socket so that the receiving end
# loses some of our messages
L = await asyncio.wait_for(comm.read(), 5)
count += 1
results.extend(r["i"] for r in L)
assert count == b.batch_count == e.count
assert b.message_count == nsends
assert results == list(range(nsends))
await comm.close() # external closing
await b.close()
