Stuff

colony/core/roles/explorer/role.py

@@ -1,19 +1,23 @@
 import datetime
+from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
 
-from colony.core.network.requests import CommandRequest
 from colony.core.storage import DEFAULT_STORAGE_FILES
 from colony.core.roles.role import Role
 from colony.core.roles.unit.role import Unit
 from colony.core.roles.vaultwarden.role import VaultWarden
+from colony.settings import configuration
 from colony.transport.discover import Discover
-from colony.transport.whisperer import Whisperer
+from colony.utils.network import get_local_subnet
 
 
 class Explorer(Role):
     CONTINUOUS = True
     LOOP_SECONDS_WAIT = 60
 
+    SCAN_WORKERS = 64
+    SCAN_TIMEOUT = 1
+
     class COMMANDS(Enum):
         pass
 
@@ -26,29 +30,77 @@ class Explorer(Role):
 
     def loop(self):
         self.logger.info("Getting network snapshot")
-        snapshot = self.get_snapshot()['result'] or {}
+        snapshot = self.load_snapshot()
+        if snapshot is None:
+            return
         self.logger.debug(f"Loaded snapshot")
 
         discovered_peers = self.discover_peers()
         self.logger.info(f"Discovered peers: {list(discovered_peers.keys())}")
 
+        reachable = set()
         for ip, port in discovered_peers.items():
             self.logger.info(f"Sending SHOW_STATS command to {ip}:{port}")
 
             stats = self.fetch_stats_from_peer(ip, port)
+            if stats is None:
+                continue
 
             self.logger.debug(f"Received stats from {ip}")
-            snapshot[ip] = stats | {'last_update': datetime.datetime.now().isoformat(), 'last_seen': datetime.datetime.now().isoformat(), 'status': 'reachable'}
+            now = datetime.datetime.now().isoformat()
+            snapshot[ip] = {
+                'status': 'reachable',
+                'result': stats,
+                'last_update': now,
+                'last_seen': now,
+            }
+            reachable.add(ip)
 
-        missing_ips = set(snapshot.keys()) - set(discovered_peers.keys())
-        for ip in missing_ips:
+        for ip in snapshot.keys() - reachable:
             self.logger.warning(f"Peer {ip} was in snapshot but not reachable. Marking as unreachable")
             snapshot[ip]['last_update'] = datetime.datetime.now().isoformat()
             snapshot[ip]['status'] = 'unreachable'
 
-        self.set_snapshot(snapshot)
-        self.logger.info("Snapshot updated successfully")
+        if self._stop_event.is_set():
+            self.logger.info("Role was removed during the scan, not writing the snapshot")
+            return
 
+        self.save_snapshot(snapshot)
+
+    def load_snapshot(self):
+        try:
+            resp = self.get_snapshot()
+        except Exception as e:
+            self.logger.warning(f"Could not read the snapshot: {e}")
+            return None
+
+        if not isinstance(resp, dict) or resp.get('status') != 'ok':
+            self.logger.warning("Could not read the snapshot: VaultWarden returned an error")
+            return None
+
+        snapshot = resp.get('result')
+        if snapshot is None:
+            return {}
+        if not isinstance(snapshot, dict):
+            self.logger.warning("Stored snapshot is not an object. Refusing to overwrite it")
+            return None
+
+        for ip in [ip for ip, entry in snapshot.items() if not isinstance(entry, dict)]:
+            self.logger.warning(f"Dropping malformed snapshot entry for {ip}")
+            del snapshot[ip]
+        return snapshot
+
+    def save_snapshot(self, snapshot):
+        try:
+            resp = self.set_snapshot(snapshot)
+        except Exception as e:
+            self.logger.error(f"Could not write the snapshot: {e}")
+            return
+
+        if not isinstance(resp, dict) or resp.get('status') != 'ok':
+            self.logger.error("Could not write the snapshot: VaultWarden returned an error")
+            return
+        self.logger.info("Snapshot updated successfully")
 
     def get_snapshot(self):
         return self.context.execute(VaultWarden.COMMANDS.OPEN, DEFAULT_STORAGE_FILES.COLONY_SNAPSHOT.value)
@@ -56,17 +108,32 @@ class Explorer(Role):
     def set_snapshot(self, snapshot):
         return self.context.execute(VaultWarden.COMMANDS.OVERWRITE, DEFAULT_STORAGE_FILES.COLONY_SNAPSHOT.value, snapshot)
 
-    def discover_peers(self):
-        peers = {}
+    def _hosts_to_scan(self):
+        configured = configuration.get("explorer.hosts")
+        if isinstance(configured, (list, tuple)) and configured:
+            return [str(host) for host in configured]
+        return [str(host) for host in get_local_subnet().hosts()]
 
-        # for ip in get_local_subnet().hosts():
-        for ip in ['192.168.65.101', '192.168.65.102']:
-            try:
-                port = Discover.request(str(ip))
-                peers[str(ip)] = port
-            except (TimeoutError, ConnectionRefusedError, OSError):
-                continue
-        return peers
+    def _probe(self, ip):
+        try:
+            return ip, Discover.request(ip, timeout=self.SCAN_TIMEOUT)
+        except Exception:
+            return ip, None
+
+    def discover_peers(self):
+        with ThreadPoolExecutor(max_workers=self.SCAN_WORKERS) as pool:
+            results = list(pool.map(self._probe, self._hosts_to_scan()))
+
+        return {ip: port for ip, port in results if port is not None}
 
     def fetch_stats_from_peer(self, ip, port):
-        return self.context.execute_to_unregistered_node(ip, port, Unit.COMMANDS.SHOW_STATS)
+        try:
+            resp = self.context.execute_to_unregistered_node(ip, port, Unit.COMMANDS.SHOW_STATS)
+        except Exception as e:
+            self.logger.warning(f"Could not fetch stats from {ip}:{port}: {e}")
+            return None
+
+        if not isinstance(resp, dict) or resp.get('status') != 'ok' or not isinstance(resp.get('result'), dict):
+            self.logger.warning(f"Peer {ip}:{port} answered SHOW_STATS with an error")
+            return None
+        return resp['result']