omlish 0.0.0.dev225__py3-none-any.whl → 0.0.0.dev227__py3-none-any.whl

omlish/lite/timeouts.py

@@ -0,0 +1,202 @@
+# ruff: noqa: UP006 UP007
+"""
+TODO:
+ - Event (/ Predicate)
+"""
+import abc
+import time
+import typing as ta
+
+
+TimeoutLike = ta.Union['Timeout', 'Timeout.Default', ta.Iterable['TimeoutLike'], float]  # ta.TypeAlias
+
+
+##
+
+
+class Timeout(abc.ABC):
+    @property
+    @abc.abstractmethod
+    def can_expire(self) -> bool:
+        """Indicates whether or not this timeout will ever expire."""
+
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def expired(self) -> bool:
+        """Return whether or not this timeout has expired."""
+
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def remaining(self) -> float:
+        """Returns the time (in seconds) remaining until the timeout expires. May be negative and/or infinite."""
+
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def __call__(self) -> float:
+        """Returns the time (in seconds) remaining until the timeout expires, or raises if the timeout has expired."""
+
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def or_(self, o: ta.Any) -> ta.Any:
+        """Evaluates time remaining via remaining() if this timeout can expire, otherwise returns `o`."""
+
+        raise NotImplementedError
+
+    #
+
+    @classmethod
+    def _now(cls) -> float:
+        return time.time()
+
+    #
+
+    class Default:
+        def __new__(cls, *args, **kwargs):  # noqa
+            raise TypeError
+
+    class _NOT_SPECIFIED:  # noqa
+        def __new__(cls, *args, **kwargs):  # noqa
+            raise TypeError
+
+    @classmethod
+    def of(
+            cls,
+            obj: ta.Optional[TimeoutLike],
+            default: ta.Union[TimeoutLike, ta.Type[_NOT_SPECIFIED]] = _NOT_SPECIFIED,
+    ) -> 'Timeout':
+        if obj is None:
+            return InfiniteTimeout()
+
+        elif isinstance(obj, Timeout):
+            return obj
+
+        elif isinstance(obj, (float, int)):
+            return DeadlineTimeout(cls._now() + obj)
+
+        elif isinstance(obj, ta.Iterable):
+            return CompositeTimeout(*[Timeout.of(c) for c in obj])
+
+        elif obj is Timeout.Default:
+            if default is Timeout._NOT_SPECIFIED or default is Timeout.Default:
+                raise RuntimeError('Must specify a default timeout')
+
+            else:
+                return Timeout.of(default)  # type: ignore[arg-type]
+
+        else:
+            raise TypeError(obj)
+
+    @classmethod
+    def of_deadline(cls, deadline: float) -> 'DeadlineTimeout':
+        return DeadlineTimeout(deadline)
+
+    @classmethod
+    def of_predicate(cls, expired_fn: ta.Callable[[], bool]) -> 'PredicateTimeout':
+        return PredicateTimeout(expired_fn)
+
+
+class DeadlineTimeout(Timeout):
+    def __init__(
+            self,
+            deadline: float,
+            exc: ta.Union[ta.Type[BaseException], BaseException] = TimeoutError,
+    ) -> None:
+        super().__init__()
+
+        self.deadline = deadline
+        self.exc = exc
+
+    @property
+    def can_expire(self) -> bool:
+        return True
+
+    def expired(self) -> bool:
+        return not (self.remaining() > 0)
+
+    def remaining(self) -> float:
+        return self.deadline - self._now()
+
+    def __call__(self) -> float:
+        if (rem := self.remaining()) > 0:
+            return rem
+        raise self.exc
+
+    def or_(self, o: ta.Any) -> ta.Any:
+        return self()
+
+
+class InfiniteTimeout(Timeout):
+    @property
+    def can_expire(self) -> bool:
+        return False
+
+    def expired(self) -> bool:
+        return False
+
+    def remaining(self) -> float:
+        return float('inf')
+
+    def __call__(self) -> float:
+        return float('inf')
+
+    def or_(self, o: ta.Any) -> ta.Any:
+        return o
+
+
+class CompositeTimeout(Timeout):
+    def __init__(self, *children: Timeout) -> None:
+        super().__init__()
+
+        self.children = children
+
+    @property
+    def can_expire(self) -> bool:
+        return any(c.can_expire for c in self.children)
+
+    def expired(self) -> bool:
+        return any(c.expired() for c in self.children)
+
+    def remaining(self) -> float:
+        return min(c.remaining() for c in self.children)
+
+    def __call__(self) -> float:
+        return min(c() for c in self.children)
+
+    def or_(self, o: ta.Any) -> ta.Any:
+        if self.can_expire:
+            return self()
+        return o
+
+
+class PredicateTimeout(Timeout):
+    def __init__(
+            self,
+            expired_fn: ta.Callable[[], bool],
+            exc: ta.Union[ta.Type[BaseException], BaseException] = TimeoutError,
+    ) -> None:
+        super().__init__()
+
+        self.expired_fn = expired_fn
+        self.exc = exc
+
+    @property
+    def can_expire(self) -> bool:
+        return True
+
+    def expired(self) -> bool:
+        return self.expired_fn()
+
+    def remaining(self) -> float:
+        return float('inf')
+
+    def __call__(self) -> float:
+        if not self.expired_fn():
+            return float('inf')
+        raise self.exc
+
+    def or_(self, o: ta.Any) -> ta.Any:
+        return self()

