Stuff
This commit is contained in:
1 parent
b954d9d448
commit
01f4c13954
22 files changed
+929
-170
No files matched your search
@@ -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()
|
||||
|
||||
|
||||
@@ -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 {}
|
||||
@@ -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
|
||||
@@ -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(),
|
||||
ip=get_local_ip(),
|
||||
hostname=hostname or get_hostname(),
|
||||
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] + [
|
||||
node for node in self.load_network_from_snapshot()
|
||||
if (node.ip, node.port) != (self.self_node.ip, self.self_node.port)
|
||||
self.membership = membership or Membership(
|
||||
self.self_node.ip, self.self_node.port, self.self_node.hostname
|
||||
)
|
||||
|
||||
@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
|
||||
|
||||
@@ -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()]
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
_HARDWARE_SCORE_KEY = {
|
||||
HardwareType.DISK_SPACE: "disk",
|
||||
HardwareType.CPU: "cpu",
|
||||
HardwareType.GPU: "gpu",
|
||||
HardwareType.RAM: "memory",
|
||||
HardwareType.LONGEVITY: "longevity",
|
||||
}
|
||||
from colony.settings import setting
|
||||
|
||||
|
||||
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_node = self.context.self()
|
||||
by_key = {}
|
||||
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)]
|
||||
self.acquire_delay = setting("foreman.acquire_delay")
|
||||
self.proposal_stagger = setting("foreman.proposal_stagger")
|
||||
self.proposal_cooldown = setting("foreman.proposal_cooldown")
|
||||
|
||||
def _call(self, peer, command, *args):
|
||||
try:
|
||||
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
|
||||
self._needed_since = {}
|
||||
self._last_proposal = {}
|
||||
|
||||
def loop(self):
|
||||
role_registry = self.context.get_role_registry()
|
||||
roles = role_registry.roles
|
||||
self.tick()
|
||||
|
||||
general_roles = sorted(
|
||||
name for name, stats in roles.items() if stats.priority == RolePriority.GENERAL
|
||||
)
|
||||
critical_roles = sorted(
|
||||
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)
|
||||
)
|
||||
def tick(self):
|
||||
membership = self.context.membership()
|
||||
if membership.denied:
|
||||
return None
|
||||
|
||||
peers = self._peers_for_assignment()
|
||||
if not peers:
|
||||
return
|
||||
now = membership.now()
|
||||
records = membership.alive_records()
|
||||
roles = self.context.get_role_registry().roles
|
||||
|
||||
live_roles = self._fetch_live_roles(peers)
|
||||
peers = [peer for peer in peers if peer in live_roles]
|
||||
if not peers:
|
||||
return
|
||||
unmet = advice.unmet_roles(records, roles)
|
||||
self._needed_since = {name: self._needed_since.get(name, now) for name in unmet}
|
||||
|
||||
critical_holders = self._assign_critical_holders(
|
||||
peers, critical_roles, live_roles
|
||||
)
|
||||
alone = len(records) == 1
|
||||
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 critical_role, holder in critical_holders.items():
|
||||
desired[holder].append(critical_role)
|
||||
for role_name in unmet:
|
||||
if now - self._needed_since[role_name] < wait:
|
||||
continue
|
||||
|
||||
hardware_scores = self._fetch_hardware_scores(peers) if optional_roles else {}
|
||||
for role_name in optional_roles:
|
||||
role_stat = roles[role_name]
|
||||
best_peer = self._best_peer_for_optional(peers, role_stat, hardware_scores)
|
||||
if best_peer is not None:
|
||||
desired[best_peer].append(role_name)
|
||||
target = advice.pick_candidate(role_name, roles[role_name], records)
|
||||
if target is None:
|
||||
self.logger.warning(f"No node is willing to take role {role_name}")
|
||||
continue
|
||||
|
||||
for peer in peers:
|
||||
self._call(
|
||||
peer,
|
||||
Unit.COMMANDS.SET_ROLES,
|
||||
self._dedupe_preserve_order(desired[peer]),
|
||||
last = self._last_proposal.get((role_name, target.key))
|
||||
if last is not None and now - last < self.proposal_cooldown:
|
||||
continue
|
||||
|
||||
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)
|
||||
Whitespace-only changes.
@@ -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])
|
||||
@@ -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}
|
||||
@@ -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
|
||||
@@ -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):
|
||||
self.context.set_self_node_roles(roles)
|
||||
def request_role(self, role):
|
||||
return self.context.get_role_manager().request_role(role)
|
||||
@@ -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")
|
||||
self.logger.warning(f"Could not use {self.STATE_FILE} ({error}). Keeping a copy as {backup} and resetting to defaults")
|
||||
backup = self._path.with_name(self._path.name + ".corrupt")
|
||||
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:
|
||||
|
||||
@@ -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])
|
||||
@@ -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):
|
||||
with socket.create_connection((target_host, Discover.DISCOVER_PORT), timeout=timeout) as s:
|
||||
def request(target_host, timeout=3, discover_port=None):
|
||||
with socket.create_connection((target_host, discover_port or Discover.DISCOVER_PORT), timeout=timeout) as s:
|
||||
s.sendall(Discover.MAGIC)
|
||||
|
||||
data = b""
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
+28
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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:
|
||||
|
||||
+6
-1
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user