colony/colony/core/network/membership.py

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()]