Repository: dask/distributed

Issue:
Worker creates excessive number of connections despite ConnectionPool
We observed that sometimes, workers crate a large number of tcp connections to the scheduler, well above the limit of 512 the ConnectionPool should impose.

To show the issue, we created an example on `ConnectionPool` level, limiting the number of connections to 2. If running this code, one can see from the print output that all 10 connections are established, despite this limit. Commenting in the `asyncio.sleep` line will make the `ConnectionPool` "work" again.

```python
from distributed.core import ConnectionPool
import asyncio
import random


pool = ConnectionPool(limit=2)
n_conns = 0


async def send_data_forever(addr):
    # await asyncio.sleep(random.random())
    conn = await pool.connect(addr)
    global n_conns
    n_conns += 1
    print(f"Have {n_conns} connections")
    while True:
        await conn.write(b'0')


async def client_connected_cb(client_reader, client_writer):
    print(f"[server] entering connection handler")
    try:
        data = await client_reader.read()
    finally:
        print(f"[server] exiting connection handler")


port = 2222

loop = asyncio.get_event_loop()
server_fut = asyncio.ensure_future(asyncio.start_server(client_connected_cb, port=port))
for i in range(10):
    asyncio.ensure_future(send_data_forever(addr=f'tcp://127.0.0.1:{port}'))
loop.run_forever()
```

Worker creates excessive number of connections despite ConnectionPool
We observed that sometimes, workers crate a large number of tcp connections to the scheduler, well above the limit of 512 the ConnectionPool should impose.

To show the issue, we created an example on `ConnectionPool` level, limiting the number of connections to 2. If running this code, one can see from the print output that all 10 connections are established, despite this limit. Commenting in the `asyncio.sleep` line will make the `ConnectionPool` "work" again.

```python
from distributed.core import ConnectionPool
import asyncio
import random


pool = ConnectionPool(limit=2)
n_conns = 0


async def send_data_forever(addr):
    # await asyncio.sleep(random.random())
    conn = await pool.connect(addr)
    global n_conns
    n_conns += 1
    print(f"Have {n_conns} connections")
    while True:
        await conn.write(b'0')


async def client_connected_cb(client_reader, client_writer):
    print(f"[server] entering connection handler")
    try:
        data = await client_reader.read()
    finally:
        print(f"[server] exiting connection handler")


port = 2222

loop = asyncio.get_event_loop()
server_fut = asyncio.ensure_future(asyncio.start_server(client_connected_cb, port=port))
for i in range(10):
    asyncio.ensure_future(send_data_forever(addr=f'tcp://127.0.0.1:{port}'))
loop.run_forever()
```



Discussion:

