Stuff

colony/core/context.py

@@ -16,6 +16,9 @@ class RoleContext:
     def peers(self):
         return self._network.get_all_nodes()
 
+    def membership(self):
+        return self._network.membership
+
     def self(self) -> LocalNode:
         return self._network.get_self()
 

colony/core/hardware_scores.py

@@ -0,0 +1,8 @@
+from colony.core.network.hardware import HardwareStatistics, HardwareScorer
+
+
+def collect_scores():
+    try:
+        return HardwareScorer.from_object(HardwareStatistics.collect())
+    except Exception:
+        return {}

colony/core/network/advice.py

@@ -0,0 +1,67 @@
+from colony.role_registry import HardwareType, RolePriority
+
+HARDWARE_SCORE_KEY = {
+    HardwareType.DISK_SPACE: "disk",
+    HardwareType.CPU: "cpu",
+    HardwareType.GPU: "gpu",
+    HardwareType.RAM: "memory",
+    HardwareType.LONGEVITY: "longevity",
+}
+
+
+def required_holders(role_stat):
+    return 1
+
+
+def assignable_roles(roles):
+    names = [name for name, stats in roles.items() if stats.priority != RolePriority.GENERAL]
+    return sorted(names, key=lambda name: (roles[name].priority.value, name))
+
+
+def claimants(records, role_name):
+    return [record for record in records if role_name in record.roles]
+
+
+def unmet_roles(records, roles):
+    return [
+        name for name in assignable_roles(roles)
+        if len(claimants(records, name)) < required_holders(roles[name])
+    ]
+
+
+def hardware_score(record, role_stat):
+    preferences = role_stat.hardware_preferences
+    if not preferences:
+        return 0.0
+    return sum(
+        record.scores.get(HARDWARE_SCORE_KEY[hardware], 0) * weight
+        for hardware, weight in preferences.priorities.items()
+        if hardware in HARDWARE_SCORE_KEY
+    )
+
+
+def pick_candidate(role_name, role_stat, records):
+    candidates = [
+        record for record in records
+        if role_name not in record.roles and role_name not in record.declined
+    ]
+    if not candidates:
+        return None
+    return min(candidates, key=lambda record: (-hardware_score(record, role_stat), record.key))
+
+
+def rank_of(key, records):
+    return sorted(record.key for record in records).index(key)
+
+
+def roles_to_release(self_key, records, roles):
+    release = []
+    for name in assignable_roles(roles):
+        holders = sorted(
+            claimants(records, name),
+            key=lambda record: (record.role_since.get(name, float("inf")), record.key),
+        )
+        keys = [record.key for record in holders]
+        if self_key in keys and keys.index(self_key) >= required_holders(roles[name]):
+            release.append(name)
+    return release

colony/core/network/core.py

@@ -2,6 +2,7 @@ import json
 from enum import Enum
 from typing import Optional
 
+from colony.core.network.membership import Membership
 from colony.core.network.nodes import LocalNode, Node
 from colony.core.network.requests import CommandRequest
 from colony.core.logger.logging import get_logger
@@ -11,21 +12,29 @@ from colony.utils.network import get_local_ip, get_hostname
 
 
 class Network:
-    def __init__(self):
+    def __init__(self, ip=None, port=None, hostname=None, membership=None):
         self.logger = get_logger(self.__class__.__name__)
 
         self.self_node = LocalNode(
-            hostname=get_hostname(),
+            hostname=hostname or get_hostname(),
-            ip=get_local_ip(),
+            ip=ip or get_local_ip(),
             roles=[],
-            port=configuration.listener.port
+            port=configuration.listener.port if port is None else port
         )
         self.logger.info(f"Self node created: {self.self_node.id}")
 
