ominfra 0.0.0.dev126__py3-none-any.whl → 0.0.0.dev128__py3-none-any.whl

ominfra/supervisor/supervisor.py

@@ -3,30 +3,40 @@ import signal
 import time
 import typing as ta
 
-from omlish.lite.cached import cached_nullary
 from omlish.lite.check import check_isinstance
 from omlish.lite.check import check_not_none
 from omlish.lite.logs import log
-from omlish.lite.typing import Func
+from omlish.lite.typing import Func1
 
 from .configs import ProcessGroupConfig
 from .context import ServerContextImpl
+from .dispatchers import Dispatchers
 from .events import TICK_EVENTS
 from .events import EventCallbacks
 from .events import SupervisorRunningEvent
 from .events import SupervisorStoppingEvent
 from .groups import ProcessGroup
-from .groups import ProcessGroups
+from .groups import ProcessGroupManager
 from .poller import Poller
-from .signals import SignalReceiver
-from .signals import sig_name
+from .process import PidHistory
+from .setup import SupervisorSetup
 from .states import SupervisorState
-from .types import Dispatcher
+from .types import OutputDispatcher
 from .types import Process
-from .utils import ExitNow
-from .utils import as_string
-from .utils import decode_wait_status
-from .utils import timeslice
+from .utils.os import decode_wait_status
+from .utils.signals import SignalReceiver
+from .utils.signals import sig_name
+
+
+##
+
+
+class ExitNow(Exception):  # noqa
+    pass
+
+
+def timeslice(period: int, when: float) -> int:
+    return int(when - (when % period))
 
 
 ##
@@ -38,7 +48,7 @@ class SignalHandler:
             *,
             context: ServerContextImpl,
             signal_receiver: SignalReceiver,
-            process_groups: ProcessGroups,
+            process_groups: ProcessGroupManager,
     ) -> None:
         super().__init__()
 
@@ -78,8 +88,10 @@ class SignalHandler:
         elif sig == signal.SIGUSR2:
             log.info('received %s indicating log reopen request', sig_name(sig))
 
-            for group in self._process_groups:
-                group.reopen_logs()
+            for p in self._process_groups.all_processes():
+                for d in p.get_dispatchers():
+                    if isinstance(d, OutputDispatcher):
+                        d.reopen_logs()
 
         else:
             log.debug('received %s indicating nothing', sig_name(sig))
@@ -88,7 +100,8 @@ class SignalHandler:
 ##
 
 
-ProcessGroupFactory = ta.NewType('ProcessGroupFactory', Func[ProcessGroup])  # (config: ProcessGroupConfig)
+class ProcessGroupFactory(Func1[ProcessGroupConfig, ProcessGroup]):
+    pass
 
 
 class Supervisor:
@@ -97,10 +110,12 @@ class Supervisor:
             *,
             context: ServerContextImpl,
             poller: Poller,
-            process_groups: ProcessGroups,
+            process_groups: ProcessGroupManager,
             signal_handler: SignalHandler,
             event_callbacks: EventCallbacks,
             process_group_factory: ProcessGroupFactory,
+            pid_history: PidHistory,
+            setup: SupervisorSetup,
     ) -> None:
         super().__init__()
 
@@ -110,6 +125,8 @@ class Supervisor:
         self._signal_handler = signal_handler
         self._event_callbacks = event_callbacks
         self._process_group_factory = process_group_factory
+        self._pid_history = pid_history
+        self._setup = setup
 
         self._ticks: ta.Dict[int, float] = {}
         self._stop_groups: ta.Optional[ta.List[ProcessGroup]] = None  # list used for priority ordered shutdown
@@ -127,31 +144,13 @@ class Supervisor:
 
     #
 
