piker/piker/service/_actor_runtime.py

349 lines
9.4 KiB
Python
Raw Normal View History

# 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/>.
"""
``tractor`` wrapping + default config to bootstrap the `pikerd`.
"""
from __future__ import annotations
from pprint import pformat
from functools import partial
import os
from typing import (
Optional,
Any,
ClassVar,
)
from contextlib import (
asynccontextmanager as acm,
)
import tractor
import trio
from ._util import (
log, # sub-sys logger
get_console_log,
)
from ._mngr import (
Services,
)
from ._registry import ( # noqa
_tractor_kwargs,
_default_reg_addr,
open_registry,
)
def get_tractor_runtime_kwargs() -> dict[str, Any]:
'''
Deliver ``tractor`` related runtime variables in a `dict`.
'''
return _tractor_kwargs
@acm
async def open_piker_runtime(
name: str,
enable_modules: list[str] = [],
loglevel: Optional[str] = None,
# XXX NOTE XXX: you should pretty much never want debug mode
# for data daemons when running in production.
debug_mode: bool = False,
registry_addr: None | tuple[str, int] = None,
# TODO: once we have `rsyscall` support we will read a config
# and spawn the service tree distributed per that.
start_method: str = 'trio',
tractor_runtime_overrides: dict | None = None,
**tractor_kwargs,
) -> tuple[
tractor.Actor,
tuple[str, int],
]:
'''
Start a piker actor who's runtime will automatically sync with
existing piker actors on the local link based on configuration.
Can be called from a subactor or any program that needs to start
a root actor.
'''
try:
# check for existing runtime
actor = tractor.current_actor().uid
except tractor._exceptions.NoRuntime:
tractor._state._runtime_vars[
'piker_vars'] = tractor_runtime_overrides
registry_addr = registry_addr or _default_reg_addr
async with (
tractor.open_root_actor(
# passed through to ``open_root_actor``
arbiter_addr=registry_addr,
name=name,
loglevel=loglevel,
debug_mode=debug_mode,
start_method=start_method,
# TODO: eventually we should be able to avoid
# having the root have more then permissions to
# spawn other specialized daemons I think?
enable_modules=enable_modules,
**tractor_kwargs,
) as _,
open_registry(registry_addr, ensure_exists=False) as addr,
):
yield (
tractor.current_actor(),
addr,
)
else:
async with open_registry(registry_addr) as addr:
yield (
actor,
addr,
)
_root_dname = 'pikerd'
_root_modules = [
__name__,
'piker.service._daemon',
'piker.clearing._ems',
'piker.clearing._client',
'piker.data._sampling',
]
@acm
async def open_pikerd(
loglevel: str | None = None,
# XXX: you should pretty much never want debug mode
# for data daemons when running in production.
debug_mode: bool = False,
registry_addr: None | tuple[str, int] = None,
# db init flags
tsdb: bool = False,
es: bool = False,
drop_root_perms_for_ahab: bool = True,
**kwargs,
) -> Services:
'''
Start a root piker daemon with an indefinite lifetime.
A root actor nursery is created which can be used to create and keep
alive underling services (see below).
'''
async with (
open_piker_runtime(
name=_root_dname,
# TODO: eventually we should be able to avoid
# having the root have more then permissions to
# spawn other specialized daemons I think?
enable_modules=_root_modules,
loglevel=loglevel,
debug_mode=debug_mode,
registry_addr=registry_addr,
**kwargs,
) as (root_actor, reg_addr),
tractor.open_nursery() as actor_nursery,
trio.open_nursery() as service_nursery,
):
if root_actor.accept_addr != reg_addr:
raise RuntimeError(
f'`pikerd` failed to bind on {reg_addr}!\n'
'Maybe you have another daemon already running?'
)
# assign globally for future daemon/task creation
Services.actor_n = actor_nursery
Services.service_n = service_nursery
Services.debug_mode = debug_mode
if tsdb:
from ._ahab import start_ahab
from .marketstore import start_marketstore
log.info('Spawning `marketstore` supervisor')
ctn_ready, config, (cid, pid) = await service_nursery.start(
partial(
start_ahab,
'marketstored',
start_marketstore,
loglevel=loglevel,
drop_root_perms=drop_root_perms_for_ahab,
)
)
log.info(
f'`marketstored` up!\n'
f'pid: {pid}\n'
f'container id: {cid[:12]}\n'
f'config: {pformat(config)}'
)
if es:
from ._ahab import start_ahab
from .elastic import start_elasticsearch
log.info('Spawning `elasticsearch` supervisor')
ctn_ready, config, (cid, pid) = await service_nursery.start(
partial(
start_ahab,
'elasticsearch',
start_elasticsearch,
loglevel=loglevel,
drop_root_perms=drop_root_perms_for_ahab,
)
)
log.info(
f'`elasticsearch` up!\n'
f'pid: {pid}\n'
f'container id: {cid[:12]}\n'
f'config: {pformat(config)}'
)
try:
yield Services
finally:
# TODO: is this more clever/efficient?
# if 'samplerd' in Services.service_tasks:
# await Services.cancel_service('samplerd')
service_nursery.cancel_scope.cancel()
# TODO: do we even need this?
# @acm
# async def maybe_open_runtime(
# loglevel: Optional[str] = None,
# **kwargs,
# ) -> None:
# '''
# Start the ``tractor`` runtime (a root actor) if none exists.
# '''
# name = kwargs.pop('name')
# if not tractor.current_actor(err_on_no_runtime=False):
# async with open_piker_runtime(
# name,
# loglevel=loglevel,
# **kwargs,
# ) as (_, addr):
# yield addr,
# else:
# async with open_registry() as addr:
# yield addr
@acm
async def maybe_open_pikerd(
loglevel: Optional[str] = None,
registry_addr: None | tuple = None,
tsdb: bool = False,
es: bool = False,
drop_root_perms_for_ahab: bool = True,
**kwargs,
) -> tractor._portal.Portal | ClassVar[Services]:
'''
If no ``pikerd`` daemon-root-actor can be found start it and
yield up (we should probably figure out returning a portal to self
though).
'''
if loglevel:
get_console_log(loglevel)
# subtle, we must have the runtime up here or portal lookup will fail
query_name = kwargs.pop(
'name',
f'piker_query_{os.getpid()}',
)
# TODO: if we need to make the query part faster we could not init
# an actor runtime and instead just hit the socket?
# from tractor._ipc import _connect_chan, Channel
# async with _connect_chan(host, port) as chan:
# async with open_portal(chan) as arb_portal:
# yield arb_portal
async with (
open_piker_runtime(
name=query_name,
registry_addr=registry_addr,
loglevel=loglevel,
**kwargs,
) as _,
tractor.find_actor(
_root_dname,
arbiter_sockaddr=registry_addr,
) as portal
):
# connect to any existing daemon presuming
# its registry socket was selected.
if (
portal is not None
):
yield portal
return
# presume pikerd role since no daemon could be found at
# configured address
async with open_pikerd(
loglevel=loglevel,
registry_addr=registry_addr,
# ahabd (docker super) specific controls
tsdb=tsdb,
es=es,
drop_root_perms_for_ahab=drop_root_perms_for_ahab,
# passthrough to ``tractor`` init
**kwargs,
) as service_manager:
# in the case where we're starting up the
# tractor-piker runtime stack in **this** process
# we return no portal to self.
assert service_manager
yield service_manager