-        self.nodes = [self.self_node] + [
+        self.membership = membership or Membership(
-            node for node in self.load_network_from_snapshot()
+            self.self_node.ip, self.self_node.port, self.self_node.hostname
-            if (node.ip, node.port) != (self.self_node.ip, self.self_node.port)
+        )
+
+    @property
+    def nodes(self):
+        peers = [
+            Node(hostname=record.hostname, ip=record.ip, roles=list(record.roles), port=record.port)
+            for record in self.membership.alive_records()
+            if record.key != self.membership.key
         ]
+        return [self.self_node] + peers
 
     def get_self(self):
         return self.self_node

colony/core/network/membership.py

@@ -0,0 +1,206 @@
+import copy
+import ipaddress
+import threading
+import time
+from dataclasses import dataclass, field, asdict
+
+from colony.settings import setting
+
+ADMISSION_POLICIES = {
+    "open": lambda record: True,
+}
+
+MAX_TEXT = 255
+MAX_LIST = 64
+
+
+def _is_int(value):
+    return isinstance(value, int) and not isinstance(value, bool)
+
+
+def _is_number(value):
+    return isinstance(value, (int, float)) and not isinstance(value, bool)
+
+
+def _text(value, name):
+    if not isinstance(value, str) or len(value) > MAX_TEXT:
+        raise ValueError(f"Member record field {name} must be a short string")
+    return value
+
+
+def _text_list(value, name):
+    if not isinstance(value, list) or len(value) > MAX_LIST:
+        raise ValueError(f"Member record field {name} must be a short list")
+    return [_text(item, name) for item in value]
+
+
+@dataclass
+class MemberRecord:
+    ip: str
+    port: int
+    hostname: str
+    roles: list = field(default_factory=list)
+    role_since: dict = field(default_factory=dict)
+    declined: list = field(default_factory=list)
+    scores: dict = field(default_factory=dict)
+    incarnation: int = 0
+    seq: int = 0
+
+    @property
+    def key(self):
+        return (self.ip, self.port)
+
+    @property
+    def version(self):
+        return (self.incarnation, self.seq)
+
+    def to_dict(self):
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data):
+        if not isinstance(data, dict):
+            raise ValueError("Member record must be an object")
+
+        try:
+            ip = _text(data["ip"], "ip")
+            port = data["port"]
+            hostname = _text(data["hostname"], "hostname")
+            roles = _text_list(data.get("roles", []), "roles")
+            declined = _text_list(data.get("declined", []), "declined")
+            role_since = data.get("role_since", {})
+            scores = data.get("scores", {})
+            incarnation = data["incarnation"]
+            seq = data["seq"]
+        except KeyError as e:
+            raise ValueError(f"Member record has no {e}") from e
+
+        try:
+            ip = str(ipaddress.ip_address(ip))
+        except ValueError as e:
+            raise ValueError(f"Member record has a bad ip: {e}") from e
+
+        if not _is_int(port) or not 1 <= port <= 65535:
+            raise ValueError("Member record has a bad port")
+        if not _is_int(incarnation) or not _is_int(seq) or incarnation < 0 or seq < 0:
+            raise ValueError("Member record has a bad version")
+        if not isinstance(role_since, dict) or len(role_since) > MAX_LIST or \
+                not all(isinstance(k, str) and _is_int(v) for k, v in role_since.items()):
+            raise ValueError("Member record has bad role_since")
+        if not isinstance(scores, dict) or len(scores) > MAX_LIST or \
+                not all(isinstance(k, str) and _is_number(v) for k, v in scores.items()):
+            raise ValueError("Member record has bad scores")
+
+        return cls(ip, port, hostname, roles, dict(role_since), declined, dict(scores), incarnation, seq)
+
+
+def resolve_admission(admission):
+    if callable(admission):
+        return admission
+    try:
+        return ADMISSION_POLICIES[admission]
+    except KeyError:
+        raise ValueError(f"Unknown admission policy {admission!r}. Known: {sorted(ADMISSION_POLICIES)}") from None
+
+
+class Membership:
+    def __init__(self, ip, port, hostname, member_timeout=None, forget_after=None, max_members=None,
+                 admission=None, clock=time.monotonic, wall=time.time):
+        self.clock = clock
+        self.wall = wall
+        self.member_timeout = setting("membership.member_timeout") if member_timeout is None else member_timeout
+        self.forget_after = setting("membership.forget_after") if forget_after is None else forget_after
+        self.max_members = setting("membership.max_members") if max_members is None else max_members
+        self._admission = resolve_admission(setting("membership.admission") if admission is None else admission)
+        self.denied = False
+
+        self._lock = threading.RLock()
+        self._self = MemberRecord(ip, port, hostname, incarnation=int(wall() * 1000), seq=1)
+        self._records = {}
+        self._advanced = {}
+        self._declines = {}
+
+    @property
+    def key(self):
+        return self._self.key
+
+    def now(self):
+        return self.clock()
+
+    def admit(self, record):
+        return bool(self._admission(record))
+
+    def _current_self(self):
+        now = self.clock()
+        self._declines = {role: expiry for role, expiry in self._declines.items() if expiry > now}
+        record = copy.deepcopy(self._self)
+        record.declined = sorted(self._declines)
+        return record
+
+    def self_record(self):
+        with self._lock:
+            return self._current_self()
+
+    def update_self(self, roles=None, scores=None):
+        with self._lock:
+            if roles is not None:
+                now_ms = int(self.wall() * 1000)
+                self._self.role_since = {role: self._self.role_since.get(role, now_ms) for role in roles}
+                self._self.roles = list(roles)
+            if scores is not None:
+                self._self.scores = dict(scores)
+            self._self.seq += 1
+
+    def decline(self, role, cooldown):
+        with self._lock:
+            self._declines[role] = self.clock() + cooldown
+            self._self.seq += 1
+
+    def clear_decline(self, role):
+        with self._lock:
+            if self._declines.pop(role, None) is not None:
+                self._self.seq += 1
+
+    def is_declined(self, role):
+        with self._lock:
+            expiry = self._declines.get(role)
+            return expiry is not None and expiry > self.clock()
+
+    def merge(self, records):
+        changed = 0
+        with self._lock:
+            now = self.clock()
+            for record in records:
+                if record.key == self.key:
+                    continue
+                current = self._records.get(record.key)
+                if current is not None and record.version <= current.version:
+                    continue
+                if current is None and len(self._records) >= self.max_members:
+                    continue
+                self._records[record.key] = record
+                self._advanced[record.key] = now
+                changed += 1
+            self._purge(now)
+        return changed
+
+    def _purge(self, now):
+        for key in [key for key, advanced in self._advanced.items() if now - advanced > self.forget_after]:
+            del self._records[key]
+            del self._advanced[key]
+
+    def alive_records(self):
+        with self._lock:
+            now = self.clock()
+            alive = [self._current_self()]
+            alive += [record for key, record in self._records.items()
+                      if now - self._advanced[key] <= self.member_timeout]
+            return sorted(alive, key=lambda record: record.key)
+
+    def known_addresses(self):
+        with self._lock:
+            self._purge(self.clock())
+            return sorted(self._records)
+
+    def export(self):
+        return [record.to_dict() for record in self.alive_records()]

colony/core/roles/explorer/role.py

@@ -6,9 +6,9 @@ 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.settings import configuration, setting
 from colony.transport.discover import Discover
-from colony.utils.network import get_local_subnet
+from colony.utils.network import hosts_in_range
 
 
 class Explorer(Role):
@@ -112,11 +112,11 @@ class Explorer(Role):
         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()]
+        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)
+            return ip, Discover.request(ip, timeout=self.SCAN_TIMEOUT, discover_port=setting("discovery.discover_port"))
         except Exception:
             return ip, None
 

colony/core/roles/foreman/role.py

@@ -1,22 +1,14 @@
 from enum import Enum
 
-from colony.core.network.hardware import HardwareStatistics, HardwareScorer
+from colony.core.network import advice
 from colony.core.roles.role import Role
 from colony.core.roles.unit.role import Unit
-from colony.role_registry import HardwareType, RolePriority
+from colony.settings import setting
-
-_HARDWARE_SCORE_KEY = {
-    HardwareType.DISK_SPACE: "disk",
-    HardwareType.CPU: "cpu",
-    HardwareType.GPU: "gpu",
-    HardwareType.RAM: "memory",
-    HardwareType.LONGEVITY: "longevity",
-}
 
 
 class Foreman(Role):
     CONTINUOUS = True