-    class DiffToActive(ta.NamedTuple):
-        added: ta.List[ProcessGroupConfig]
-        changed: ta.List[ProcessGroupConfig]
-        removed: ta.List[ProcessGroupConfig]
-
-    def diff_to_active(self) -> DiffToActive:
-        new = self._context.config.groups or []
-        cur = [group.config for group in self._process_groups]
-
-        curdict = dict(zip([cfg.name for cfg in cur], cur))
-        newdict = dict(zip([cfg.name for cfg in new], new))
-
-        added = [cand for cand in new if cand.name not in curdict]
-        removed = [cand for cand in cur if cand.name not in newdict]
-
-        changed = [cand for cand in new if cand != curdict.get(cand.name, cand)]
-
-        return Supervisor.DiffToActive(added, changed, removed)
-
     def add_process_group(self, config: ProcessGroupConfig) -> bool:
         if self._process_groups.get(config.name) is not None:
             return False
 
         group = check_isinstance(self._process_group_factory(config), ProcessGroup)
-        group.after_setuid()
+        for process in group:
+            process.after_setuid()
 
         self._process_groups.add(group)
 
@@ -165,11 +164,7 @@ class Supervisor:
 
         return True
 
-    def get_process_map(self) -> ta.Dict[int, Dispatcher]:
-        process_map: ta.Dict[int, Dispatcher] = {}
-        for group in self._process_groups:
-            process_map.update(group.get_dispatchers())
-        return process_map
+    #
 
     def shutdown_report(self) -> ta.List[Process]:
         unstopped: ta.List[Process] = []
@@ -181,7 +176,7 @@ class Supervisor:
             # throttle 'waiting for x to die' reports
             now = time.time()
             if now > (self._last_shutdown_report + 3):  # every 3 secs
-                names = [as_string(p.config.name) for p in unstopped]
+                names = [p.config.name for p in unstopped]
                 namestr = ', '.join(names)
                 log.info('waiting for %s to die', namestr)
                 self._last_shutdown_report = now
@@ -192,25 +187,12 @@ class Supervisor:
 
     #
 
