basis/components/models/cluster.py

import asyncio
import socket

from components.cluster.ssl import get_ssl_context
from components.models import *
from components.utils.datetimes import ntime_utc_now
from config import defaults
from contextlib import closing


class Role(Enum):
    LEADER = 1
    FOLLOWER = 2


class ConnectionStatus(Enum):
    CONNECTED = 0
    REFUSED = 1
    SOCKET_REFUSED = 2
    ALL_AVAILABLE_FAILED = 3
    OK = 4
    OK_WITH_PREVIOUS_ERRORS = 5


class CritErrors(Enum):
    NOT_READY = "CRIT:NOT_READY"
    TABLE_HASH_MISMATCH = "CRIT:TABLE_HASH_MISMATCH"
    CANNOT_APPLY = "CRIT:CANNOT_APPLY"
    NOTHING_TO_COMMIT = "CRIT:NOTHING_TO_COMMIT"
    INVALID_FILE_PATH = "CRIT:INVALID_FILE_PATH"
    START_BEHIND_FILE_END = "CRIT:START_BEHIND_FILE_END"
    PEERS_MISMATCH = "CRIT:PEERS_MISMATCH"
    DOC_MISMATCH = "CRIT:DOC_MISMATCH"
    ZOMBIE = "CRIT:ZOMBIE"

    @property
    def response(self):
        return f"ACK {self.value}"


class LocalPeer(BaseModel):
    @model_validator(mode="before")
    @classmethod
    def pre_init(cls, data: Any) -> Any:
        if not data["ip4"] and not data["ip6"]:
            raise ValueError("Neither a IPv4 nor a IPv6 address was provided")
        return data

    @field_validator("is_self")
    def local_self_validator(cls, v):
        if v != True:
            raise ValueError("LocalPeer does not have is_self flag")
        return v

    is_self: bool
    name: constr(pattern=r"^[a-zA-Z0-9\-_\.]+$", min_length=3)
    ip4: IPvAnyAddress | None = None
    ip6: IPvAnyAddress | None = None
    cli_bindings: list[IPvAnyAddress] = defaults.CLUSTER_CLI_BINDINGS
    leader: str | None = None
    role: Role = Role.FOLLOWER
    swarm: str = ""
    started: float = ntime_utc_now()
    swarm_complete: bool = False

    @computed_field
    @property
    def _bindings_as_str(self) -> str:
        return [str(ip) for key in ("ip4", "ip6") if (ip := getattr(self, key))]

    @computed_field
    @property
    def _all_bindings_as_str(self) -> str:
        return [
            str(ip) for key in ("ip4", "ip6") if (ip := getattr(self, key))
        ] + self.cli_bindings

    @model_validator(mode="after")
    def cli_bindings_validator(self):
        for ip in self.cli_bindings:
            if ip == self.ip4 or ip == self.ip6:
                raise ValueError("CLI bindings overlap local bindings")
        return self


class Streams(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    out: tuple[asyncio.StreamReader, asyncio.StreamWriter] | None = None
    _in: tuple[asyncio.StreamReader, asyncio.StreamWriter] | None = None


class RemotePeer(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    @model_validator(mode="before")
    @classmethod
    def pre_init(cls, data: Any) -> Any:
        if not data["ip4"] and not data["ip6"]:
            raise ValueError("Neither a IPv4 nor a IPv6 address was provided")
        return data

    @field_validator("is_self")
    def local_self_validator(cls, v):
        if v:
            raise ValueError("RemotePeer has is_self flag")
        return v

    is_self: bool = False
    sending_lock: asyncio.Lock = asyncio.Lock()
    swarm: str = ""
    leader: str | None = None
    started: float | None = None
    name: constr(pattern=r"^[a-zA-Z0-9\-_\.]+$", min_length=3)
    ip4: IPvAnyAddress | None = None
    ip6: IPvAnyAddress | None = None
    nat_ip4: IPvAnyAddress | None = None
    streams: Streams = Streams()
    port: int = 2102

    def reset(self):
        self.streams._in = None
        self.streams.out = None
        self.leader = None
        self.started = None
        self.swarm = ""
        return self

    def _eval_ip(
        self,
    ) -> tuple[IPvAnyAddress | None, tuple[ConnectionStatus, dict]]:
        errors = dict()
        peer_ips = [ip for ip in [self.ip4, self.ip6] if ip is not None]
        for ip in peer_ips:
            with closing(
                socket.socket(
                    socket.AF_INET if ip.version == 4 else socket.AF_INET6,
                    socket.SOCK_STREAM,
                )
            ) as sock:
                sock.settimeout(defaults.CLUSTER_PEERS_TIMEOUT)
                connection_return = sock.connect_ex((str(ip), self.port))
                if connection_return != 0:
                    errors[ip] = (
                        ConnectionStatus.SOCKET_REFUSED,
                        socket.errno.errorcode.get(connection_return),
                    )
                else:
                    if errors:
                        return ip, (ConnectionStatus.OK_WITH_PREVIOUS_ERRORS, errors)
                    return ip, (ConnectionStatus.OK, errors)
        else:
            return None, (ConnectionStatus.ALL_AVAILABLE_FAILED, errors)

    async def connect(
        self,
        ip: IPvAnyAddress | None = None,
    ) -> tuple[
        tuple[asyncio.StreamReader, asyncio.StreamWriter] | None,
        tuple[ConnectionStatus, Any],
    ]:
        if not self.streams.out:
            if not ip:
                ip, status = self._eval_ip()
                if not ip:
                    return None, status
            try:
                self.streams.out = await asyncio.open_connection(
                    str(ip), self.port, ssl=get_ssl_context("client")
                )
            except ConnectionRefusedError as e:
                return None, (ConnectionStatus.REFUSED, e)

        return self.streams.out, (ConnectionStatus.CONNECTED, None)

    @computed_field
    @property
    def _all_ips_as_str(self) -> str:
        return [str(ip) for key in ("ip4", "ip6") if (ip := getattr(self, key))]

    @computed_field
    @property
    def _fully_established(self) -> str:
        return (
            True
            if self.streams.out
            and self.streams._in
            and self.swarm
            and self.started
            and self.leader
            else False
        )