omlish/os/fcntl.py

@@ -0,0 +1,59 @@
+# ruff: noqa: UP006 UP007
+# @omlish-lite
+import dataclasses as dc
+import os
+import struct
+import sys
+import typing as ta
+
+
+@dc.dataclass(frozen=True)
+class FcntlLockData:
+    # cmd = {F_SETLK, F_SETLKW, F_GETLK}
+
+    type: int  # {F_RDLCK, F_WRLCK, F_UNLCK}
+    whence: int = os.SEEK_SET
+    start: int = 0
+    len: int = 0
+    pid: int = 0
+
+    #
+
+    _STRUCT_PACK_BY_PLATFORM: ta.ClassVar[ta.Mapping[str, ta.Sequence[ta.Tuple[str, str]]]] = {
+        'linux': [
+            ('type', 'h'),
+            ('whence', 'h'),
+            ('start', 'q'),
+            ('len', 'q'),
+            ('pid', 'i'),
+        ],
+        'darwin': [
+            ('start', 'q'),
+            ('len', 'q'),
+            ('pid', 'i'),
+            ('type', 'h'),
+            ('whence', 'h'),
+        ],
+    }
+
+    def pack(self) -> bytes:
+        try:
+            pack = self._STRUCT_PACK_BY_PLATFORM[sys.platform]
+        except KeyError:
+            raise OSError from None
+
+        fmt = ''.join(f for _, f in pack)
+        tup = [getattr(self, a) for a, _ in pack]
+        return struct.pack(fmt, *tup)
+
+    @classmethod
+    def unpack(cls, data: bytes) -> 'FcntlLockData':
+        try:
+            pack = cls._STRUCT_PACK_BY_PLATFORM[sys.platform]
+        except KeyError:
+            raise OSError from None
+
+        fmt = ''.join(f for _, f in pack)
+        tup = struct.unpack(fmt, data)
+        kw = {a: v for (a, _), v in zip(pack, tup)}
+        return FcntlLockData(**kw)

omlish/os/pidfiles/manager.py

@@ -14,6 +14,17 @@ from .pidfile import Pidfile
 
 
 class _PidfileManager:
+    """
+    Manager for controlled inheritance of Pidfiles across forks in the presence of multiple threads. There is of course
+    no safe or correct way to mix the use of fork and multiple active threads, and one should never write code which
+    does so, but in the Real World one may still find oneself in such a situation outside of their control (such as when
+    running under Pycharm's debugger which forces the use of forked multiprocessing).
+
+    Not implemented as an instantiated class as there is no way to unregister at_fork listeners, and because Pidfiles
+    may be pickled and there must be no possibility of accidentally unpickling and instantiating a new instance of the
+    manager.
+    """
+
     def __new__(cls, *args, **kwargs):  # noqa
         raise TypeError
 
