From 90b1405b34671f955ae0f183d3b8f7e0dc3eaa5c Mon Sep 17 00:00:00 2001
From: Danijar Hafner <mail@danijar.com>
Date: Wed, 7 Aug 2024 21:59:19 +0000
Subject: [PATCH] Better connection timeout message

---
 tests/test_server.py | 10 +++++-----
 zerofun/__init__.py  |  2 +-
 zerofun/client.py    | 13 ++++++++++---
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/tests/test_server.py b/tests/test_server.py
index 4d6fd03..28edc28 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -45,7 +45,7 @@ def test_multiple_clients(self, Server, addr):
     with server:
       clients = []
       for i in range(10):
-        client = zerofun.Client(addr, i, pings=0, maxage=1)
+        client = zerofun.Client(addr, str(i), pings=0, maxage=1)
         client.connect()
         clients.append(client)
       futures = [
@@ -95,7 +95,7 @@ def test_future_order(self, Server, addr):
     server = Server(addr)
     server.bind('function', lambda data: data)
     with server:
-      client = zerofun.Client(addr, 0, pings=0, maxage=1)
+      client = zerofun.Client(addr, pings=0, maxage=1)
       client.connect(retry=False, timeout=1)
       future1 = client.function({'foo': 1})
       future2 = client.function({'foo': 2})
@@ -112,7 +112,7 @@ def test_future_cleanup(self, Server, addr):
     server.bind('function', lambda data: data)
     with server:
       client = zerofun.Client(
-          addr, 0, pings=0, maxage=1, maxinflight=None, errors=False)
+          addr, pings=0, maxage=1, maxinflight=None, errors=False)
       client.connect(retry=False, timeout=1)
       client.function({'foo': 1})
       client.function({'foo': 2})
@@ -141,7 +141,7 @@ def workfn(data):
 
     with server:
       client = zerofun.Client(
-          addr, 0, pings=0, maxage=1, maxinflight=2)
+          addr, pings=0, maxage=1, maxinflight=2)
       client.connect(retry=False, timeout=1)
       futures = [client.function({'foo': i}) for i in range(4)]
       results = [future.result()['foo'] for future in futures]
@@ -154,7 +154,7 @@ def test_future_cleanup_errors(self, Server, addr):
     server = Server(addr)
     server.bind('function', lambda data: data)
     with server:
-      client = zerofun.Client(addr, 0, pings=0, maxage=1, errors=True)
+      client = zerofun.Client(addr, pings=0, maxage=1, errors=True)
       client.connect(retry=False, timeout=1)
       client.function({'foo': 1})
       client.function({'foo': 2})
diff --git a/zerofun/__init__.py b/zerofun/__init__.py
index c59dce5..2cd68be 100644
--- a/zerofun/__init__.py
+++ b/zerofun/__init__.py
@@ -1,4 +1,4 @@
-__version__ = '2.1.3'
+__version__ = '2.1.4'
 
 import multiprocessing as mp
 try:
diff --git a/zerofun/client.py b/zerofun/client.py
index 87de366..c08e497 100644
--- a/zerofun/client.py
+++ b/zerofun/client.py
@@ -17,6 +17,7 @@ def __init__(
       self, address, name='Client', ipv6=False, identity=None,
       pings=10, maxage=120, maxinflight=16, errors=True,
       connect=False):
+    assert isinstance(name, str)
     if identity is None:
       identity = int(np.random.randint(2 ** 32))
     self.address = address
@@ -52,9 +53,10 @@ def stats(self):
 
   @elements.timer.section('client_connect')
   def connect(self, retry=True, timeout=10):
+    self._print(f'Connecting to {self.resolved}')
+    warned = False
     while True:
       self.resolved = self._resolve(self.address)
-      self._print(f'Connecting to {self.resolved}')
       try:
         self.socket.connect(self.resolved, timeout)
         self._print('Connection established')
@@ -63,7 +65,12 @@ def connect(self, retry=True, timeout=10):
       except sockets.ProtocolError as e:
         self._print(f'Ignoring unexpected message: {e}')
       except sockets.ConnectError:
-        pass
+        if not warned:
+          self._print(
+              '\n  +------------------------------------------+' +
+              '\n  | Could not connect yet; retrying forever. |' +
+              '\n  +------------------------------------------+')
+          warned = True
       if retry:
         continue
       else:
@@ -99,7 +106,7 @@ def close(self):
     return self.socket.close()
 
   def _receive(self, rid, retry):
-    with elements.timer.section(f'client_{self.name.lower()}_receive'):
+    with elements.timer.section(f'{self.name}_client_receive'):
       while rid in self.futures and not self.futures[rid].done():
         result = self._listen()
         if result is None and not retry: