omlish 0.0.0.dev1__py3-none-any.whl → 0.0.0.dev3__py3-none-any.whl

omlish/diag/procfs.py

@@ -0,0 +1,334 @@
+"""
+TODO:
+ - dataclasses
+"""
+import argparse
+import logging
+import os
+import re
+import resource
+import struct
+import sys
+import typing as ta
+
+from .. import iterators as it
+from .. import json
+from .. import lang
+from .. import os as oos
+
+
+log = logging.getLogger(__name__)
+
+
+PidLike = int | str
+
+
+RLIMIT_RESOURCES = {
+    getattr(resource, k): k
+    for k in dir(resource)
+    if k.startswith('RLIMIT_')
+}
+
+
+def parse_size(s: str) -> int:
+    if ' ' not in s:
+        return int(s)
+    us = {'kB': 1024, 'mB': 1024 * 1024}
+    v, u = s.split()
+    return int(v) * us[u]
+
+
+class ProcStat(lang.Namespace):
+    PID = 0
+    COMM = 1
+    STATE = 2
+    PPID = 3
+    PGRP = 4
+    SESSION = 5
+    TTY_NR = 6
+    TPGID = 7
+    FLAGS = 8
+    MINFLT = 9
+    CMINFLT = 10
+    MAJFLT = 11
+    CMAJFLT = 12
+    UTIME = 13
+    STIME = 14
+    CUTIME = 15
+    CSTIME = 16
+    PRIORITY = 17
+    NICE = 18
+    NUM_THREADS = 19
+    ITREALVALUE = 20
+    STARTTIME = 21
+    VSIZE = 22
+    RSS = 23
+    RSSLIM = 24
+    STARTCODE = 25
+    ENDCODE = 26
+    STARTSTACK = 27
+    KSTKESP = 28
+    KSTKEIP = 29
+    SIGNAL = 30
+    BLOCKED = 31
+    SIGIGNORE = 32
+    SIGCATCH = 33
+    WCHAN = 34
+    NSWAP = 35
+    CNSWAP = 36
+    EXIT_SIGNAL = 37
+    PROCESSOR = 38
+    RT_PRIORITY = 39
+    POLICY = 40
+    DELAYACCT_BLKIO_TICKS = 41
+    GUEST_TIME = 42
+    CGUEST_TIME = 43
+    START_DATA = 44
+    END_DATA = 45
+    START_BRK = 46
+    ARG_START = 47
+    ARG_END = 48
+    ENV_START = 49
+    ENV_END = 50
+    EXIT_CODE = 51
+
+
+def _check_linux() -> None:
+    if sys.platform != 'linux':
+        raise OSError
+
+
+def get_process_stats(pid: PidLike = 'self') -> list[str]:
+    """http://man7.org/linux/man-pages/man5/proc.5.html -> /proc/[pid]/stat"""
+
+    _check_linux()
+    with open(f'/proc/{pid}/stat') as f:
+        buf = f.read()
+    l, _, r = buf.rpartition(')')
+    pid, _, comm = l.partition('(')
+    return [pid.strip(), comm, *r.strip().split(' ')]
+
+
+def get_process_chain(pid: PidLike = 'self') -> list[tuple[int, str]]:
+    _check_linux()
+    lst = []
+    while pid:
+        process_stats = get_process_stats(pid)
+        lst.append((int(process_stats[ProcStat.PID]), process_stats[ProcStat.COMM]))
+        pid = int(process_stats[ProcStat.PPID])
+    return lst
+
+
+def get_process_start_time(pid: PidLike = 'self') -> int:
+    """https://stackoverflow.com/questions/2598145/how-to-retrieve-the-process-start-time-or-uptime-in-python"""
+
+    _check_linux()
+    hz = os.sysconf(os.sysconf_names['SC_CLK_TCK'])
+    with open('/proc/stat') as f:
+        system_stats = f.readlines()
+    for line in system_stats:
+        if line.startswith('btime'):
+            boot_timestamp = int(line.split()[1])
+            break
+    else:
+        raise ValueError
+    process_stats = get_process_stats(pid)
+    age_from_boot_jiffies = int(process_stats[ProcStat.STARTTIME])
+    age_from_boot_timestamp = age_from_boot_jiffies // hz
+    return boot_timestamp + age_from_boot_timestamp
+
+
+def get_process_rss(pid: PidLike = 'self') -> int:
+    return int(get_process_stats(pid)[ProcStat.RSS])
+
+
+def set_process_oom_score_adj(score: str, pid: PidLike = 'self') -> None:
+    _check_linux()
+    with open(f'/proc/{pid}/oom_score_adj', 'w') as f:
+        f.write(str(score))
+
+
+MAP_LINE_RX = re.compile(
+    r'^'
+    r'(?P<address>[A-Fa-f0-9]+)-(?P<end_address>[A-Fa-f0-9]+)\s+'
+    r'(?P<permissions>\S+)\s+'
+    r'(?P<offset>[A-Fa-f0-9]+)\s+'
+    r'(?P<device>\S+)\s+'
+    r'(?P<inode>\d+)\s+'
+    r'(?P<path>.*)'
+    r'$'
+)
+
+
+def get_process_maps(pid: PidLike = 'self', sharing: bool = False) -> ta.Iterator[dict[str, ta.Any]]:
+    """http://man7.org/linux/man-pages/man5/proc.5.html -> /proc/[pid]/maps"""
+
+    _check_linux()
+    with open(f'/proc/{pid}/{"smaps" if sharing else "maps"}') as map_file:
+        while True:
+            line = map_file.readline()
+            if not line:
+                break
+            m = MAP_LINE_RX.match(line)
+            if not m:
+                raise ValueError(line)
+            address = int(m.group('address'), 16)
+            end_address = int(m.group('end_address'), 16)
+            d = {
+                'address': address,
+                'end_address': end_address,
+                'size': end_address - address,
+                'permissions': [x for x in m.group('permissions') if x != '-'],
+                'offset': int(m.group('offset'), 16),
+                'device': m.group('device'),
+                'inode': int(m.group('inode')),
+                'path': m.group('path'),
+            }
+            if sharing:
+                s: dict[str, ta.Any] = {}
+                while True:
+                    line = map_file.readline()
+                    k, v = line.split(':')
+                    if k.lower() == 'vmflags':
+                        break
+                    s[k.lower()] = parse_size(v.strip())
+                _, v = line.split(':')
+                s['vmflags'] = [p for p in [j.strip() for j in v.split(' ')] if p]
+                d['sharing'] = s
+            yield d
+
+
+PAGEMAP_KEYS = (
+    'address',
+    'pfn',
+    'swap_type',
+    'swap_offset',
+    'pte_soft_dirty',
+    'file_page_or_shared_anon',
+    'page_swapped',
+    'page_present',
+)
+
+
+def get_process_range_pagemaps(start: int, end: int, pid: PidLike = 'self') -> ta.Iterable[dict[str, int]]:
+    """https://www.kernel.org/doc/Documentation/vm/pagemap.txt"""
+
+    _check_linux()
+    offset = (start // oos.PAGE_SIZE) * 8
+    npages = ((end - start) // oos.PAGE_SIZE)
+    size = npages * 8
+    with open(f'/proc/{pid}/pagemap', 'rb') as pagemap_file:
+        pagemap_file.seek(offset)
+        pagemap_buf = pagemap_file.read(size)
+    if not pagemap_buf:
+        return
+    _struct_unpack = struct.unpack
+    for pagenum in range(npages):
+        [packed] = _struct_unpack('Q', pagemap_buf[pagenum * 8:(pagenum + 1) * 8])
+        yield {
+            'address': start + (pagenum * oos.PAGE_SIZE),
+            'pfn': (packed & ((1 << (54 + 1)) - 1)),
+            'swap_type': (packed & ((1 << (4 + 1)) - 1)),
+            'swap_offset': (packed & ((1 << (54 + 1)) - 1)) >> 5,
+            'pte_soft_dirty': ((packed >> 55) & 1) > 0,
+            'file_page_or_shared_anon': ((packed >> 61) & 1) > 0,
+            'page_swapped': ((packed >> 62) & 1) > 0,
+            'page_present': ((packed >> 63) & 1) > 0,
+        }
+
+
+def get_process_pagemaps(pid: PidLike = 'self') -> ta.Iterable[dict[str, int]]:
+    _check_linux()
+    for m in get_process_maps(pid):
+        yield from get_process_range_pagemaps(m['address'], m['end_address'], pid)
+
+
+def _dump_cmd(args: ta.Any) -> None:
+    total = 0
+    dirty_total = 0
+    for m in get_process_maps(args.pid, sharing=True):
+        total += m['sharing']['rss']
+        sys.stdout.write(json.dumps({'map': m}))
+        sys.stdout.write('\n')
+        for pm in get_process_range_pagemaps(m['address'], m['end_address'], args.pid):
+            if pm['pte_soft_dirty']:
+                dirty_total += oos.PAGE_SIZE
+            sys.stdout.write(json.dumps({'page': tuple(pm[k] for k in PAGEMAP_KEYS)}))
+            sys.stdout.write('\n')
+    dct = {
+        'total': total,
+        'dirty_total': dirty_total,
+    }
+    sys.stdout.write(json.dumps(dct))
+    sys.stdout.write('\n')
+
+
+def _cmp_cmd(args: ta.Any) -> None:
+    if len(args.pids) == 1:
+        [rpid] = args.pids
+        lpid = get_process_chain(rpid)[1][0]
+    elif len(args.pids) == 2:
+        lpid, rpid = args.pids
+    else:
+        raise TypeError('Invalid arguments')
+
+    def g(pid: int) -> ta.Iterator[dict[str, int]]:
+        for m in get_process_maps(pid, sharing=True):
+            yield from get_process_range_pagemaps(m['address'], m['end_address'], pid)
+
+    lpms, rpms = (g(pid) for pid in (lpid, rpid))
+
+    l_pages = 0
+    r_pages = 0
+    c_pages = 0
+    for _, ps in it.merge_on(lambda pm: pm['address'], lpms, rpms):
+        l, r = it.expand_indexed_pairs(ps, None, width=2)
+        if l is not None and r is None:
+            l_pages += 1
+        elif l is None and r is not None:
+            r_pages += 1
+        elif l['pfn'] != r['pfn']:  # type: ignore
+            c_pages += 1
+        else:
+            continue
+        if not args.quiet:
+            sys.stdout.write(json.dumps([l, r]))
+            sys.stdout.write('\n')
+    l_pages += c_pages
+    r_pages += c_pages
+    dct = {
+        'l_pages': l_pages,
+        'l_bytes': l_pages * oos.PAGE_SIZE,
+        'r_pages': r_pages,
+        'r_bytes': r_pages * oos.PAGE_SIZE,
+        'c_pages': c_pages,
+        'c_bytes': c_pages * oos.PAGE_SIZE,
+    }
+    sys.stdout.write(json.dumps(dct))
+    sys.stdout.write('\n')
+
+
+def _main() -> None:
+    _check_linux()
+
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument('-q', '--quiet', action='store_true')
+    arg_subparsers = arg_parser.add_subparsers()
+
+    dump_arg_parser = arg_subparsers.add_parser('dump')
+    dump_arg_parser.add_argument('pid', type=int)
+    dump_arg_parser.set_defaults(func=_dump_cmd)
+
+    cmp_arg_parser = arg_subparsers.add_parser('cmp')
+    cmp_arg_parser.add_argument('pids', type=int, nargs='*')
+    cmp_arg_parser.set_defaults(func=_cmp_cmd)
+
+    args = arg_parser.parse_args()
+    if not hasattr(args, 'func'):
+        arg_parser.print_help()
+    else:
+        args.func(args)
+
+
+if __name__ == '__main__':
+    _main()

omlish/diag/ps.py

@@ -0,0 +1,47 @@
+import dataclasses as dc
+import os
+import subprocess
+
+from .. import lang
+
+
+@dc.dataclass(frozen=True)
+class PsItem:
+    pid: int
+    ppid: int
+    cmd: str
+
+
+def get_ps_item(pid: int, timeout: lang.Timeout | None = None) -> PsItem:
+    timeout = lang.timeout(timeout)
+    out = subprocess.check_output(
+        ['ps', '-o', 'pid=,ppid=,command=', str(int(pid))],
+        timeout=timeout.or_(None),
+    ).decode().strip()
+    opid, _, rest = out.partition(' ')
+    ppid, _, cmd = rest.strip().partition(' ')
+    return PsItem(
+        int(opid),
+        int(ppid),
+        cmd.strip(),
+    )
+
+
+def get_ps_lineage(pid: int, timeout: lang.Timeout | None = None) -> list[PsItem]:
+    timeout = lang.timeout(timeout)
+    ret: list[PsItem] = []
+    while True:
+        cur = get_ps_item(pid, timeout)
+        if cur.ppid < 1:
+            break
+        ret.append(cur)
+        pid = cur.ppid
+    return ret
+
+
+def _main() -> None:
+    print(get_ps_lineage(os.getpid()))
+
+
+if __name__ == '__main__':
+    _main()

omlish/{replserver → diag/replserver}/console.py

@@ -25,13 +25,13 @@ import traceback
 import types
 import typing as ta
 
-from .. import check
+from ... import check
 
 
 log = logging.getLogger(__name__)
 
 
-class DisconnectException(Exception):
+class DisconnectError(Exception):
     pass
 
 
@@ -43,13 +43,13 @@ class InteractiveSocketConsole:
     def __init__(
             self,
             conn: sock.socket,
-            locals: ta.Optional[dict[str, ta.Any]] = None,
-            filename: str = '<console>'
+            locals: dict[str, ta.Any] | None = None,  # noqa
+            filename: str = '<console>',
     ) -> None:
         super().__init__()
 
         if locals is None:
-            locals = {
+            locals = {  # noqa
                 '__name__': '__console__',
                 '__doc__': None,
                 '__console__': self,
@@ -73,7 +73,7 @@ class InteractiveSocketConsole:
 
     CPRT = 'Type "help", "copyright", "credits" or "license" for more information.'
 
-    def interact(self, banner: ta.Optional[str] = None, exitmsg: ta.Optional[str] = None) -> None:
+    def interact(self, banner: str | None = None, exitmsg: str | None = None) -> None:
         log.info(f'Console {id(self)} on thread {threading.current_thread().ident} interacting')
 
         try:
@@ -81,11 +81,9 @@ class InteractiveSocketConsole:
             ps2 = getattr(sys, 'ps2', '... ')
 
             if banner is None:
-                self.write(
-                    'Python %s on %s\n%s\n(%s)\n' %
-                    (sys.version, sys.platform, self.CPRT, self.__class__.__name__))
+                self.write(f'Python {sys.version} on {sys.platform}\n{self.CPRT}\n({self.__class__.__name__})\n')
             elif banner:
-                self.write('%s\n' % (str(banner),))
+                self.write(f'{banner!s}\n')
 
             more = False
             while True:
@@ -104,12 +102,12 @@ class InteractiveSocketConsole:
                     more = False
 
             if exitmsg is None:
-                self.write('now exiting %s...\n' % self.__class__.__name__)
+                self.write(f'now exiting {self.__class__.__name__}...\n')
 
             elif exitmsg != '':
-                self.write('%s\n' % exitmsg)
+                self.write(f'{exitmsg}\n')
 
-        except DisconnectException:
+        except DisconnectError:
             pass
 
         except OSError as oe:
@@ -133,7 +131,7 @@ class InteractiveSocketConsole:
         while True:
             b = self._conn.recv(1)
             if not b:
-                raise DisconnectException
+                raise DisconnectError
             if b == b'\n':
                 break
             buf += b
@@ -144,10 +142,10 @@ class InteractiveSocketConsole:
 
     def compile(
             self,
-            source: ta.Union[str, ast.AST],
+            source: str | ast.AST,
             filename: str = '<input>',
-            symbol: str = 'single'
-    ) -> ta.Optional[types.CodeType]:
+            symbol: str = 'single',
+    ) -> types.CodeType | None:
         if isinstance(source, ast.AST):
             return self._compiler.compiler(source, filename, symbol)  # type: ignore
         else:
@@ -155,7 +153,7 @@ class InteractiveSocketConsole:
 
     def run_source(
             self,
-            source: ta.Union[str, ast.AST],
+            source: str | ast.AST,
             filename: str = '<input>',
             symbol: str = 'single',
     ) -> bool:
@@ -195,7 +193,7 @@ class InteractiveSocketConsole:
                         expr.value,
                         lineno=expr.lineno,
                         col_offset=expr.col_offset,
-                    )
+                    ),
                 ],
             )
             ast.fix_missing_locations(source)
