221 lines
8.4 KiB
Python
221 lines
8.4 KiB
Python
import datetime
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from enum import Enum
|
|
|
|
from colony.core.network.membership import MemberRecord
|
|
from colony.core.storage import DEFAULT_STORAGE_FILES
|
|
from colony.core.roles.gossip.role import Gossip
|
|
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, setting
|
|
from colony.transport.discover import Discover
|
|
from colony.utils.network import hosts_in_range
|
|
|
|
|
|
class Explorer(Role):
|
|
CONTINUOUS = True
|
|
LOOP_SECONDS_WAIT = 60
|
|
|
|
SCAN_WORKERS = 64
|
|
SCAN_TIMEOUT = 1
|
|
|
|
class COMMANDS(Enum):
|
|
pass
|
|
|
|
def __init__(self, context):
|
|
commands = {
|
|
|
|
}
|
|
|
|
super().__init__(commands, context)
|
|
|
|
def loop(self):
|
|
self.logger.info("Getting network snapshot")
|
|
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}")
|
|
now = datetime.datetime.now().isoformat()
|
|
snapshot[ip] = {
|
|
'status': 'reachable',
|
|
'result': stats,
|
|
'last_update': now,
|
|
'last_seen': now,
|
|
}
|
|
reachable.add(ip)
|
|
|
|
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'
|
|
|
|
if self._stop_event.is_set():
|
|
self.logger.info("Role was removed during the scan, not writing the snapshot")
|
|
return
|
|
|
|
self.save_snapshot(snapshot)
|
|
self.introduce_unknown_peers(discovered_peers)
|
|
|
|
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)
|
|
|
|
def set_snapshot(self, snapshot):
|
|
return self.context.execute(VaultWarden.COMMANDS.OVERWRITE, DEFAULT_STORAGE_FILES.COLONY_SNAPSHOT.value, snapshot)
|
|
|
|
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 hosts_in_range(setting("discovery.range_start"), setting("discovery.range_end"))
|
|
|
|
def _probe(self, ip):
|
|
try:
|
|
return ip, Discover.request(ip, timeout=self.SCAN_TIMEOUT, discover_port=setting("discovery.discover_port"))
|
|
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):
|
|
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']
|
|
|
|
def fetch_members_from_peer(self, ip, port):
|
|
try:
|
|
resp = self.context.execute_to_unregistered_node(ip, port, Unit.COMMANDS.SHOW_CLUSTER)
|
|
except Exception as e:
|
|
self.logger.warning(f"Could not fetch cluster members from {ip}:{port}: {e}")
|
|
return []
|
|
|
|
result = resp.get('result') if isinstance(resp, dict) and resp.get('status') == 'ok' else None
|
|
rows = result.get('members') if isinstance(result, dict) else None
|
|
if not isinstance(rows, list):
|
|
self.logger.warning(f"Peer {ip}:{port} answered SHOW_CLUSTER with an error")
|
|
return []
|
|
|
|
members = []
|
|
for row in rows[:setting("membership.max_members")]:
|
|
try:
|
|
record = MemberRecord.from_dict(row)
|
|
except ValueError as e:
|
|
self.logger.debug(f"Ignoring bad member from {ip}:{port}: {e}")
|
|
continue
|
|
if row.get('alive') is True:
|
|
members.append(record.key)
|
|
return members
|
|
|
|
def _ping(self, address):
|
|
ip, _ = address
|
|
try:
|
|
port = Discover.request(ip, timeout=self.SCAN_TIMEOUT, discover_port=setting("discovery.discover_port"))
|
|
except Exception:
|
|
return None
|
|
return (ip, port) if port is not None else None
|
|
|
|
def introduce_unknown_peers(self, discovered_peers):
|
|
membership = self.context.membership()
|
|
known = {record.key for record in membership.alive_records()} | set(discovered_peers.items())
|
|
|
|
candidates = {}
|
|
for ip, port in discovered_peers.items():
|
|
for address in self.fetch_members_from_peer(ip, port):
|
|
if address not in known:
|
|
candidates[address] = None
|
|
if not candidates or self._stop_event.is_set():
|
|
return
|
|
|
|
addresses = list(candidates)[:setting("membership.max_members")]
|
|
with ThreadPoolExecutor(max_workers=min(self.SCAN_WORKERS, len(addresses))) as pool:
|
|
pinged = [address for address in pool.map(self._ping, addresses) if address is not None]
|
|
|
|
for ip, port in pinged:
|
|
self.logger.info(f"Peers know {ip}:{port}, asking it to become our peer")
|
|
self.send_hello(ip, port)
|
|
|
|
def send_hello(self, ip, port):
|
|
membership = self.context.membership()
|
|
sender = membership.self_record().to_dict()
|
|
known = [record.to_dict() for record in membership.alive_records() if record.key != membership.key]
|
|
|
|
try:
|
|
resp = self.context.execute_to_unregistered_node(
|
|
ip, port, Gossip.COMMANDS.HELLO, sender, known, timeout=setting("membership.hello_timeout")
|
|
)
|
|
except Exception as e:
|
|
self.logger.warning(f"Hello to {ip}:{port} failed: {e}")
|
|
return False
|
|
|
|
result = resp.get('result') if isinstance(resp, dict) and resp.get('status') == 'ok' else None
|
|
if not isinstance(result, dict) or not result.get('admitted'):
|
|
self.logger.warning(f"{ip}:{port} did not admit this node")
|
|
return False
|
|
|
|
try:
|
|
members = [MemberRecord.from_dict(item) for item in result.get('members', [])]
|
|
except (ValueError, TypeError) as e:
|
|
self.logger.warning(f"Ignoring bad member list from {ip}:{port}: {e}")
|
|
return False
|
|
membership.merge(members[:setting("membership.max_members")])
|
|
return True
|