Remove hard coded connect handshake timeouts

distributed/comm/core.py

@@ -214,8 +214,8 @@ async def _():
     async def on_connection(self, comm: Comm, handshake_overrides=None):
         local_info = {**comm.handshake_info(), **(handshake_overrides or {})}
         try:
-            write = await asyncio.wait_for(comm.write(local_info), 1)
-            handshake = await asyncio.wait_for(comm.read(), 1)
+            write = await comm.write(local_info)
+            handshake = await comm.read()
             # This would be better, but connections leak if worker is closed quickly
             # write, handshake = await asyncio.gather(comm.write(local_info), comm.read())
         except Exception as e:
@@ -263,78 +263,61 @@ async def connect(
 
     start = time()
     deadline = start + timeout
-    error = None
-
-    def _raise(error):
-        error = error or "connect() didn't finish in time"
-        msg = "Timed out trying to connect to %r after %s s: %s" % (
-            addr,
-            timeout,
-            error,
-        )
-        raise IOError(msg)
-
-    backoff = 0.01
-    if timeout and timeout / 20 < backoff:
-        backoff = timeout / 20
 
-    retry_timeout_backoff = random.randrange(140, 160) / 100
+    backoff_base = 0.01
+    attempt = 0
 
-    # This starts a thread
+    # Prefer multiple small attempts than one long attempt.
+    intermediate_cap = timeout / 5
     while True:
         try:
-            while deadline - time() > 0:
-
-                async def _():
-                    comm = await connector.connect(
-                        loc, deserialize=deserialize, **connection_args
-                    )
-                    local_info = {
-                        **comm.handshake_info(),
-                        **(handshake_overrides or {}),
-                    }
-                    try:
-                        handshake = await asyncio.wait_for(comm.read(), 1)
-                        write = await asyncio.wait_for(comm.write(local_info), 1)
-                        # This would be better, but connections leak if worker is closed quickly
-                        # write, handshake = await asyncio.gather(comm.write(local_info), comm.read())
-                    except Exception as e:
-                        with suppress(Exception):
-                            await comm.close()
-                        raise CommClosedError() from e
-
-                    comm.remote_info = handshake
-                    comm.remote_info["address"] = comm._peer_addr
-                    comm.local_info = local_info
-                    comm.local_info["address"] = comm._local_addr
-
-                    comm.handshake_options = comm.handshake_configuration(
-                        comm.local_info, comm.remote_info
-                    )
-                    return comm
-
-                with suppress(TimeoutError):
-                    comm = await asyncio.wait_for(
-                        _(), timeout=min(deadline - time(), retry_timeout_backoff)
-                    )
-                    break
-            if not comm:
-                _raise(error)
+            comm = await asyncio.wait_for(
+                connector.connect(loc, deserialize=deserialize, **connection_args),
+                timeout=min(intermediate_cap, deadline - time()),
+            )
+            break
         except FatalCommClosedError:
             raise
-        except EnvironmentError as e:
-            error = str(e)
+        # CommClosed, EnvironmentError inherit from OSError
+        except (TimeoutError, OSError) as exc:
+
             if time() < deadline:
-                logger.debug("Could not connect, waiting before retrying")
+                # FullJitter see https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
+                backoff = random.uniform(0, backoff_base * 2 ** attempt)
+                # If there is no more time left, just raise the exception
+                if backoff > deadline - time():
+                    raise IOError(
+                        f"Timed out trying to connect to {addr} after {timeout} s"
+                    ) from exc
+                attempt += 1
+                logger.debug(
+                    "Could not connect, waiting for %s before retrying", backoff
+                )
                 await asyncio.sleep(backoff)
-                backoff *= random.randrange(140, 160) / 100
-                retry_timeout_backoff *= random.randrange(140, 160) / 100
-                backoff = min(backoff, 1)  # wait at most one second
-            else:
-                _raise(error)
-        else:
-            break
 
+    local_info = {
+        **comm.handshake_info(),
+        **(handshake_overrides or {}),
+    }
+    try:
+        # This would be better, but connections leak if worker is closed quickly
+        # write, handshake = await asyncio.gather(comm.write(local_info), comm.read())
+        handshake = await asyncio.wait_for(comm.read(), deadline - time())
+        await asyncio.wait_for(comm.write(local_info), deadline - time())
+    except Exception as exc:
+        comm.abort()
+        raise IOError(
+            f"Timed out during handshake while connecting to {addr} after {timeout} s"
+        ) from exc
+
+    comm.remote_info = handshake
+    comm.remote_info["address"] = comm._peer_addr
+    comm.local_info = local_info
+    comm.local_info["address"] = comm._local_addr
+
+    comm.handshake_options = comm.handshake_configuration(
+        comm.local_info, comm.remote_info
+    )
     return comm
 
 

distributed/comm/tcp.py

@@ -1,6 +1,7 @@
 import errno
 import logging
 import socket
+from ssl import SSLError
 import struct
 import sys
 from tornado import gen
@@ -349,7 +350,6 @@ async def connect(self, address, deserialize=True, **connection_args):
             stream = await self.client.connect(
                 ip, port, max_buffer_size=MAX_BUFFER_SIZE, **kwargs
             )
-
             # Under certain circumstances tornado will have a closed connnection with an error and not raise
             # a StreamClosedError.
             #
@@ -360,6 +360,8 @@ async def connect(self, address, deserialize=True, **connection_args):
         except StreamClosedError as e:
             # The socket connect() call failed
             convert_stream_closed_error(self, e)
+        except SSLError as err:
+            raise FatalCommClosedError() from err
 
         local_address = self.prefix + get_stream_address(stream)
         comm = self.comm_class(

distributed/comm/tests/test_comms.py

@@ -1,45 +1,43 @@
 import asyncio
-import types
-from functools import partial
+from build.lib.distributed.comm.tcp import TCPConnector
 import os
 import sys
 import threading
+import types
 import warnings
+from functools import partial
 
+import distributed
 import pkg_resources
 import pytest
-
-from tornado import ioloop
-from tornado.concurrent import Future
-
-import distributed
-from distributed.metrics import time
-from distributed.utils import get_ip, get_ipv6
-from distributed.utils_test import (
-    requires_ipv6,
-    has_ipv6,
-    get_cert,
-    get_server_ssl_context,
-    get_client_ssl_context,
-)
-from distributed.utils_test import loop  # noqa: F401
-
-from distributed.protocol import to_serialize, Serialized, serialize, deserialize
-
-from distributed.comm.registry import backends, get_backend
 from distributed.comm import (
-    tcp,
-    inproc,
+    CommClosedError,
     connect,
+    get_address_host,
+    get_local_address_for,
+    inproc,
     listen,
-    CommClosedError,
     parse_address,
     parse_host_port,
-    unparse_host_port,
     resolve_address,
-    get_address_host,
-    get_local_address_for,
+    tcp,
+    unparse_host_port,
+)
+from distributed.comm.registry import backends, get_backend
+from distributed.comm.tcp import TCP, TCPBackend
+from distributed.metrics import time
+from distributed.protocol import Serialized, deserialize, serialize, to_serialize
+from distributed.utils import get_ip, get_ipv6
+from distributed.utils_test import loop  # noqa: F401
+from distributed.utils_test import (
+    get_cert,
+    get_client_ssl_context,
+    get_server_ssl_context,
+    has_ipv6,
+    requires_ipv6,
 )
+from tornado import ioloop
+from tornado.concurrent import Future
 
 EXTERNAL_IP4 = get_ip()
 if has_ipv6():
@@ -665,7 +663,8 @@ async def handle_comm(comm):
 
     with pytest.raises(EnvironmentError) as excinfo:
         await connect(listener.contact_address, timeout=2, ssl_context=cli_ctx)
-    assert "certificate verify failed" in str(excinfo.value)
+
+    assert "certificate verify failed" in str(excinfo.value.__cause__)
 
 
 #
@@ -797,10 +796,90 @@ async def handle_comm(comm):
 #
 
 
+async def echo(comm):
+    message = await comm.read()
+    await comm.write(message)
+
+
+@pytest.mark.asyncio
+async def test_retry_connect(monkeypatch):
+    async def echo(comm):
+        message = await comm.read()
+        await comm.write(message)
+
+    class UnreliableConnector(TCPConnector):
+        def __init__(self):
+
+            self.num_failures = 2
+            self.failures = 0
+            super().__init__()
+
+        async def connect(self, address, deserialize=True, **connection_args):
+            if self.failures > self.num_failures:
+                return await super().connect(address, deserialize, **connection_args)
+            else:
+                self.failures += 1
+                raise IOError()
+
+    class UnreliableBackend(TCPBackend):
+        """Don't try this at home"""
+
+        _connector_class = UnreliableConnector
+
+    monkeypatch.setitem(backends, "tcp", UnreliableBackend())
+
+    listener = await listen("tcp://127.0.0.1:1234", echo)
+    comm = await connect(listener.contact_address)
+    await comm.write(b"test")
+    msg = await comm.read()
+    assert msg == b"test"
+
+
+@pytest.mark.asyncio
+async def test_handshake_slow_comm(monkeypatch):
+    class SlowComm(TCP):
+        def __init__(self, *args, delay_in_comm=0.5, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.delay_in_comm = delay_in_comm
+
+        async def read(self, *args, **kwargs):
+            await asyncio.sleep(self.delay_in_comm)
+            return await super().read(*args, **kwargs)
+
+        async def write(self, *args, **kwargs):
+            await asyncio.sleep(self.delay_in_comm)
+            res = await super(type(self), self).write(*args, **kwargs)
+            return res
+
+    class SlowConnector(TCPConnector):
+        comm_class = SlowComm
+
+    class SlowBackend(TCPBackend):
+        """Don't try this at home"""
+
+        _connector_class = SlowConnector
+
+    monkeypatch.setitem(backends, "tcp", SlowBackend())
+
+    listener = await listen("tcp://127.0.0.1:1234", echo)
+    comm = await connect(listener.contact_address)
+    await comm.write(b"test")
+    msg = await comm.read()
+    assert msg == b"test"
+
+    import dask
+
+    with dask.config.set({"distributed.comm.timeouts.connect": "100ms"}):
+        with pytest.raises(
+            TimeoutError, match="Timed out during handshake while connecting to"
+        ):
+            await connect(listener.contact_address)
+
+
 async def check_connect_timeout(addr):
     t1 = time()
     with pytest.raises(IOError):
-        await connect(addr, timeout=0.15)
+        await connect(addr, timeout=0.3)
     dt = time() - t1
     assert 1 >= dt >= 0.1