@@ -227,21 +225,21 @@ class InteractiveSocketConsole:
         finally:
             last_tb = ei = None  # type: ignore  # noqa
 
-    def show_syntax_error(self, filename: ta.Optional[str] = None) -> None:
-        type, value, tb = sys.exc_info()
-        sys.last_type = type
-        sys.last_value = value
+    def show_syntax_error(self, filename: str | None = None) -> None:
+        et, e, tb = sys.exc_info()
+        sys.last_type = et
+        sys.last_value = e
         sys.last_traceback = tb
-        if filename and type is SyntaxError:
+        if filename and et is SyntaxError:
             # Work hard to stuff the correct filename in the exception
             try:
-                msg, (dummy_filename, lineno, offset, line) = value.args  # type: ignore
+                msg, (dummy_filename, lineno, offset, line) = e.args  # type: ignore
             except ValueError:
                 # Not the format we expect; leave it alone
                 pass
             else:
                 # Stuff in the right filename
-                value = SyntaxError(msg, (filename, lineno, offset, line))
-                sys.last_value = value
-        lines = traceback.format_exception_only(type, value)
+                e = SyntaxError(msg, (filename, lineno, offset, line))
+                sys.last_value = e
+        lines = traceback.format_exception_only(et, e)
         self.write(''.join(lines))