-    LOOP_SECONDS_WAIT = 60
+    LOOP_SECONDS_WAIT = 5
 
     class COMMANDS(Enum):
         pass
@@ -26,138 +18,64 @@ class Foreman(Role):
 
         }
 
+        self.LOOP_SECONDS_WAIT = setting("foreman.interval")
         super().__init__(commands, context)
 
-    def _peers_for_assignment(self):
+        self.acquire_delay = setting("foreman.acquire_delay")
-        self_node = self.context.self()
+        self.proposal_stagger = setting("foreman.proposal_stagger")
-        by_key = {}
+        self.proposal_cooldown = setting("foreman.proposal_cooldown")
-        for p in self.context.peers():
-            by_key[(p.ip, p.port)] = p
-        k = (self_node.ip, self_node.port)
-        if k not in by_key:
-            by_key[k] = self_node
-        return [by_key[key] for key in sorted(by_key)]
 
-    def _call(self, peer, command, *args):
+        self._needed_since = {}
-        try:
+        self._last_proposal = {}
-            resp = peer.execute(command, *args)
-        except Exception as e:
-            self.logger.warning("%s failed for %s:%s: %s", command.value, peer.ip, peer.port, e)
-            return None
-
-        if not isinstance(resp, dict) or resp.get("status") != "ok":
-            self.logger.warning("%s failed for %s:%s", command.value, peer.ip, peer.port)
-            return None
-        return resp
-
-    def _fetch_live_roles(self, peers):
-        live = {}
-        for peer in peers:
-            resp = self._call(peer, Unit.COMMANDS.SHOW_ROLES)
-            if resp is None:
-                continue
-            if not isinstance(resp.get("result"), list):
-                self.logger.warning("Bad role list from %s:%s, leaving it out of this round", peer.ip, peer.port)
-                continue
-            live[peer] = list(resp["result"])
-        return live
-
-    def _assign_critical_holders(self, peers, critical_roles, live_roles):
-        peer_list = list(peers)
-        n = len(peer_list)
-        holders = {}
-        for i, critical_role in enumerate(critical_roles):
-            current = next(
-                (p for p in peer_list if critical_role in live_roles.get(p, ())),
-                None,
-            )
-            if current is not None:
-                holders[critical_role] = current
-            else:
-                holders[critical_role] = peer_list[i % n]
-        return holders
-
-    def _fetch_hardware_scores(self, peers):
-        scores = {}
-        for peer in peers:
-            resp = self._call(peer, Unit.COMMANDS.SHOW_HARDWARE)
-            if resp is None:
-                continue
-            hw_obj = HardwareStatistics.from_dict(resp["result"])
-            scores[peer] = HardwareScorer.from_object(hw_obj)
-        return scores
-
-    def _best_peer_for_optional(self, peers, role_stat, hardware_scores):
-        best_score = float("-inf")
-        best_peer = None
-        for peer in peers:
-            hs = hardware_scores.get(peer)
-            if hs is None:
-                continue
-            score = 0.0
-            if role_stat.hardware_preferences:
-                for hw_type, weight in role_stat.hardware_preferences.priorities.items():
-                    key = _HARDWARE_SCORE_KEY.get(hw_type)
-                    if key is not None:
-                        score += hs.get(key, 0) * weight
-            if score > best_score:
-                best_score = score
-                best_peer = peer
-        return best_peer
-
-    @staticmethod
-    def _dedupe_preserve_order(role_names):
-        seen = set()
-        out = []
-        for r in role_names:
-            if r not in seen:
-                seen.add(r)
-                out.append(r)
-        return out
 
     def loop(self):
-        role_registry = self.context.get_role_registry()
+        self.tick()
-        roles = role_registry.roles
 
-        general_roles = sorted(
+    def tick(self):
-            name for name, stats in roles.items() if stats.priority == RolePriority.GENERAL
+        membership = self.context.membership()
-        )
+        if membership.denied:
-        critical_roles = sorted(
+            return None
-            name for name, stats in roles.items() if stats.priority == RolePriority.CRITICAL
-        )
-        optional_roles = sorted(
-            name
-            for name, stats in roles.items()
-            if stats.priority not in (RolePriority.GENERAL, RolePriority.CRITICAL)
-        )
 
-        peers = self._peers_for_assignment()
+        now = membership.now()
-        if not peers:
+        records = membership.alive_records()
-            return
+        roles = self.context.get_role_registry().roles
 
-        live_roles = self._fetch_live_roles(peers)
+        unmet = advice.unmet_roles(records, roles)
-        peers = [peer for peer in peers if peer in live_roles]
+        self._needed_since = {name: self._needed_since.get(name, now) for name in unmet}
-        if not peers:
-            return
 
-        critical_holders = self._assign_critical_holders(
+        alone = len(records) == 1
-            peers, critical_roles, live_roles
+        rank = advice.rank_of(membership.key, records)
-        )
+        wait = 0 if alone else self.acquire_delay + min(rank, 3) * self.proposal_stagger
 
-        desired = {peer: list(general_roles) for peer in peers}
+        for role_name in unmet:
-        for critical_role, holder in critical_holders.items():
+            if now - self._needed_since[role_name] < wait:
-            desired[holder].append(critical_role)
+                continue
 
-        hardware_scores = self._fetch_hardware_scores(peers) if optional_roles else {}
+            target = advice.pick_candidate(role_name, roles[role_name], records)
-        for role_name in optional_roles:
+            if target is None:
-            role_stat = roles[role_name]
+                self.logger.warning(f"No node is willing to take role {role_name}")
-            best_peer = self._best_peer_for_optional(peers, role_stat, hardware_scores)
+                continue
-            if best_peer is not None:
-                desired[best_peer].append(role_name)
 