Resolving patch:
diff --git a/distributed/core.py b/distributed/core.py
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -826,7 +826,9 @@ def __init__(
         self.deserializers = deserializers if deserializers is not None else serializers
         self.connection_args = connection_args
         self.timeout = timeout
-        self.event = Event()
+        self._n_connecting = 0
+        # Invariant: semaphore._value == limit - open - _n_connecting
+        self.semaphore = asyncio.Semaphore(self.limit)
         self.server = weakref.ref(server) if server else None
         self._created = weakref.WeakSet()
         self._instances.add(self)
@@ -840,7 +842,11 @@ def open(self):
         return self.active + sum(map(len, self.available.values()))
 
     def __repr__(self):
-        return "<ConnectionPool: open=%d, active=%d>" % (self.open, self.active)
+        return "<ConnectionPool: open=%d, active=%d, connecting=%d>" % (
+            self.open,
+            self.active,
+            self._n_connecting,
+        )
 
     def __call__(self, addr=None, ip=None, port=None):
         """ Cached rpc objects """
@@ -861,10 +867,11 @@ async def connect(self, addr, timeout=None):
                 occupied.add(comm)
                 return comm
 
-        while self.open >= self.limit:
-            self.event.clear()
+        if self.semaphore.locked():
             self.collect()
-            await self.event.wait()
+
+        self._n_connecting += 1
+        await self.semaphore.acquire()
 
         try:
             comm = await connect(
@@ -877,11 +884,12 @@ async def connect(self, addr, timeout=None):
             comm._pool = weakref.ref(self)
             self._created.add(comm)
         except Exception:
+            self.semaphore.release()
             raise
-        occupied.add(comm)
+        finally:
+            self._n_connecting -= 1
 
-        if self.open >= self.limit:
-            self.event.clear()
+        occupied.add(comm)
 
         return comm
 
@@ -889,30 +897,34 @@ def reuse(self, addr, comm):
         """
         Reuse an open communication to the given address.  For internal use.
         """
-        try:
-            self.occupied[addr].remove(comm)
-        except KeyError:
-            pass
+        # if the pool is asked to re-use a comm it does not know about, ignore
+        # this comm: just close it.
+        if comm not in self.occupied[addr]:
+            IOLoop.current().add_callback(comm.close)
         else:
+            self.occupied[addr].remove(comm)
             if comm.closed():
-                if self.open < self.limit:
-                    self.event.set()
+                self.semaphore.release()
             else:
                 self.available[addr].add(comm)
+                if self.semaphore.locked() and self._n_connecting > 0:
+                    self.collect()
 
     def collect(self):
         """
         Collect open but unused communications, to allow opening other ones.
         """
         logger.info(
-            "Collecting unused comms.  open: %d, active: %d", self.open, self.active
+            "Collecting unused comms.  open: %d, active: %d, connecting: %d",
+            self.open,
+            self.active,
+            self._n_connecting,
         )
         for addr, comms in self.available.items():
             for comm in comms:
                 IOLoop.current().add_callback(comm.close)
+                self.semaphore.release()
             comms.clear()
-        if self.open < self.limit:
-            self.event.set()
 
     def remove(self, addr):
         """
@@ -923,12 +935,12 @@ def remove(self, addr):
             comms = self.available.pop(addr)
             for comm in comms:
                 IOLoop.current().add_callback(comm.close)
+                self.semaphore.release()
         if addr in self.occupied:
             comms = self.occupied.pop(addr)
             for comm in comms:
                 IOLoop.current().add_callback(comm.close)
-        if self.open < self.limit:
-            self.event.set()
+                self.semaphore.release()
 
     def close(self):
         """
@@ -937,8 +949,10 @@ def close(self):
         for comms in self.available.values():
             for comm in comms:
                 comm.abort()
+                self.semaphore.release()
         for comms in self.occupied.values():
             for comm in comms:
+                self.semaphore.release()
                 comm.abort()
 
         for comm in self._created:


Test patch:
diff --git a/distributed/tests/test_core.py b/distributed/tests/test_core.py
--- a/distributed/tests/test_core.py
+++ b/distributed/tests/test_core.py
@@ -1,10 +1,10 @@
+import asyncio
 from contextlib import contextmanager
 import os
 import socket
 import threading
 import weakref
 
-from tornado import gen
 import pytest
 
 import dask
@@ -82,8 +82,7 @@ def test_server(loop):
     Simple Server test.
     """
 
-    @gen.coroutine
-    def f():
+    async def f():
         server = Server({"ping": pingpong})
         with pytest.raises(ValueError):
             server.port
@@ -92,20 +91,20 @@ def f():
         assert server.address == ("tcp://%s:8881" % get_ip())
 
         for addr in ("127.0.0.1:8881", "tcp://127.0.0.1:8881", server.address):
-            comm = yield connect(addr)
+            comm = await connect(addr)
 
-            n = yield comm.write({"op": "ping"})
+            n = await comm.write({"op": "ping"})
             assert isinstance(n, int)
             assert 4 <= n <= 1000
 
-            response = yield comm.read()
+            response = await comm.read()
             assert response == b"pong"
 
-            yield comm.write({"op": "ping", "close": True})
-            response = yield comm.read()
+            await comm.write({"op": "ping", "close": True})
+            response = await comm.read()
             assert response == b"pong"
 
-            yield comm.close()
+            await comm.close()
 
         server.stop()
 
@@ -113,20 +112,19 @@ def f():
 
 
 def test_server_raises_on_blocked_handlers(loop):
-    @gen.coroutine
-    def f():
+    async 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()
+        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"])
 
-        yield comm.close()
+        await comm.close()
         server.stop()
 
     res = loop.run_sync(f)
@@ -251,21 +249,20 @@ def listen_on(cls, *args, **kwargs):
         yield assert_cannot_connect(inproc_addr2)
 
 
-@gen.coroutine
-def check_rpc(listen_addr, rpc_addr=None, listen_args=None, connection_args=None):
+async 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()
+        response = await remote.ping()
         assert response == b"pong"
         assert remote.comms
 
-        response = yield remote.ping(close=True)
+        response = await remote.ping(close=True)
         assert response == b"pong"
-        response = yield remote.ping()
+        response = await remote.ping()
         assert response == b"pong"
 
     assert not remote.comms
@@ -311,8 +308,7 @@ def test_rpc_inputs():
         r.close_rpc()
 
 
-@gen.coroutine
-def check_rpc_message_lifetime(*listen_args):
+async def check_rpc_message_lifetime(*listen_args):
     # Issue #956: rpc arguments and result shouldn't be kept alive longer
     # than necessary
     server = Server({"echo": echo_serialize})
@@ -324,15 +320,15 @@ def check_rpc_message_lifetime(*listen_args):
     del obj
     start = time()
     while CountedObject.n_instances != 0:
-        yield gen.sleep(0.01)
+        await asyncio.sleep(0.01)
         assert time() < start + 1
 
     with rpc(server.address) as remote:
         obj = CountedObject()
-        res = yield remote.echo(x=to_serialize(obj))
+        res = await remote.echo(x=to_serialize(obj))
         assert isinstance(res["result"], CountedObject)
         # Make sure resource cleanup code in coroutines runs
-        yield gen.sleep(0.05)
+        await asyncio.sleep(0.05)
 
         w1 = weakref.ref(obj)
         w2 = weakref.ref(res["result"])
@@ -361,18 +357,17 @@ def test_rpc_message_lifetime_inproc():
     yield check_rpc_message_lifetime("inproc://")
 
 
-@gen.coroutine
-def check_rpc_with_many_connections(listen_arg):
-    @gen.coroutine
-    def g():
+async def check_rpc_with_many_connections(listen_arg):
+    async def g():
         for i in range(10):
-            yield remote.ping()
+            await remote.ping()
 
     server = Server({"ping": pingpong})
     server.listen(listen_arg)
 
     with rpc(server.address) as remote:
-        yield [g() for i in range(10)]
+        for i in range(10):
+            await g()
 
         server.stop()
 
@@ -390,19 +385,18 @@ def test_rpc_with_many_connections_inproc():
     yield check_rpc_with_many_connections("inproc://")
 
 
-@gen.coroutine
-def check_large_packets(listen_arg):
+async 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)
+    result = await conn.echo(x=data)
     assert result == data
 
     d = {"x": data}
-    result = yield conn.echo(x=d)
+    result = await conn.echo(x=d)
     assert result == d
 
     conn.close_comms()
@@ -420,14 +414,13 @@ def test_large_packets_inproc():
     yield check_large_packets("inproc://")
 
 
-@gen.coroutine
-def check_identity(listen_arg):
+async def check_identity(listen_arg):
     server = Server({})
     server.listen(listen_arg)
 
     with rpc(server.address) as remote:
-        a = yield remote.identity()
-        b = yield remote.identity()
+        a = await remote.identity()
+        b = await remote.identity()
         assert a["type"] == "Server"
         assert a["id"] == b["id"]
 
@@ -489,7 +482,7 @@ def test_errors():
 
 @gen_test()
 def test_connect_raises():
-    with pytest.raises((gen.TimeoutError, IOError)):
+    with pytest.raises(IOError):
         yield connect("127.0.0.1:58259", timeout=0.01)
 
 
@@ -519,10 +512,9 @@ def test_coerce_to_address():
 
 @gen_test()
 def test_connection_pool():
-    @gen.coroutine
-    def ping(comm, delay=0.1):
-        yield gen.sleep(delay)
-        raise gen.Return("pong")
+    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:
@@ -553,12 +545,35 @@ def ping(comm, delay=0.1):
     rpc.collect()
     start = time()
     while any(rpc.available.values()):
-        yield gen.sleep(0.01)
+        yield asyncio.sleep(0.01)
         assert time() < start + 2
 
     rpc.close()
 
 
+@gen_test()
+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:
+        server.listen(0)
+
+    pool = ConnectionPool(limit=limit)
+
+    yield [do_ping(pool, s.port) for s in servers]
+
+
 @gen_test()
 def test_connection_pool_tls():
     """
@@ -568,10 +583,9 @@ def test_connection_pool_tls():
     connection_args = sec.get_connection_args("client")
     listen_args = sec.get_listen_args("scheduler")
 
-    @gen.coroutine
-    def ping(comm, delay=0.01):
-        yield gen.sleep(delay)
-        raise gen.Return("pong")
+    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:
@@ -589,10 +603,9 @@ def ping(comm, delay=0.01):
 
 @gen_test()
 def test_connection_pool_remove():
-    @gen.coroutine
-    def ping(comm, delay=0.01):
-        yield gen.sleep(delay)
-        raise gen.Return("pong")
+    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:
@@ -617,6 +630,9 @@ def ping(comm, delay=0.01):
     assert rpc.open == 4
 
     rpc.collect()
+
+    # this pattern of calls (esp. `reuse` after `remove`)
+    # can happen in case of worker failures:
     comm = yield rpc.connect(serv.address)
     rpc.remove(serv.address)
     rpc.reuse(serv.address, comm)
@@ -642,7 +658,7 @@ def test_counters():
 @gen_cluster()
 def test_ticks(s, a, b):
     pytest.importorskip("crick")
-    yield gen.sleep(0.1)
+    yield asyncio.sleep(0.1)
     c = s.digests["tick-duration"]
     assert c.size()
     assert 0.01 < c.components[0].quantile(0.5) < 0.5
@@ -657,7 +673,7 @@ def test_tick_logging(s, a, b):
     core.tick_maximum_delay = 0.001
     try:
         with captured_logger("distributed.core") as sio:
-            yield gen.sleep(0.1)
+            yield asyncio.sleep(0.1)
 
         text = sio.getvalue()
         assert "unresponsive" in text
@@ -671,14 +687,13 @@ def test_tick_logging(s, a, b):
 def test_compression(compression, serialize, loop):
     with dask.config.set(compression=compression):
 
-        @gen.coroutine
-        def f():
+        async def f():
             server = Server({"echo": serialize})
             server.listen("tcp://")
 
             with rpc(server.address) as r:
                 data = b"1" * 1000000
-                result = yield r.echo(x=to_serialize(data))
+                result = await r.echo(x=to_serialize(data))
                 assert result == {"result": data}
 
             server.stop()
@@ -687,17 +702,16 @@ def f():
 
 
 def test_rpc_serialization(loop):
-    @gen.coroutine
-    def f():
+    async def f():
         server = Server({"echo": echo_serialize})
         server.listen("tcp://")
 
         with rpc(server.address, serializers=["msgpack"]) as r:
             with pytest.raises(TypeError):
-                yield r.echo(x=to_serialize(inc))
+                await r.echo(x=to_serialize(inc))
 
         with rpc(server.address, serializers=["msgpack", "pickle"]) as r:
-            result = yield r.echo(x=to_serialize(inc))
+            result = await r.echo(x=to_serialize(inc))
             assert result == {"result": inc}
 
         server.stop()