omlish/{replserver → diag/replserver}/server.py

@@ -23,8 +23,8 @@ import threading
 import typing as ta
 import weakref
 
-from .. import check
-from .. import dataclasses as dc
+from ... import check
+from ... import dataclasses as dc
 from .console import InteractiveSocketConsole
 
 
@@ -38,7 +38,7 @@ class ReplServer:
     @dc.dataclass(frozen=True)
     class Config:
         path: str
-        file_mode: ta.Optional[int] = None
+        file_mode: int | None = None
         poll_interval: float = 0.5
         exit_timeout: float = 10.0
 
@@ -51,7 +51,7 @@ class ReplServer:
         check.not_empty(config.path)
         self._config = check.isinstance(config, ReplServer.Config)
 
-        self._socket: ta.Optional[sock.socket] = None
+        self._socket: sock.socket | None = None
         self._is_running = False
         self._consoles_by_threads: ta.MutableMapping[threading.Thread, InteractiveSocketConsole] = \
             weakref.WeakKeyDictionary()  # noqa
@@ -62,12 +62,12 @@ class ReplServer:
     def path(self) -> str:
         return self._config.path
 
-    def __enter__(self):
+    def __enter__(self) -> ta.Self:
         check.state(not self._is_running)
         check.state(not self._is_shutdown.is_set())
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         if not self._is_shutdown.is_set():
             self.shutdown(True, self._config.exit_timeout)
 
