piker/piker/service/_registry.py

# piker: trading gear for hackers
# Copyright (C) Tyler Goodlet (in stewardship for pikers)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""
Inter-actor "discovery" (protocol) layer.

"""
from __future__ import annotations
from contextlib import (
    asynccontextmanager as acm,
)
from typing import (
    Any,
    TypeAlias,
)

import tractor
from tractor import (
    msg,
    Actor,
    Portal,
)
from tractor.discovery._addr import (
    LegacyUnwrappedAddress,
    UnwrappedAddress,
    wrap_address,
)

from piker.log import get_logger

log = get_logger(name=__name__)

# Piker config still accepts legacy socket pairs while tractor now
# emits tagged `UnwrappedAddress` values.
RegistryAddress: TypeAlias = (
    LegacyUnwrappedAddress
    |UnwrappedAddress
)

# TODO? default path-space for UDS registry?
# [ ] needs to be Xplatform tho!
# _default_registry_path: Path = (
#     Path(os.environ['XDG_RUNTIME_DIR'])
#     /'piker'
# )

_default_registry_host: str = '127.0.0.1'
_default_registry_port: int = 6116
_default_reg_addr: RegistryAddress = (
    _default_registry_host,
    _default_registry_port,
)


# NOTE: this value is set as an actor-global once the first endpoint
# who is capable, spawns a `pikerd` service tree.
_registry: Registry | None = None


class Registry:
    # TODO: should this be a set or should we complain
    # on duplicates?
    addrs: list[RegistryAddress] = []

    # TODO: table of uids to sockaddrs
    peers: dict[
        tuple[str, str],
        RegistryAddress,
    ] = {}


_tractor_kwargs: dict[str, Any] = {}


def _normalize_addresses(
    addrs: list[RegistryAddress],
) -> list[RegistryAddress]:
    '''
    Normalize legacy and tagged addresses through tractor's codec.

    '''
    return [
        wrap_address(addr).unwrap()
        for addr in addrs
    ]


@acm
async def open_registry(
    addrs: list[RegistryAddress],
    ensure_exists: bool = True,

) -> list[RegistryAddress]:
    '''
    Open the service-actor-discovery registry by returning a set of
    tranport socket-addrs to registrar actors which may be
    contacted and queried for similar addresses for other
    non-registrar actors.

    '''
    global _tractor_kwargs
    actor: Actor = tractor.current_actor()
    aid: msg.Aid = actor.aid
    uid: tuple[str, str] = aid.uid
    addrs = _normalize_addresses(addrs)
    preset_reg_addrs: list[RegistryAddress] = (
        _normalize_addresses(Registry.addrs)
    )
    if preset_reg_addrs:
        Registry.addrs = preset_reg_addrs
    if (
        preset_reg_addrs
        and
        addrs
    ):
        if preset_reg_addrs != addrs:
            # if any(addr in preset_reg_addrs for addr in addrs):
            diff: set[RegistryAddress] = (
                set(preset_reg_addrs)
                -set(addrs)
            )
            if diff:
                log.warning(
                    f'`{uid}` requested only subset of registrars: {addrs}\n'
                    f'However there are more @{diff}'
                )
            else:
                raise RuntimeError(
                    f'`{uid}` has non-matching registrar addresses?\n'
                    f'request: {addrs}\n'
                    f'already set: {preset_reg_addrs}'
                )

    was_set: bool = False
    if (
        not tractor.is_root_process()
        and
        not Registry.addrs
    ):
        Registry.addrs.extend(
            _normalize_addresses(actor.reg_addrs)
        )

    if (
        ensure_exists
        and
        not Registry.addrs
    ):
        raise RuntimeError(
            f"`{uid}` registry should already exist but doesn't?"
        )

    if not Registry.addrs:
        was_set = True
        Registry.addrs = (
            addrs
            or
            _normalize_addresses([_default_reg_addr])
        )

    # NOTE: only spot this seems currently used is inside
    # `.ui._exec` which is the (eventual qtloops) bootstrapping
    # with guest mode.
    reg_addrs: list[RegistryAddress] = Registry.addrs
    # !TODO, a struct-API to stringently allow this only in special
    # cases?
    # -> better would be to have some way to (atomically) rewrite
    #    and entire `RuntimeVars`?? ideas welcome obvi..
    _tractor_kwargs['registry_addrs'] = reg_addrs

    try:
        yield Registry.addrs
    finally:
        # XXX: always clear the global addr if we set it so that the
        # next (set of) calls will apply whatever new one is passed
        # in.
        if was_set:
            Registry.addrs = []


@acm
async def find_service(
    service_name: str,
    registry_addrs: list[RegistryAddress] | None = None,

    first_only: bool = True,

) -> (
    Portal
    | list[Portal]
    | None
):
    # try:
    reg_addrs: list[RegistryAddress]
    async with open_registry(
        addrs=(
            registry_addrs
            # NOTE: if no addr set is passed assume the registry has
            # already been opened and use the previously applied
            # startup set.
            or Registry.addrs
        ),
    ) as reg_addrs:

        log.info(
            f'Scanning for service {service_name!r}'
        )

        # attach to existing daemon by name if possible
        maybe_portals: list[Portal]|Portal|None
        async with tractor.find_actor(
            service_name,
            registry_addrs=reg_addrs,
            only_first=first_only,  # if set only returns single ref
        ) as maybe_portals:
            if not maybe_portals:
                log.info(
                    f'Could NOT find service {service_name!r} -> {maybe_portals!r}'
                )
                yield None
                return

            log.info(
                f'Found service {service_name!r} -> {maybe_portals}'
            )
            yield maybe_portals

    # except BaseException as _berr:
    #     berr = _berr
    #     log.exception(
    #         'tractor.find_actor() failed with,\n'
    #     )
    #     raise berr


async def check_for_service(
    service_name: str,
) -> None|RegistryAddress:
    '''
    Service daemon "liveness" predicate.

    '''
    async with (
        open_registry(
            addrs=Registry.addrs,
            ensure_exists=False,
        ) as reg_addrs,
        tractor.query_actor(
            service_name,
            regaddr=reg_addrs[0],
        ) as (sockaddr, _),
    ):
        return sockaddr