@@ -71,6 +82,11 @@ class _PidfileManager:
             inheritable: bool = True,
             **kwargs: ta.Any,
     ) -> ta.Iterator[Pidfile]:
+        """
+        A contextmanager for creating and managing a Pidfile which will only be inherited by forks of the calling /
+        creating thread.
+        """
+
         check.arg(inheritable)
 
         cls.install()
@@ -84,6 +100,7 @@ class _PidfileManager:
         with cls._lock:
             cls._pidfile_threads[pf] = threading.current_thread()
         try:
+
             with pf:
                 os.set_inheritable(check.not_none(pf.fileno()), True)
                 yield pf

omlish/os/pidfiles/pidfile.py

@@ -10,6 +10,7 @@ TODO:
    - 3) recheck current pid of flock holder == that pid
   - racy as to if it's a different actual process as initial check, just with same pid, but due to 'identity' / semantic
     meaning of the named pidfile the processes are considered equivalent
+  - read_checked(), contextmanager
 """
 import fcntl
 import os
@@ -32,6 +33,14 @@ class Pidfile:
         self._path = path
         self._inheritable = inheritable
 
+    @property
+    def path(self) -> str:
+        return self._path
+
+    @property
+    def inheritable(self) -> bool:
+        return self._inheritable
+
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}({self._path!r})'
 
@@ -75,6 +84,7 @@ class Pidfile:
         state = self.__dict__.copy()
 
         if '_f' in state:
+            # self._inheritable may be decoupled from actual file inheritability - for example when using the manager.
             if os.get_inheritable(fd := state.pop('_f').fileno()):
                 state['__fd'] = fd
 
@@ -99,7 +109,7 @@ class Pidfile:
         del self._f
         return True
 
-    def try_lock(self) -> bool:
+    def try_acquire_lock(self) -> bool:
         try:
             fcntl.flock(self._f, fcntl.LOCK_EX | fcntl.LOCK_NB)
             return True
@@ -107,14 +117,29 @@ class Pidfile:
         except OSError:
             return False
 
-    def ensure_locked(self) -> None:
-        if not self.try_lock():
-            raise RuntimeError('Could not get lock')
+    #
+
+    class Error(Exception):
+        pass
+
+    class LockedError(Error):
+        pass
+
+    def acquire_lock(self) -> None:
+        if not self.try_acquire_lock():
+            raise self.LockedError
+
+    class NotLockedError(Error):
+        pass
+
+    def ensure_cannot_lock(self) -> None:
+        if self.try_acquire_lock():
+            raise self.NotLockedError
 
     #
 
     def write(self, pid: ta.Optional[int] = None) -> None:
-        self.ensure_locked()
+        self.acquire_lock()
 
         if pid is None:
             pid = os.getpid()
@@ -125,18 +150,23 @@ class Pidfile:
         self._f.flush()
 
     def clear(self) -> None:
-        self.ensure_locked()
+        self.acquire_lock()
 
         self._f.seek(0)
         self._f.truncate()
 
-    def read(self) -> int:
-        if self.try_lock():
-            raise RuntimeError('Got lock')
+    #
+
+    def read(self) -> ta.Optional[int]:
+        self.ensure_cannot_lock()
 
         self._f.seek(0)
-        return int(self._f.read())  # FIXME: could be empty or hold old value, race w proc start
+        buf = self._f.read()
+        if not buf:
+            return None
+        return int(buf)
 
     def kill(self, sig: int = signal.SIGTERM) -> None:
-        pid = self.read()
-        os.kill(pid, sig)  # FIXME: Still racy - pidfd_send_signal?
+        if (pid := self.read()) is None:
+            raise self.Error(f'Pidfile locked but empty')
+        os.kill(pid, sig)

omlish/os/pidfiles/pinning.py

@@ -0,0 +1,246 @@
+# ruff: noqa: UP006 UP007
+# @omlish-lite
+"""
+Strategies:
+ - linux
+  - get pid of owner (lslocks or F_GETLK)
+  - open pidfd to owner pid
+  - re-check pid of owner
+ - darwin
+  - get pids of referrers (lsof)
+  - read pid from file
+  - ensure pid is in referrers
+  - optionally loop
+  - ? setup pid death watcher? still a race
+"""
+import abc
+import contextlib
+import os.path
+import shutil
+import sys
+import time
+import typing as ta
+
+from ...diag.lslocks import LslocksCommand
+from ...diag.lsof import LsofCommand
+from ...lite.timeouts import Timeout
+from ...lite.timeouts import TimeoutLike
+from ...subprocesses.sync import subprocesses  # noqa
+from .pidfile import Pidfile
+
+
+##
+
+
+class PidfilePinner(abc.ABC):
+    def __init__(
+            self,
+            *,
+            sleep_s: float = .1,
+    ) -> None:
+        super().__init__()
+
+        self._sleep_s = sleep_s
+
+    @classmethod
+    @abc.abstractmethod
+    def is_available(cls) -> bool:
+        raise NotImplementedError
+
+    class NoOwnerError(Exception):
+        pass
+
+    @abc.abstractmethod
+    def _pin_pidfile_owner(self, pidfile: Pidfile, timeout: Timeout) -> ta.ContextManager[int]:
+        raise NotImplementedError
+
+    @contextlib.contextmanager
+    def pin_pidfile_owner(
+            self,
+            path: str,
+            *,
+            timeout: ta.Optional[TimeoutLike] = None,
+    ) -> ta.Iterator[int]:
+        timeout = Timeout.of(timeout)
+
+        if not os.path.isfile(path):
+            raise self.NoOwnerError
+
+        with Pidfile(path, inheritable=False) as pf:
+            try:
+                with self._pin_pidfile_owner(pf, timeout) as pid:
+                    yield pid
+
+            except Pidfile.NotLockedError:
+                raise self.NoOwnerError from None
+
+    @classmethod
+    def default_impl(cls) -> ta.Type['PidfilePinner']:
+        for impl in [
+            LslocksPidfdPidfilePinner,
+            LsofPidfilePinner,
+        ]:
+            if impl.is_available():
+                return impl
+        return UnverifiedPidfilePinner
+
+
+##
+
+
+class UnverifiedPidfilePinner(PidfilePinner):
+    @classmethod
+    def is_available(cls) -> bool:
+        return True
+
+    @contextlib.contextmanager
+    def _pin_pidfile_owner(self, pidfile: Pidfile, timeout: Timeout) -> ta.Iterator[int]:
+        while (pid := pidfile.read()) is None:
+            time.sleep(self._sleep_s)
+            timeout()
+
+        yield pid
+
+
+##
+
+
+class LsofPidfilePinner(PidfilePinner):
+    """
+    Fundamentally wrong, but still better than nothing. Simply reads the file contents and ensures a valid contained pid
+    has the file open via `lsof`.
+    """
+
+    @classmethod
+    def is_available(cls) -> bool:
+        return shutil.which('lsof') is not None
+
+    def _try_read_and_verify(self, pf: Pidfile, timeout: Timeout) -> ta.Optional[int]:
+        if (initial_pid := pf.read()) is None:
+            return None
+
+        lsof_output = LsofCommand(
+            # pid=initial_pid,
+            file=os.path.abspath(pf.path),
+        ).run(
+            timeout=timeout,
+        )
+
+        lsof_pids: ta.Set[int] = set()
+        for li in lsof_output:
+            if li.pid is None:
+                continue
+            try:
+                li_pid = int(li.pid)
+            except ValueError:
+                continue
+            lsof_pids.add(li_pid)
+
+        if initial_pid not in lsof_pids:
+            return None
+
+        if (reread_pid := pf.read()) is None or reread_pid != initial_pid:
+            return None
+
+        return reread_pid
+
+    @contextlib.contextmanager
+    def _pin_pidfile_owner(self, pidfile: Pidfile, timeout: Timeout) -> ta.Iterator[int]:
+        while (pid := self._try_read_and_verify(pidfile, timeout)) is None:
+            time.sleep(self._sleep_s)
+            timeout()
+
+        yield pid
+
+
+##
+
+
+class LslocksPidfdPidfilePinner(PidfilePinner):
+    """
+    Finds the locking pid via `lslocks`, opens a pidfd, then re-runs `lslocks` and rechecks the locking pid is the same.
+    """
+
+    @classmethod
+    def is_available(cls) -> bool:
+        return sys.platform == 'linux' and shutil.which('lslocks') is not None
+
+    def _read_locking_pid(self, path: str, timeout: Timeout) -> int:
+        lsl_output = LslocksCommand().run(timeout=timeout)
+
+        lsl_pids = {
+            li.pid
+            for li in lsl_output
+            if li.path == path
+            and li.type == 'FLOCK'
+        }
+        if not lsl_pids:
+            raise self.NoOwnerError
+        if len(lsl_pids) != 1:
+            raise RuntimeError(f'Multiple locks on file: {path}')
+
+        [pid] = lsl_pids
+        return pid
+
+    class _Result(ta.NamedTuple):
+        pid: int
+        pidfd: int
+
+    def _try_read_and_verify(
+            self,
+            pidfile: Pidfile,
+            timeout: Timeout,
+    ) -> ta.Optional[_Result]:
+        path = os.path.abspath(pidfile.path)
+        initial_pid = self._read_locking_pid(path, timeout)
+
+        try:
+            pidfd = os.open(f'/proc/{initial_pid}', os.O_RDONLY)
+        except FileNotFoundError:
+            raise self.NoOwnerError from None
+
+        try:
+            reread_pid = self._read_locking_pid(path, timeout)
+            if reread_pid != initial_pid:
+                os.close(pidfd)
+                return None
+
+            return self._Result(initial_pid, pidfd)
+
+        except BaseException:
+            os.close(pidfd)
+            raise
+
+    @contextlib.contextmanager
+    def _pin_pidfile_owner(
+            self,
+            pidfile: Pidfile,
+            timeout: Timeout,
+    ) -> ta.Iterator[int]:
+        while (res := self._try_read_and_verify(pidfile, timeout)) is None:
+            time.sleep(self._sleep_s)
+            timeout()
+
+        try:
+            yield res.pid
+        finally:
+            os.close(res.pidfd)
+
+
+##
+
+
+if __name__ == '__main__':
+    def _main() -> None:
+        argparse = __import__('argparse')
+        parser = argparse.ArgumentParser()
+        parser.add_argument('file')
+        args = parser.parse_args()
+
+        with PidfilePinner.default_impl()().pin_pidfile_owner(
+                args.file,
+                timeout=5.,
+        ) as pid:
+            print(pid)
+
+    _main()

omlish/subprocesses/async_.py

@@ -4,6 +4,7 @@ import abc
 import sys
 import typing as ta
 
+from ..lite.timeouts import TimeoutLike
 from .base import BaseSubprocesses
 from .run import SubprocessRun
 from .run import SubprocessRunOutput
@@ -21,7 +22,7 @@ class AbstractAsyncSubprocesses(BaseSubprocesses):
             self,
             *cmd: str,
             input: ta.Any = None,  # noqa
-            timeout: ta.Optional[float] = None,
+            timeout: ta.Optional[TimeoutLike] = None,
             check: bool = False,
             capture_output: ta.Optional[bool] = None,
             **kwargs: ta.Any,

omlish/subprocesses/base.py

@@ -8,6 +8,7 @@ import subprocess
 import time
 import typing as ta
 
+from ..lite.timeouts import Timeout
 from .wrap import subprocess_maybe_shell_wrap_exec
 
 
@@ -111,6 +112,11 @@ class BaseSubprocesses(abc.ABC):  # noqa
 
         #
 
+        if 'timeout' in kwargs:
+            kwargs['timeout'] = Timeout.of(kwargs['timeout']).or_(None)
+
+        #
+
         return cmd, dict(
             env=env,
             shell=shell,