This commit is contained in:
ciomek committed 2026-09-19 20:39:12 +02:00
1 parent b954d9d448
commit b2bd4d25f9
20 files changed
+804 -170

No files matched your search

+3
View File
@@ -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()
+8
View File
@@ -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 {}
+67
View File
@@ -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
+16 -7
View File
@@ -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
+206
View File
@@ -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()]
+4 -4
View File
@@ -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
+51 -133
View File
@@ -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.
+237
View File
@@ -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])
+41 -1
View File
@@ -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}
+44
View File
@@ -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
+4 -4
View File
@@ -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)
+11 -8
View File
@@ -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:
+27
View File
@@ -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])
+4 -4
View File
@@ -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""
+2 -2
View File
@@ -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)
+18
View File
@@ -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
View File
@@ -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
}
}
+27 -6
View File
@@ -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
View File
@@ -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,