@@ -91,7 +91,7 @@ class ReplServer:
                 while not self._should_shutdown:
                     try:
                         conn, _ = self._socket.accept()
-                    except sock.timeout:
+                    except TimeoutError:
                         continue
 
                     log.info(f'Got repl server connection on file {self._config.path}')
@@ -106,7 +106,7 @@ class ReplServer:
                             log.info(
                                 f'Starting console {id(console)} repl server connection '
                                 f'on file {self._config.path} '
-                                f'on thread {threading.current_thread().ident}'
+                                f'on thread {threading.current_thread().ident}',
                             )
                             self._consoles_by_threads[threading.current_thread()] = console
                             console.interact()
@@ -117,13 +117,13 @@ class ReplServer:
                         name=self.CONNECTION_THREAD_NAME)
                     thread.start()
 
-                for thread, console in self._consoles_by_threads.items():
+                for console in self._consoles_by_threads.values():
                     try:
                         console.conn.close()
                     except Exception:
                         log.exception('Error shutting down')
 
-                for thread in self._consoles_by_threads.keys():
+                for thread in self._consoles_by_threads:
                     try:
                         thread.join(self._config.exit_timeout)
                     except Exception:
@@ -135,12 +135,12 @@ class ReplServer:
                 self._is_shutdown.set()
                 self._is_running = False
 