-        for peer in peers:
+            last = self._last_proposal.get((role_name, target.key))
-            self._call(
+            if last is not None and now - last < self.proposal_cooldown:
-                peer,
+                continue
-                Unit.COMMANDS.SET_ROLES,
+
-                self._dedupe_preserve_order(desired[peer]),
+            self._last_proposal[(role_name, target.key)] = now
+            return self.propose(role_name, target)
+
+        return None
+
+    def propose(self, role_name, target):
+        self.logger.info(f"Asking {target.ip}:{target.port} to take role {role_name}")
+        try:
+            resp = self.context.execute_to_unregistered_node(
+                target.ip, target.port, Unit.COMMANDS.REQUEST_ROLE, role_name
             )
+        except Exception as e:
+            self.logger.warning(f"Could not ask {target.ip}:{target.port} for role {role_name}: {e}")
+            return (role_name, target.key, None)
+
+        result = resp.get("result") if isinstance(resp, dict) and resp.get("status") == "ok" else None
+        accepted = result.get("accepted") if isinstance(result, dict) else None
+        self.logger.info(f"{target.ip}:{target.port} answered {role_name}: {result}")
+        return (role_name, target.key, accepted)

colony/core/roles/gossip/role.py

@@ -0,0 +1,237 @@
+import random
+from concurrent.futures import ThreadPoolExecutor
+from enum import Enum
+
+from colony.core.hardware_scores import collect_scores
+from colony.core.network import advice
+from colony.core.network.membership import MemberRecord
+from colony.core.roles.role import Role, Command
+from colony.core.state_manager import StateManager
+from colony.settings import setting
+from colony.transport.discover import Discover
+from colony.utils.network import hosts_in_range
+
+ADMITTED = "admitted"
+REFUSED = "refused"
+FAILED = "failed"
+
+JOINED = "joined"
+BOOTSTRAPPED = "bootstrapped"
+DENIED = "denied"
+
+
+class Gossip(Role):
+    CONTINUOUS = True
+    LOOP_SECONDS_WAIT = 5
+
+    BOOTSTRAP_REPROBE = 5
+    DENIED_BACKOFF_CAP = 60
+    SCORES_TTL = 30
+
+    class COMMANDS(Enum):
+        HELLO = 'Gossip.Hello'
+
+    def __init__(self, context):
+        commands = {
+            self.COMMANDS.HELLO.value: Command(self.hello),
+        }
+
+        self.LOOP_SECONDS_WAIT = setting("membership.gossip_interval")
+        super().__init__(commands, context)
+
+        self.reprobe_interval = setting("membership.reprobe_interval")
+        self.fanout = setting("membership.fanout")
+        self.max_members = setting("membership.max_members")
+        self.discover_port = setting("discovery.discover_port")
+        self.probe_timeout = setting("discovery.probe_timeout")
+        self.probe_workers = setting("discovery.probe_workers")
+        self.range = (setting("discovery.range_start"), setting("discovery.range_end"))
+
+        self.random = random.Random()
+        self.next_reprobe = 0.0
+        self.denied_attempts = 0
+        self.outcome = None
+        self._scores = {}
+        self._scores_at = None
+        self._persisted_peers = None
+
+    @property
+    def membership(self):
+        return self.context.membership()
+
+    def loop(self):
+        self.tick()
+
+    def hello(self, sender, known):
+        record = MemberRecord.from_dict(sender)
+        if not self.membership.admit(record):
+            return {"admitted": False, "members": []}
+
+        records = self._parse_records(known)
+        self.membership.merge([record] + records)
+        return {"admitted": True, "members": self.membership.export()}
+
+    def _parse_records(self, raw):
+        if not isinstance(raw, list) or len(raw) > self.max_members:
+            raise ValueError("Member list must be a list of bounded size")
+        return [MemberRecord.from_dict(item) for item in raw]
+
+    def scan_targets(self):
+        own_ip = self.context.self().ip
+        return [(host, self.discover_port) for host in hosts_in_range(*self.range, exclude=[own_ip])]
+
+    def discover_target(self, ip, port):
+        return (ip, self.discover_port)
+
+    def seed_addresses(self):
+        state = self.context.get_role_manager().state_manager
+        seeds = []
+        for entry in state.get(StateManager.Data.Peers):
+            if isinstance(entry, list) and len(entry) == 2 and isinstance(entry[0], str) and isinstance(entry[1], int):
+                seeds.append((entry[0], entry[1]))
+        return seeds
+
+    def probe_targets(self):
+        membership = self.membership
+        alive = {record.key for record in membership.alive_records()}
+        remembered = (set(membership.known_addresses()) | set(self.seed_addresses())) - alive
+
+        targets = list(self.scan_targets())
+        for ip, port in sorted(remembered):
+            target = self.discover_target(ip, port)
+            if target not in targets:
+                targets.append(target)
+        return targets
+
+    def _probe(self, target):
+        host, discover_port = target
+        try:
+            return host, Discover.request(host, timeout=self.probe_timeout, discover_port=discover_port)
+        except Exception:
+            return host, None
+
+    def scan(self, targets):
+        if not targets:
+            return []
+        with ThreadPoolExecutor(max_workers=min(self.probe_workers, len(targets))) as pool:
+            results = list(pool.map(self._probe, targets))
+        return [(host, port) for host, port in results if port is not None]
+
+    def find_peers(self):
+        membership = self.membership
+        alive = {record.key for record in membership.alive_records()}
+        found = self.scan(self.probe_targets())
+        return [address for address in dict.fromkeys(found) if address not in alive]
+
+    def exchange(self, ip, port):
+        membership = self.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, self.COMMANDS.HELLO, sender, known)
+        except Exception as e:
+            self.logger.debug(f"Hello to {ip}:{port} failed: {e}")
+            return FAILED
+
+        result = resp.get("result") if isinstance(resp, dict) and resp.get("status") == "ok" else None
+        if not isinstance(result, dict):
+            self.logger.debug(f"Hello to {ip}:{port} was answered with an error")
+            return FAILED
+        if not result.get("admitted"):
+            self.logger.warning(f"{ip}:{port} refused to admit this node")
+            return REFUSED
+
+        try:
+            membership.merge(self._parse_records(result.get("members", [])))
+        except ValueError as e:
+            self.logger.warning(f"Ignoring bad member list from {ip}:{port}: {e}")
+            return FAILED
+        return ADMITTED
+
+    def scores(self):
+        now = self.membership.now()
+        if self._scores_at is None or now - self._scores_at >= self.SCORES_TTL:
+            self._scores = self.collect_scores()
+            self._scores_at = now
+        return self._scores
+
+    def collect_scores(self):
+        return collect_scores()
+
+    def refresh_self(self):
+        self.membership.update_self(roles=list(self.context.self().roles), scores=self.scores())
+
+    def _after_probe(self, results, now):
+        membership = self.membership
+        has_peers = any(record.key != membership.key for record in membership.alive_records())
+
+        if ADMITTED in results or has_peers:
+            outcome, delay = JOINED, self.reprobe_interval
+            membership.denied = False
+            self.denied_attempts = 0
+        elif REFUSED in results:
+            self.denied_attempts += 1
+            outcome = DENIED
+            delay = min(self.BOOTSTRAP_REPROBE * 2 ** self.denied_attempts, self.DENIED_BACKOFF_CAP)
+            membership.denied = True
+        else:
+            outcome, delay = BOOTSTRAPPED, self.BOOTSTRAP_REPROBE
+            membership.denied = False
+
+        self.next_reprobe = now + delay
+        self.outcome = outcome
+        return outcome
+
+    def join_or_bootstrap(self):
+        self.refresh_self()
+        results = [self.exchange(ip, port) for ip, port in self.find_peers()]
+        outcome = self._after_probe(results, self.membership.now())
+        self.logger.info(f"Startup probe finished: {outcome}")
+        return outcome
+
+    def reprobe(self):
+        results = [self.exchange(ip, port) for ip, port in self.find_peers()]
+        self._after_probe(results, self.membership.now())
+
+    def tick(self):
+        membership = self.membership
+        self.refresh_self()
+
+        peers = [record for record in membership.alive_records() if record.key != membership.key]
+        for record in self.random.sample(peers, min(self.fanout, len(peers))):
+            self.exchange(*record.key)
+
+        if membership.now() >= self.next_reprobe:
+            self.reprobe()
+
+        self.persist_peers()
+        self.release_roles()
+
+    def persist_peers(self):
+        addresses = self.membership.known_addresses()
+        if addresses == self._persisted_peers:
+            return
+        state = self.context.get_role_manager().state_manager
+        state.set(StateManager.Data.Peers, [list(address) for address in addresses])
+        self._persisted_peers = addresses
+
+    def release_roles(self):
+        manager = self.context.get_role_manager()
+        membership = self.membership
+        roles = self.context.get_role_registry().roles
+        policy = manager.role_policy
+
+        held = list(manager.get_role_instances())
+        scores = membership.self_record().scores
+        assignable = advice.assignable_roles(roles)
+        forbidden = [name for name in held if name in assignable and not policy.should_keep(name, scores)]
+        over_held = advice.roles_to_release(membership.key, membership.alive_records(), roles)
+
+        for name in forbidden:
+            membership.decline(name, policy.decline_cooldown)
+
+        drop = set(forbidden) | set(over_held)
+        if drop:
+            self.logger.info(f"Releasing roles {sorted(drop)}")
+            manager.set_roles([name for name in held if name not in drop])

