Stuff

colony/transport/discover.py

@@ -1,32 +1,50 @@
 import socket
 import struct
+import time
+
+from colony.core.logger.logging import get_logger
+
 
 class Discover:
     DISCOVER_PORT = 28149
     MAGIC = b"OK"
+    CLIENT_TIMEOUT = 2
 
     @staticmethod
     def listen(real_port, host="0.0.0.0"):
+        logger = get_logger(Discover.__name__)
+
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         s.bind((host, Discover.DISCOVER_PORT))
         s.listen()
 
         while True:
-            conn, _ = s.accept()
             try:
-                data = conn.recv(2)
-                if data == Discover.MAGIC:
-                    conn.sendall(struct.pack("!H", real_port))
-            finally:
-                conn.close()
+                conn, _ = s.accept()
+            except OSError as e:
+                logger.error(f"Accept failed: {e}")
+                time.sleep(1)
+                continue
+
+            with conn:
+                try:
+                    conn.settimeout(Discover.CLIENT_TIMEOUT)
+                    if conn.recv(len(Discover.MAGIC)) == Discover.MAGIC:
+                        conn.sendall(struct.pack("!H", real_port))
+                except OSError as e:
+                    logger.debug(f"Discovery request failed: {e}")
 
     @staticmethod
     def request(target_host, timeout=3):
-        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        s.settimeout(timeout)
-        s.connect((target_host, Discover.DISCOVER_PORT))
-        s.sendall(Discover.MAGIC)
-        data = s.recv(2)
-        s.close()
+        with socket.create_connection((target_host, Discover.DISCOVER_PORT), timeout=timeout) as s:
+            s.sendall(Discover.MAGIC)
+
+            data = b""
+            while len(data) < 2:
+                chunk = s.recv(2 - len(data))
+                if not chunk:
+                    raise ConnectionError("Discovery response was cut short")
+                data += chunk
+
         return struct.unpack("!H", data)[0]