-    def shutdown(self, block: bool = False, timeout: ta.Optional[float] = None) -> None:
+    def shutdown(self, block: bool = False, timeout: float | None = None) -> None:
         self._should_shutdown = True
         if block:
             self._is_shutdown.wait(timeout=timeout)
 
 
-def run():
+def run() -> None:
     with ReplServer(ReplServer.Config('repl.sock')) as repl_server:
         repl_server.run()

omlish/dispatch/dispatch.py

@@ -1,4 +1,5 @@
 import abc
+import contextlib
 import typing as ta
 import weakref
 
@@ -13,33 +14,32 @@ T = ta.TypeVar('T')
 ##
 
 
-_IMPL_FUNC_CLS_SET_CACHE: ta.MutableMapping[ta.Callable, ta.FrozenSet[type]] = weakref.WeakKeyDictionary()
+_IMPL_FUNC_CLS_SET_CACHE: ta.MutableMapping[ta.Callable, frozenset[type]] = weakref.WeakKeyDictionary()
 
 
-def get_impl_func_cls_set(func: ta.Callable) -> ta.FrozenSet[type]:
-    try:
+def get_impl_func_cls_set(func: ta.Callable) -> frozenset[type]:
+    with contextlib.suppress(KeyError):
         return _IMPL_FUNC_CLS_SET_CACHE[func]