colony/core/roles/manager.py

@@ -1,19 +1,32 @@
 import threading
 
 from colony.core.context import RoleContext
+from colony.core.roles.policy import DefaultRolePolicy
 from colony.core.state_manager import StateManager
+from colony.role_registry import RolePriority
 from colony.core.logger.logging import get_logger
 
 
 class RoleManager:
-    def __init__(self, network, role_registry, state_manager: StateManager):
+    def __init__(self, network, role_registry, state_manager: StateManager, role_policy=None):
         self.network = network
         self.role_registry = role_registry
         self.state_manager = state_manager
+        self.role_policy = role_policy or DefaultRolePolicy()
 
         self.logger = get_logger(self.__class__.__name__)
         self._lock = threading.RLock()
 
+    def ensure_general_roles(self):
+        general = sorted(
+            name for name, stats in self.role_registry.roles.items()
+            if stats.priority == RolePriority.GENERAL
+        )
+        current = self.state_manager.get(StateManager.Data.Roles)
+        wanted = list(dict.fromkeys(general + current))
+        if wanted != current:
+            self.state_manager.set(StateManager.Data.Roles, wanted)
+
     def initialize_roles(self):
         self.logger.info("Initializing roles...")
         context = RoleContext(self.network, self)
@@ -31,6 +44,7 @@ class RoleManager:
 
             self.network.self_node.roles = list(roles)
             self.network.self_node.set_role_instances(roles)
+            self.network.membership.update_self(roles=list(roles))
         self.logger.info(f"Roles initialized: {list(roles.keys())}. ")
 
     def get_role_instances(self):
@@ -80,6 +94,7 @@ class RoleManager:
             self.state_manager.set(StateManager.Data.Roles, list(final_roles.keys()))
             self.network.self_node.set_role_instances(final_roles)
             self.network.self_node.roles = list(final_roles.keys())
+            self.network.membership.update_self(roles=list(final_roles))
 
             for role_name in roles_to_remove:
                 self.logger.debug(f"Stopping removed role {role_name}")
@@ -90,3 +105,28 @@ class RoleManager:
                 role.start()
 
         self.logger.info(f"Roles updated to: {list(final_roles.keys())}")
+
+    def request_role(self, role_name):
+        membership = self.network.membership
+
+        with self._lock:
+            if self.role_registry.get_role_class_by_name(role_name) is None:
+                return {"accepted": False, "reason": "unknown role"}
+
+            held = list(self.get_role_instances())
+            if role_name in held:
+                return {"accepted": True, "reason": "already held"}
+
+            if membership.is_declined(role_name):
+                return {"accepted": False, "reason": "recently declined"}
+
+            scores = membership.self_record().scores
+            accepted, reason = self.role_policy.decide(role_name, held, scores, self.role_registry.roles)
+            if not accepted:
+                self.logger.info(f"Declining role {role_name}: {reason}")
+                membership.decline(role_name, self.role_policy.decline_cooldown)
+                return {"accepted": False, "reason": reason}
+
+            self.set_roles(held + [role_name])
+            membership.clear_decline(role_name)
+            return {"accepted": True, "reason": reason}