-    def main(self) -> None:
-        self.setup()
-        self.run()
-
-    @cached_nullary
-    def setup(self) -> None:
-        if not self._context.first:
-            # prevent crash on libdispatch-based systems, at least for the first request
-            self._context.cleanup_fds()
-
-        self._context.set_uid_or_exit()
-
-        if self._context.first:
-            self._context.set_rlimits_or_exit()
-
-        # this sets the options.logger object delay logger instantiation until after setuid
-        if not self._context.config.nocleanup:
-            # clean up old automatic logs
-            self._context.clear_auto_child_logdir()
+    def main(self, **kwargs: ta.Any) -> None:
+        self._setup.setup()
+        try:
+            self.run(**kwargs)
+        finally:
+            self._setup.cleanup()
 
     def run(
             self,
@@ -228,12 +210,6 @@ class Supervisor:
 
             self._signal_handler.set_signals()
 
-            if not self._context.config.nodaemon and self._context.first:
-                self._context.daemonize()
-
-            # writing pid file needs to come *after* daemonizing or pid will be wrong
-            self._context.write_pidfile()
-
             self._event_callbacks.notify(SupervisorRunningEvent())
 
             while True:
@@ -243,7 +219,7 @@ class Supervisor:
                 self._run_once()
 
         finally:
-            self._context.cleanup()
+            self._poller.close()
 
     #
 
@@ -272,18 +248,24 @@ class Supervisor:
                 # down, so push it back on to the end of the stop group queue
                 self._stop_groups.append(group)
 
+    def get_dispatchers(self) -> Dispatchers:
+        return Dispatchers(
+            d
+            for p in self._process_groups.all_processes()
+            for d in p.get_dispatchers()
+        )
+
     def _poll(self) -> None:
-        combined_map = {}
-        combined_map.update(self.get_process_map())
+        dispatchers = self.get_dispatchers()
 
-        pgroups = list(self._process_groups)
-        pgroups.sort()
+        sorted_groups = list(self._process_groups)
+        sorted_groups.sort()
 
         if self._context.state < SupervisorState.RUNNING:
             if not self._stopping:
                 # first time, set the stopping flag, do a notification and set stop_groups
                 self._stopping = True
-                self._stop_groups = pgroups[:]
+                self._stop_groups = sorted_groups[:]
                 self._event_callbacks.notify(SupervisorStoppingEvent())
 
             self._ordered_stop_groups_phase_1()
@@ -292,7 +274,7 @@ class Supervisor:
                 # if there are no unstopped processes (we're done killing everything), it's OK to shutdown or reload
                 raise ExitNow
 
-        for fd, dispatcher in combined_map.items():
+        for fd, dispatcher in dispatchers.items():
             if dispatcher.readable():
                 self._poller.register_readable(fd)
             if dispatcher.writable():
@@ -302,9 +284,9 @@ class Supervisor:
         r, w = self._poller.poll(timeout)
 
         for fd in r:
-            if fd in combined_map:
+            if fd in dispatchers:
                 try:
-                    dispatcher = combined_map[fd]
+                    dispatcher = dispatchers[fd]
                     log.debug('read event caused by %r', dispatcher)
                     dispatcher.handle_read_event()
                     if not dispatcher.readable():
@@ -312,9 +294,9 @@ class Supervisor:
                 except ExitNow:
                     raise
                 except Exception:  # noqa
-                    combined_map[fd].handle_error()
+                    dispatchers[fd].handle_error()
             else:
-                # if the fd is not in combined_map, we should unregister it. otherwise, it will be polled every
+                # if the fd is not in combined map, we should unregister it. otherwise, it will be polled every
                 # time, which may cause 100% cpu usage
                 log.debug('unexpected read event from fd %r', fd)
                 try:
@@ -323,9 +305,9 @@ class Supervisor:
                     pass
 
         for fd in w:
-            if fd in combined_map:
+            if fd in dispatchers:
                 try:
-                    dispatcher = combined_map[fd]
+                    dispatcher = dispatchers[fd]
                     log.debug('write event caused by %r', dispatcher)
                     dispatcher.handle_write_event()
                     if not dispatcher.writable():
@@ -333,7 +315,7 @@ class Supervisor:
                 except ExitNow:
                     raise
                 except Exception:  # noqa
-                    combined_map[fd].handle_error()
+                    dispatchers[fd].handle_error()
             else:
                 log.debug('unexpected write event from fd %r', fd)
                 try:
@@ -341,8 +323,9 @@ class Supervisor:
                 except Exception:  # noqa
                     pass
 
-        for group in pgroups:
-            group.transition()
+        for group in sorted_groups:
+            for process in group:
+                process.transition()
 
     def _reap(self, *, once: bool = False, depth: int = 0) -> None:
         if depth >= 100:
@@ -352,13 +335,13 @@ class Supervisor:
         if not pid:
             return
 
-        process = self._context.pid_history.get(pid, None)
+        process = self._pid_history.get(pid, None)
         if process is None:
             _, msg = decode_wait_status(check_not_none(sts))
             log.info('reaped unknown pid %s (%s)', pid, msg)
         else:
             process.finish(check_not_none(sts))
-            del self._context.pid_history[pid]
+            del self._pid_history[pid]
 
         if not once:
             # keep reaping until no more kids to reap, but don't recurse infinitely

ominfra/supervisor/types.py

@@ -8,6 +8,40 @@ from .configs import ProcessGroupConfig
 from .configs import ServerConfig
 from .states import ProcessState
 from .states import SupervisorState
+from .utils.collections import KeyedCollectionAccessors
+from .utils.ostypes import Fd
+from .utils.ostypes import Pid
+from .utils.ostypes import Rc
+
+
+if ta.TYPE_CHECKING:
+    from .dispatchers import Dispatchers
+
+
+##
+
+
+ServerEpoch = ta.NewType('ServerEpoch', int)
+
+
+##
+
+
+@functools.total_ordering
+class ConfigPriorityOrdered(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def config(self) -> ta.Any:
+        raise NotImplementedError
+
+    def __lt__(self, other):
+        return self.config.priority < other.config.priority
+
+    def __eq__(self, other):
+        return self.config.priority == other.config.priority
+
+
+##
 
 
 class ServerContext(abc.ABC):
@@ -25,13 +59,43 @@ class ServerContext(abc.ABC):
     def set_state(self, state: SupervisorState) -> None:
         raise NotImplementedError
 
+
+##
+
+
+class Dispatcher(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def process(self) -> 'Process':
+        raise NotImplementedError
+
     @property
     @abc.abstractmethod
-    def pid_history(self) -> ta.Dict[int, 'Process']:
+    def channel(self) -> str:
         raise NotImplementedError
 
+    @property
+    @abc.abstractmethod
+    def fd(self) -> Fd:
+        raise NotImplementedError
+
+    @property
+    @abc.abstractmethod
+    def closed(self) -> bool:
+        raise NotImplementedError
+
+    #
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def handle_error(self) -> None:
+        raise NotImplementedError
+
+    #
 
-class Dispatcher(abc.ABC):
     @abc.abstractmethod
     def readable(self) -> bool:
         raise NotImplementedError
@@ -40,26 +104,25 @@ class Dispatcher(abc.ABC):
     def writable(self) -> bool:
         raise NotImplementedError
 
+    #
+
     def handle_read_event(self) -> None:
         raise TypeError
 
     def handle_write_event(self) -> None:
         raise TypeError
 
+
+class OutputDispatcher(Dispatcher, abc.ABC):
     @abc.abstractmethod
-    def handle_error(self) -> None:
+    def remove_logs(self) -> None:
         raise NotImplementedError
 
-    @property
     @abc.abstractmethod
-    def closed(self) -> bool:
+    def reopen_logs(self) -> None:
         raise NotImplementedError
 
 
-class OutputDispatcher(Dispatcher, abc.ABC):
-    pass
-
-
 class InputDispatcher(Dispatcher, abc.ABC):
     @abc.abstractmethod
     def write(self, chars: ta.Union[bytes, str]) -> None:
@@ -70,11 +133,13 @@ class InputDispatcher(Dispatcher, abc.ABC):
         raise NotImplementedError
 
 
-@functools.total_ordering
-class Process(abc.ABC):
+##
+
+
+class Process(ConfigPriorityOrdered, abc.ABC):
     @property
     @abc.abstractmethod
-    def pid(self) -> int:
+    def name(self) -> str:
         raise NotImplementedError
 
     @property
@@ -82,27 +147,25 @@ class Process(abc.ABC):
     def config(self) -> ProcessConfig:
         raise NotImplementedError
 
-    def __lt__(self, other):
-        return self.config.priority < other.config.priority
-
-    def __eq__(self, other):
-        return self.config.priority == other.config.priority
-
     @property
     @abc.abstractmethod
-    def context(self) -> ServerContext:
+    def group(self) -> 'ProcessGroup':
         raise NotImplementedError
 
+    @property
     @abc.abstractmethod
-    def finish(self, sts: int) -> None:
+    def pid(self) -> Pid:
         raise NotImplementedError
 
+    #
+
+    @property
     @abc.abstractmethod
-    def remove_logs(self) -> None:
+    def context(self) -> ServerContext:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def reopen_logs(self) -> None:
+    def finish(self, sts: Rc) -> None:
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -122,50 +185,41 @@ class Process(abc.ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def create_auto_child_logs(self) -> None:
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def get_dispatchers(self) -> ta.Mapping[int, Dispatcher]:
+    def after_setuid(self) -> None:
         raise NotImplementedError
 
-
-@functools.total_ordering
-class ProcessGroup(abc.ABC):
-    @property
     @abc.abstractmethod
-    def config(self) -> ProcessGroupConfig:
+    def get_dispatchers(self) -> 'Dispatchers':
         raise NotImplementedError
 
-    def __lt__(self, other):
-        return self.config.priority < other.config.priority
-
-    def __eq__(self, other):
-        return self.config.priority == other.config.priority
 
-    @abc.abstractmethod
-    def transition(self) -> None:
-        raise NotImplementedError
+##
 
-    @abc.abstractmethod
-    def stop_all(self) -> None:
-        raise NotImplementedError
 
+class ProcessGroup(
+    ConfigPriorityOrdered,
+    KeyedCollectionAccessors[str, Process],
+    abc.ABC,
+):
     @property
     @abc.abstractmethod
     def name(self) -> str:
         raise NotImplementedError
 
+    @property
     @abc.abstractmethod
-    def before_remove(self) -> None:
+    def config(self) -> ProcessGroupConfig:
         raise NotImplementedError
 
+    @property
     @abc.abstractmethod
-    def get_dispatchers(self) -> ta.Mapping[int, Dispatcher]:
+    def by_name(self) -> ta.Mapping[str, Process]:
         raise NotImplementedError
 
+    #
+
     @abc.abstractmethod
-    def reopen_logs(self) -> None:
+    def stop_all(self) -> None:
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -173,5 +227,5 @@ class ProcessGroup(abc.ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def after_setuid(self) -> None:
+    def before_remove(self) -> None:
         raise NotImplementedError

ominfra/supervisor/utils/collections.py

@@ -0,0 +1,52 @@
+# ruff: noqa: UP006 UP007
+import abc
+import typing as ta
+
+
+K = ta.TypeVar('K')
+V = ta.TypeVar('V')
+
+
+class KeyedCollectionAccessors(abc.ABC, ta.Generic[K, V]):
+    @property
+    @abc.abstractmethod
+    def _by_key(self) -> ta.Mapping[K, V]:
+        raise NotImplementedError
+
+    def __iter__(self) -> ta.Iterator[V]:
+        return iter(self._by_key.values())
+
+    def __len__(self) -> int:
+        return len(self._by_key)
+
+    def __contains__(self, key: K) -> bool:
+        return key in self._by_key
+
+    def __getitem__(self, key: K) -> V:
+        return self._by_key[key]
+
+    def get(self, key: K, default: ta.Optional[V] = None) -> ta.Optional[V]:
+        return self._by_key.get(key, default)
+
+    def items(self) -> ta.Iterator[ta.Tuple[K, V]]:
+        return iter(self._by_key.items())
+
+
+class KeyedCollection(KeyedCollectionAccessors[K, V]):
+    def __init__(self, items: ta.Iterable[V]) -> None:
+        super().__init__()
+
+        by_key: ta.Dict[K, V] = {}
+        for v in items:
+            if (k := self._key(v)) in by_key:
+                raise KeyError(f'key {k} of {v} already registered by {by_key[k]}')
+            by_key[k] = v
+        self.__by_key = by_key
+
+    @property
+    def _by_key(self) -> ta.Mapping[K, V]:
+        return self.__by_key
+
+    @abc.abstractmethod
+    def _key(self, v: V) -> K:
+        raise NotImplementedError

ominfra/supervisor/utils/diag.py

@@ -0,0 +1,31 @@
+# ruff: noqa: UP006 UP007
+import sys
+import types
+import typing as ta
+
+
+def compact_traceback() -> ta.Tuple[
+    ta.Tuple[str, str, int],
+    ta.Type[BaseException],
+    BaseException,
+    types.TracebackType,
+]:
+    t, v, tb = sys.exc_info()
+    if not tb:
+        raise RuntimeError('No traceback')
+
+    tbinfo = []
+    while tb:
+        tbinfo.append((
+            tb.tb_frame.f_code.co_filename,
+            tb.tb_frame.f_code.co_name,
+            str(tb.tb_lineno),
+        ))
+        tb = tb.tb_next
+
+    # just to be safe
+    del tb
+
+    file, function, line = tbinfo[-1]
+    info = ' '.join(['[%s|%s|%s]' % x for x in tbinfo])  # noqa
+    return (file, function, line), t, v, info  # type: ignore

ominfra/supervisor/utils/fds.py

@@ -0,0 +1,46 @@
+# ruff: noqa: UP006 UP007
+import errno
+import os
+import typing as ta
+
+from .ostypes import Fd
+
+
+class PipeFds(ta.NamedTuple):
+    r: Fd
+    w: Fd
+
+
+def make_pipe() -> PipeFds:
+    return PipeFds(*os.pipe())  # type: ignore
+
+
+def read_fd(fd: Fd) -> bytes:
+    try:
+        data = os.read(fd, 2 << 16)  # 128K
+    except OSError as why:
+        if why.args[0] not in (errno.EWOULDBLOCK, errno.EBADF, errno.EINTR):
+            raise
+        data = b''
+    return data
+
+
+def close_fd(fd: Fd) -> bool:
+    try:
+        os.close(fd)
+    except OSError:
+        return False
+    return True
+
+
+def is_fd_open(fd: Fd) -> bool:
+    try:
+        n = os.dup(fd)
+    except OSError:
+        return False
+    os.close(n)
+    return True
+
+
+def get_open_fds(limit: int) -> ta.FrozenSet[Fd]:
+    return frozenset(fd for i in range(limit) if is_fd_open(fd := Fd(i)))