-    except KeyError:
-        pass
 
     ann = getattr(func, '__annotations__', {})
     if not ann:
         raise TypeError(f'Invalid impl func: {func!r}')
 
     _, cls = next(iter(ta.get_type_hints(func).items()))
-    if rfl.is_union_type(cls):
-        ret = frozenset(check.isinstance(arg, type) for arg in ta.get_args(cls))
+    rty = rfl.type_(cls)
+    if isinstance(rty, rfl.Union):
+        ret = frozenset(check.isinstance(arg, type) for arg in rty.args)
     else:
-        ret = frozenset([check.isinstance(cls, type)])
+        ret = frozenset([check.isinstance(rty, type)])
 
     _IMPL_FUNC_CLS_SET_CACHE[func] = ret
     return ret
 
 
-def find_impl(cls: type, registry: ta.Mapping[type, T]) -> ta.Optional[T]:
+def find_impl(cls: type, registry: ta.Mapping[type, T]) -> T | None:
     mro = c3.compose_mro(cls, registry.keys())
 
-    match: ta.Optional[type] = None
+    match: type | None = None
     for t in mro:
         if match is not None:
             # If *match* is an implicit ABC but there is another unrelated, equally matching implicit ABC, refuse the
@@ -71,23 +71,21 @@ class Dispatcher(ta.Generic[T]):
         impls_by_arg_cls: dict[type, T] = {}
         self._impls_by_arg_cls = impls_by_arg_cls
 
-        dispatch_cache: dict[ta.Any, ta.Optional[T]] = {}
+        dispatch_cache: dict[ta.Any, T | None] = {}
         self._get_dispatch_cache = lambda: dispatch_cache
 
         def cache_remove(k, self_ref=weakref.ref(self)):
             if (ref_self := self_ref()) is not None:
                 cache = ref_self._get_dispatch_cache()  # noqa
-                try:
+                with contextlib.suppress(KeyError):
                     del cache[k]
-                except KeyError:
-                    pass
 
         cache_token: ta.Any = None
         self._get_cache_token = lambda: cache_token
 
         weakref_ref_ = weakref.ref
 
-        def dispatch(cls: type) -> ta.Optional[T]:
+        def dispatch(cls: type) -> T | None:
             nonlocal cache_token
 
             if cache_token is not None and (current_token := abc.get_cache_token()) != cache_token:
@@ -132,6 +130,6 @@ class Dispatcher(ta.Generic[T]):
     def cache_size(self) -> int:
         return len(self._get_dispatch_cache())
 
-    dispatch: ta.Callable[[type], ta.Optional[T]]
+    dispatch: ta.Callable[[type], T | None]
 
     register: ta.Callable[[T, ta.Iterable[type]], T]