colony/core/roles/policy.py

@@ -0,0 +1,44 @@
+from colony.role_registry import RolePriority
+from colony.settings import setting
+
+
+class RolePolicy:
+    decline_cooldown = 60
+
+    def decide(self, role_name, held, scores, roles):
+        raise NotImplementedError
+
+    def should_keep(self, role_name, scores):
+        return True
+
+
+class DefaultRolePolicy(RolePolicy):
+    def __init__(self, max_roles=None, deny=None, min_scores=None, decline_cooldown=None):
+        self.max_roles = setting("role_policy.max_roles") if max_roles is None else max_roles
+        self.deny = set(setting("role_policy.deny") if deny is None else deny)
+        raw_minimums = setting("role_policy.min_scores") if min_scores is None else min_scores
+        self.min_scores = {role: dict(minimums) for role, minimums in dict(raw_minimums).items()}
+        self.decline_cooldown = setting("role_policy.decline_cooldown") if decline_cooldown is None else decline_cooldown
+
+    def _violation(self, role_name, scores):
+        if role_name in self.deny:
+            return "role is on the deny list"
+        for key, minimum in self.min_scores.get(role_name, {}).items():
+            if scores.get(key, 0) < minimum:
+                return f"{key} score is below {minimum}"
+        return None
+
+    def decide(self, role_name, held, scores, roles):
+        reason = self._violation(role_name, scores)
+        if reason:
+            return False, reason
+
+        if self.max_roles is not None:
+            assigned = [name for name in held if name in roles and roles[name].priority != RolePriority.GENERAL]
+            if len(assigned) >= self.max_roles:
+                return False, "node already holds the maximum number of roles"
+
+        return True, "accepted"
+
+    def should_keep(self, role_name, scores):
+        return self._violation(role_name, scores) is None

colony/core/roles/unit/role.py

@@ -8,14 +8,14 @@ class Unit(Role):
         SHOW_STATS = 'Unit.ShowStats'
         SHOW_ROLES = 'Unit.ShowRoles'
         SHOW_HARDWARE = 'Unit.ShowHardware'
-        SET_ROLES = 'Unit.SetRoles'
+        REQUEST_ROLE = 'Unit.RequestRole'
 
     def __init__(self, context):
         commands = {
             self.COMMANDS.SHOW_STATS.value: Command(self.show_stats),
             self.COMMANDS.SHOW_ROLES.value: Command(self.show_roles),
             self.COMMANDS.SHOW_HARDWARE.value: Command(self.show_hardware),
-            self.COMMANDS.SET_ROLES.value: Command(self.set_roles),
+            self.COMMANDS.REQUEST_ROLE.value: Command(self.request_role),
         }
 
         super().__init__(commands, context)
@@ -30,5 +30,5 @@ class Unit(Role):
         self_node = self.context.self()
         return self_node.to_dict()
 
-    def set_roles(self, roles):
+    def request_role(self, role):
-        self.context.set_self_node_roles(roles)
+        return self.context.get_role_manager().request_role(role)

colony/core/state_manager.py

@@ -13,14 +13,17 @@ from colony.utils.files import atomic_write_json
 class StateManager:
     class Data(Enum):
         Roles = "roles"
+        Peers = "peers"
 
     DEFAULTS = {
         Data.Roles: [],
+        Data.Peers: [],
     }
 
     STATE_FILE = Path("state.json")
 
-    def __init__(self):
+    def __init__(self, path=None):
+        self._path = Path(path) if path is not None else self.STATE_FILE
         self.logger = get_logger(self.__class__.__name__)
         self._lock = threading.RLock()
         self._data = {}
@@ -28,9 +31,9 @@ class StateManager:
 
     def _load(self):
         with self._lock:
-            if self.STATE_FILE.exists():
+            if self._path.exists():
                 try:
-                    with open(self.STATE_FILE, "r", encoding="utf-8") as f:
+                    with open(self._path, "r", encoding="utf-8") as f:
                         data = json.load(f)
                     if not isinstance(data, dict):
                         raise ValueError("top level is not an object")
@@ -48,12 +51,12 @@ class StateManager:
             self._apply_defaults()
 
     def _recover_from_corrupt_file(self, error):
-        backup = self.STATE_FILE.with_name(self.STATE_FILE.name + ".corrupt")
+        backup = self._path.with_name(self._path.name + ".corrupt")
-        self.logger.warning(f"Could not use {self.STATE_FILE} ({error}). Keeping a copy as {backup} and resetting to defaults")
+        self.logger.warning(f"Could not use {self._path} ({error}). Keeping a copy as {backup} and resetting to defaults")
         try:
-            shutil.copy2(self.STATE_FILE, backup)
+            shutil.copy2(self._path, backup)
         except OSError as copy_error:
-            self.logger.warning(f"Could not back up {self.STATE_FILE}: {copy_error}")
+            self.logger.warning(f"Could not back up {self._path}: {copy_error}")
         self._data = {}
 
     def _apply_defaults(self):
@@ -61,7 +64,7 @@ class StateManager:
             self._data.setdefault(key.value, copy.deepcopy(default))
 
     def _save(self):
-        atomic_write_json(self.STATE_FILE, self._data, indent=4)
+        atomic_write_json(self._path, self._data, indent=4)
 
     def get(self, key: Data) -> Any:
         with self._lock:

colony/settings.py

@@ -5,3 +5,30 @@ configuration = Dynaconf(
     settings_files=['config.json'],
     load_dotenv=False,
 )
+
+DEFAULTS = {
+    "discovery.range_start": "192.168.65.110",
+    "discovery.range_end": "192.168.65.120",
+    "discovery.discover_port": 28149,
+    "discovery.probe_timeout": 1.0,
+    "discovery.probe_workers": 64,
+    "membership.gossip_interval": 5,
+    "membership.member_timeout": 20,
+    "membership.reprobe_interval": 15,
+    "membership.fanout": 3,
+    "membership.forget_after": 600,
+    "membership.max_members": 1024,
+    "membership.admission": "open",
+    "foreman.interval": 5,
+    "foreman.acquire_delay": 10,
+    "foreman.proposal_stagger": 2,
+    "foreman.proposal_cooldown": 30,
+    "role_policy.max_roles": None,
+    "role_policy.deny": [],
+    "role_policy.min_scores": {},
+    "role_policy.decline_cooldown": 60,
+}
+
+
+def setting(key):
+    return configuration.get(key, DEFAULTS[key])

colony/transport/discover.py

@@ -11,12 +11,12 @@ class Discover:
     CLIENT_TIMEOUT = 2
 
     @staticmethod
-    def listen(real_port, host="0.0.0.0"):
+    def listen(real_port, host="0.0.0.0", discover_port=None):
         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.bind((host, discover_port or Discover.DISCOVER_PORT))
         s.listen()
 
         while True:
@@ -36,8 +36,8 @@ class Discover:
                     logger.debug(f"Discovery request failed: {e}")
 
     @staticmethod
-    def request(target_host, timeout=3):
+    def request(target_host, timeout=3, discover_port=None):
-        with socket.create_connection((target_host, Discover.DISCOVER_PORT), timeout=timeout) as s:
+        with socket.create_connection((target_host, discover_port or Discover.DISCOVER_PORT), timeout=timeout) as s:
             s.sendall(Discover.MAGIC)
 
             data = b""

colony/utils/files.py

@@ -1,6 +1,6 @@
 import json
 import os
-import stat
+import shutil
 import uuid
 
 
@@ -14,7 +14,7 @@ def atomic_write_json(path, data, indent=2):
             f.flush()
             os.fsync(f.fileno())
         if os.path.exists(path):
-            os.chmod(tmp_path, stat.S_IMODE(os.stat(path).st_mode))
+            shutil.copymode(path, tmp_path)
         os.replace(tmp_path, path)
     except BaseException:
         os.remove(tmp_path)

colony/utils/network.py

@@ -1,6 +1,8 @@
 import ipaddress
 import socket
 
+MAX_RANGE_SIZE = 1024
+
 
 def get_local_subnet():
     return ipaddress.ip_network(f"{get_local_ip()}/24", strict=False)
@@ -20,3 +22,19 @@ def get_local_ip():
 
 def get_hostname():
     return socket.gethostname()
+
+def hosts_in_range(start, end, exclude=()):
+    try:
+        first = ipaddress.IPv4Address(start)
+        last = ipaddress.IPv4Address(end)
+    except ValueError as e:
+        raise ValueError(f"Invalid discovery range {start!r}-{end!r}: {e}") from e
+
+    if first > last:
+        raise ValueError(f"Invalid discovery range: {start} is after {end}")
+    if int(last) - int(first) + 1 > MAX_RANGE_SIZE:
+        raise ValueError(f"Discovery range {start}-{end} has more than {MAX_RANGE_SIZE} addresses")
+
+    excluded = {str(host) for host in exclude}
+    hosts = (str(ipaddress.IPv4Address(value)) for value in range(int(first), int(last) + 1))
+    return [host for host in hosts if host not in excluded]

config.json

@@ -10,5 +10,33 @@
     "port": 65334,
     "cert_file": "certs/server.crt",
     "keyfile": "certs/server.key"
+  },
+  "discovery": {
+    "range_start": "192.168.65.110",
+    "range_end": "192.168.65.120",
+    "discover_port": 28149,
+    "probe_timeout": 1.0,
+    "probe_workers": 64
+  },
+  "membership": {
+    "gossip_interval": 5,
+    "member_timeout": 20,
+    "reprobe_interval": 15,
+    "fanout": 3,
+    "forget_after": 600,
+    "max_members": 1024,
+    "admission": "open"
+  },
+  "foreman": {
+    "interval": 5,
+    "acquire_delay": 10,
+    "proposal_stagger": 2,
+    "proposal_cooldown": 30
+  },
+  "role_policy": {
+    "max_roles": null,
+    "deny": [],
+    "min_scores": {},
+    "decline_cooldown": 60
   }
 }

main.py

@@ -4,23 +4,32 @@ import threading
 from colony.core.logger.logging import get_logger
 from colony.core.network.core import Network
 from colony.core.dispatcher import Dispatcher
+from colony.core.roles.foreman.role import Foreman
+from colony.core.roles.gossip.role import Gossip, BOOTSTRAPPED
 from colony.core.roles.manager import RoleManager
 from colony.core.state_manager import StateManager
 from colony.role_registry import RoleRegistry
+from colony.settings import configuration, setting
 from colony.transport.discover import Discover
 from colony.transport.listener import Listener
+from colony.utils.network import hosts_in_range
 from role_settings import available_roles
-from colony.settings import configuration
 
 
 def main():
     logger = get_logger("main")
+
+    try:
+        hosts_in_range(setting("discovery.range_start"), setting("discovery.range_end"))
+    except ValueError as e:
+        logger.error(str(e))
+        sys.exit(1)
+
     state_manager = StateManager()
-
-    if 'Unit' not in state_manager.get(StateManager.Data.Roles):
-        state_manager.set(StateManager.Data.Roles, ['Unit', 'Foreman'])
-
     network = Network()
+    if network.self_node.ip.startswith("127."):
+        logger.warning(f"This node's address is {network.self_node.ip}, other nodes will not be able to reach it")
+
     role_registry = RoleRegistry(role_list=available_roles)
     dispatcher = Dispatcher(network=network, role_registry=role_registry)
     role_manager = RoleManager(
@@ -29,6 +38,7 @@ def main():
         state_manager=state_manager,
     )
 
+    role_manager.ensure_general_roles()
     role_manager.initialize_roles()
 
     listener = Listener(dispatcher=dispatcher)
@@ -50,13 +60,24 @@ def main():
 
     t = threading.Thread(
         target=Discover.listen,
-        args=(configuration.listener.port,),
+        args=(configuration.listener.port, "0.0.0.0", setting("discovery.discover_port")),
         daemon=True,
         name="Discover"
     )
     t.start()
     threads.append(t)
 
+    roles = role_manager.get_role_instances()
+    gossip: Gossip = roles["Gossip"]
+    outcome = gossip.join_or_bootstrap()
+    gossip.release_roles()
+
+    if outcome == BOOTSTRAPPED:
+        foreman: Foreman = roles["Foreman"]
+        for _ in role_registry.roles:
+            if foreman.tick() is None:
+                break
+
     role_manager.start_all()
 
     for t in threads:

role_settings.py

@@ -2,6 +2,7 @@ from colony.core.roles.explorer.role import Explorer
 from colony.core.roles.unit.role import Unit
 from colony.role_registry import RoleStats, HardwarePreferences, RolePriority, HardwareType
 from colony.core.roles.foreman.role import Foreman
+from colony.core.roles.gossip.role import Gossip
 from colony.core.roles.vaultwarden.role import VaultWarden
 
 available_roles = [
@@ -9,9 +10,13 @@ available_roles = [
         cls=Unit,
         priority=RolePriority.GENERAL,
     ),
+    RoleStats(
+        cls=Gossip,
+        priority=RolePriority.GENERAL,
+    ),
     RoleStats(
         cls=Foreman,
-        priority=RolePriority.CRITICAL,
+        priority=RolePriority.GENERAL,
     ),
     RoleStats(
         cls=VaultWarden,

tests/helpers.py

@@ -0,0 +1,36 @@
+import socket
+import threading
+from pathlib import Path
+
+from colony.core.roles.role import Role
+from colony.transport.listener import Listener
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+CERT = str(REPO_ROOT / "certs" / "server.crt")
+KEY = str(REPO_ROOT / "certs" / "server.key")
+
+
+def free_port():
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def start_listener(dispatcher):
+    listener = Listener(dispatcher, host="127.0.0.1", port=0, cert_file=CERT, keyfile=KEY)
+    listener.bind()
+    threading.Thread(target=listener.start, daemon=True).start()
+    return listener
+
+
+class StubRole(Role):
+    def __init__(self, context):
+        super().__init__({}, context)
+
+
+class RoleA(StubRole):
+    pass
+
+
+class RoleB(StubRole):
+    pass

tests/test_tls_roundtrip.py

@@ -0,0 +1,89 @@
+import json
+import socket
+import ssl
+import unittest
+
+from colony.transport.framing import send_frame, recv_frame
+from colony.transport.listener import Listener
+from colony.transport.whisperer import Whisperer
+from tests.helpers import CERT, KEY, start_listener
+
+PAYLOAD = {"command": {"role": "X", "name": "y"}}
+
+
+class StubDispatcher:
+    def __init__(self, result=None, error=None):
+        self.result, self.error = result, error
+
+    def dispatch(self, payload):
+        if self.error:
+            raise self.error
+        return self.result
+
+
+def whisper(listener, **kwargs):
+    return Whisperer("127.0.0.1", listener.port, CERT, KEY).send(PAYLOAD, **kwargs)
+
+
+class TlsTests(unittest.TestCase):
+    def test_mutual_tls_round_trip(self):
+        listener = start_listener(StubDispatcher(result={"x": 1}))
+        self.assertEqual(whisper(listener), {"status": "ok", "result": {"x": 1}})
+
+    def test_large_payload_round_trip(self):
+        big = {"data": "y" * 500_000}
+        listener = start_listener(StubDispatcher(result=big))
+        self.assertEqual(whisper(listener)["result"], big)
+
+    def test_client_without_certificate_is_rejected(self):
+        listener = start_listener(StubDispatcher(result={"x": 1}))
+
+        context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+        context.check_hostname = False
+        context.load_verify_locations(cafile=CERT)
+        with self.assertRaises((ssl.SSLError, OSError)):
+            with socket.create_connection(("127.0.0.1", listener.port), timeout=5) as sock:
+                with context.wrap_socket(sock) as tls:
+                    send_frame(tls, json.dumps(PAYLOAD).encode())
+                    recv_frame(tls)
+
+        self.assertEqual(whisper(listener)["status"], "ok")
+
+    def test_plain_tcp_client_does_not_break_the_listener(self):
+        listener = start_listener(StubDispatcher(result=1))
+        with socket.create_connection(("127.0.0.1", listener.port)) as sock:
+            sock.sendall(b"GET / HTTP/1.0\r\n\r\n")
+        self.assertEqual(whisper(listener)["result"], 1)
+
+    def test_errors_do_not_leak_details(self):
+        listener = start_listener(StubDispatcher(error=RuntimeError("secret-path /etc/x")))
+        with self.assertLogs("Listener", level="ERROR"):
+            response = whisper(listener)
+        self.assertEqual(response, {"status": "error", "error": "RuntimeError"})
+        self.assertNotIn("secret", json.dumps(response))
+        listener.dispatcher.error = None
+        listener.dispatcher.result = 2
+        self.assertEqual(whisper(listener)["result"], 2)
+
+    def test_unserializable_result_becomes_an_error(self):
+        listener = start_listener(StubDispatcher(result={"bad": object()}))
+        with self.assertLogs("Listener", level="ERROR"):
+            self.assertEqual(whisper(listener)["status"], "error")
+
+    def test_silent_server_times_out(self):
+        server = socket.socket()
+        self.addCleanup(server.close)
+        server.bind(("127.0.0.1", 0))
+        server.listen()
+        with self.assertRaises(TimeoutError):
+            Whisperer("127.0.0.1", server.getsockname()[1], CERT, KEY).send(PAYLOAD, timeout=0.5)
+
+    def test_bind_on_used_port_raises(self):
+        first = start_listener(StubDispatcher())
+        second = Listener(StubDispatcher(), host="127.0.0.1", port=first.port, cert_file=CERT, keyfile=KEY)
+        with self.assertRaises(OSError):
+            second.bind()
+
+
+if __name__ == "__main__":
+    unittest.main()