syncing changes from ballistica/master

dist/ba_data/python/efro/logging.py

@@ -0,0 +1,733 @@
+# Released under the MIT License. See LICENSE for details.
+#
+"""Logging functionality."""
+from __future__ import annotations
+
+import sys
+import time
+import asyncio
+import logging
+import datetime
+import itertools
+from enum import Enum
+from functools import partial
+from collections import deque
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Annotated, override
+from threading import Thread, current_thread, Lock
+
+from efro.util import utc_now
+from efro.terminal import Clr, color_enabled
+from efro.dataclassio import ioprepped, IOAttrs, dataclass_to_json
+
+if TYPE_CHECKING:
+    from pathlib import Path
+    from typing import Any, Callable, TextIO, Literal
+
+
+class LogLevel(Enum):
+    """Severity level for a log entry.
+
+    These enums have numeric values so they can be compared in severity.
+    Note that these values are not currently interchangeable with the
+    logging.ERROR, logging.DEBUG, etc. values.
+    """
+
+    DEBUG = 0
+    INFO = 1
+    WARNING = 2
+    ERROR = 3
+    CRITICAL = 4
+
+    @property
+    def python_logging_level(self) -> int:
+        """Give the corresponding logging level."""
+        return LOG_LEVEL_LEVELNOS[self]
+
+    @classmethod
+    def from_python_logging_level(cls, levelno: int) -> LogLevel:
+        """Given a Python logging level, return a LogLevel."""
+        return LEVELNO_LOG_LEVELS[levelno]
+
+
+# Python logging levels from LogLevels
+LOG_LEVEL_LEVELNOS = {
+    LogLevel.DEBUG: logging.DEBUG,
+    LogLevel.INFO: logging.INFO,
+    LogLevel.WARNING: logging.WARNING,
+    LogLevel.ERROR: logging.ERROR,
+    LogLevel.CRITICAL: logging.CRITICAL,
+}
+
+# LogLevels from Python logging levels
+LEVELNO_LOG_LEVELS = {
+    logging.DEBUG: LogLevel.DEBUG,
+    logging.INFO: LogLevel.INFO,
+    logging.WARNING: LogLevel.WARNING,
+    logging.ERROR: LogLevel.ERROR,
+    logging.CRITICAL: LogLevel.CRITICAL,
+}
+
+LEVELNO_COLOR_CODES: dict[int, tuple[str, str]] = {
+    logging.DEBUG: (Clr.CYN, Clr.RST),
+    logging.INFO: ('', ''),
+    logging.WARNING: (Clr.YLW, Clr.RST),
+    logging.ERROR: (Clr.RED, Clr.RST),
+    logging.CRITICAL: (Clr.SMAG + Clr.BLD + Clr.BLK, Clr.RST),
+}
+
+
+@ioprepped
+@dataclass
+class LogEntry:
+    """Single logged message."""
+
+    name: Annotated[str, IOAttrs('n', soft_default='root', store_default=False)]
+    message: Annotated[str, IOAttrs('m')]
+    level: Annotated[LogLevel, IOAttrs('l')]
+    time: Annotated[datetime.datetime, IOAttrs('t')]
+
+    # We support arbitrary string labels per log entry which can be
+    # incorporated into custom log processing. To populate this, our
+    # LogHandler class looks for a 'labels' dict passed in the optional
+    # 'extra' dict arg to standard Python log calls.
+    labels: Annotated[dict[str, str], IOAttrs('la', store_default=False)] = (
+        field(default_factory=dict)
+    )
+
+
+@ioprepped
+@dataclass
+class LogArchive:
+    """Info and data for a log."""
+
+    # Total number of entries submitted to the log.
+    log_size: Annotated[int, IOAttrs('t')]
+
+    # Offset for the entries contained here.
+    # (10 means our first entry is the 10th in the log, etc.)
+    start_index: Annotated[int, IOAttrs('c')]
+
+    entries: Annotated[list[LogEntry], IOAttrs('e')]
+
+
+class LogHandler(logging.Handler):
+    """Fancy-pants handler for logging output.
+
+    Writes logs to disk in structured json format and echoes them
+    to stdout/stderr with pretty colors.
+    """
+
+    _event_loop: asyncio.AbstractEventLoop
+
+    # IMPORTANT: Any debug prints we do here should ONLY go to echofile.
+    # Otherwise we can get infinite loops as those prints come back to us
+    # as new log entries.
+
+    def __init__(
+        self,
+        *,
+        path: str | Path | None,
+        echofile: TextIO | None,
+        cache_size_limit: int,
+        cache_time_limit: datetime.timedelta | None,
+        echofile_timestamp_format: Literal['default', 'relative'] = 'default',
+        launch_time: float | None = None,
+    ):
+        super().__init__()
+        # pylint: disable=consider-using-with
+        self._file = None if path is None else open(path, 'w', encoding='utf-8')
+        self._echofile = echofile
+        self._echofile_timestamp_format = echofile_timestamp_format
+        self._callbacks: list[Callable[[LogEntry], None]] = []
+        self._file_chunks: dict[str, list[str]] = {'stdout': [], 'stderr': []}
+        self._file_chunk_ship_task: dict[str, asyncio.Task | None] = {
+            'stdout': None,
+            'stderr': None,
+        }
+        self._launch_time = time.time() if launch_time is None else launch_time
+        self._cache_size = 0
+        assert cache_size_limit >= 0
+        self._cache_size_limit = cache_size_limit
+        self._cache_time_limit = cache_time_limit
+        self._cache = deque[tuple[int, LogEntry]]()
+        self._cache_index_offset = 0
+        self._cache_lock = Lock()
+        self._printed_callback_error = False
+        self._thread_bootstrapped = False
+        self._thread = Thread(target=self._log_thread_main, daemon=True)
+        if __debug__:
+            self._last_slow_emit_warning_time: float | None = None
+        self._thread.start()
+
+        # Spin until our thread is up and running; otherwise we could
+        # wind up trying to push stuff to our event loop before the loop
+        # exists.
+        while not self._thread_bootstrapped:
+            time.sleep(0.001)
+
+    def add_callback(
+        self, call: Callable[[LogEntry], None], feed_existing_logs: bool = False
+    ) -> None:
+        """Add a callback to be run for each LogEntry.
+
+        Note that this callback will always run in a background thread.
+        Passing True for feed_existing_logs will cause all cached logs
+        in the handler to be fed to the callback (still in the
+        background thread though).
+        """
+
+        # Kick this over to our bg thread to add the callback and
+        # process cached entries at the same time to ensure there are no
+        # race conditions that could cause entries to be skipped/etc.
+        self._event_loop.call_soon_threadsafe(
+            partial(self._add_callback_in_thread, call, feed_existing_logs)
+        )
+
+    def _add_callback_in_thread(
+        self, call: Callable[[LogEntry], None], feed_existing_logs: bool
+    ) -> None:
+        """Add a callback to be run for each LogEntry.
+
+        Note that this callback will always run in a background thread.
+        Passing True for feed_existing_logs will cause all cached logs
+        in the handler to be fed to the callback (still in the
+        background thread though).
+        """
+        assert current_thread() is self._thread
+        self._callbacks.append(call)
+
+        # Run all of our cached entries through the new callback if desired.
+        if feed_existing_logs and self._cache_size_limit > 0:
+            with self._cache_lock:
+                for _id, entry in self._cache:
+                    self._run_callback_on_entry(call, entry)
+
+    def _log_thread_main(self) -> None:
+        self._event_loop = asyncio.new_event_loop()
+
+        # In our background thread event loop we do a fair amount of
+        # slow synchronous stuff such as mucking with the log cache.
+        # Let's avoid getting tons of warnings about this in debug mode.
+        self._event_loop.slow_callback_duration = 2.0  # Default is 0.1
+
+        # NOTE: if we ever use default threadpool at all we should allow
+        # setting it for our loop.
+        asyncio.set_event_loop(self._event_loop)
+        self._thread_bootstrapped = True
+        try:
+            if self._cache_time_limit is not None:
+                _prunetask = self._event_loop.create_task(
+                    self._time_prune_cache()
+                )
+            self._event_loop.run_forever()
+        except BaseException:
+            # If this ever goes down we're in trouble; we won't be able
+            # to log about it though. Try to make some noise however we
+            # can.
+            print('LogHandler died!!!', file=sys.stderr)
+            import traceback
+
+            traceback.print_exc()
+            raise
+
+    async def _time_prune_cache(self) -> None:
+        assert self._cache_time_limit is not None
+        while bool(True):
+            await asyncio.sleep(61.27)
+            now = utc_now()
+            with self._cache_lock:
+                # Prune the oldest entry as long as there is a first one
+                # that is too old.
+                while (
+                    self._cache
+                    and (now - self._cache[0][1].time) >= self._cache_time_limit
+                ):
+                    popped = self._cache.popleft()
+                    self._cache_size -= popped[0]
+                    self._cache_index_offset += 1
+
+    def get_cached(
+        self, start_index: int = 0, max_entries: int | None = None
+    ) -> LogArchive:
+        """Build and return an archive of cached log entries.
+
+        This will only include entries that have been processed by the
+        background thread, so may not include just-submitted logs or
+        entries for partially written stdout/stderr lines. Entries from
+        the range [start_index:start_index+max_entries] which are still
+        present in the cache will be returned.
+        """
+
+        assert start_index >= 0
+        if max_entries is not None:
+            assert max_entries >= 0
+        with self._cache_lock:
+            # Transform start_index to our present cache space.
+            start_index -= self._cache_index_offset
+            # Calc end-index in our present cache space.
+            end_index = (
+                len(self._cache)
+                if max_entries is None
+                else start_index + max_entries
+            )
+
+            # Clamp both indexes to both ends of our present space.
+            start_index = max(0, min(start_index, len(self._cache)))
+            end_index = max(0, min(end_index, len(self._cache)))
+
+            return LogArchive(
+                log_size=self._cache_index_offset + len(self._cache),
+                start_index=start_index + self._cache_index_offset,
+                entries=self._cache_slice(start_index, end_index),
+            )
+
+    def _cache_slice(
+        self, start: int, end: int, step: int = 1
+    ) -> list[LogEntry]:
+        # Deque doesn't natively support slicing but we can do it
+        # manually. It sounds like rotating the deque and pulling from
+        # the beginning is the most efficient way to do this. The
+        # downside is the deque gets temporarily modified in the process
+        # so we need to make sure we're holding the lock.
+        assert self._cache_lock.locked()
+        cache = self._cache
+        cache.rotate(-start)
+        slc = [e[1] for e in itertools.islice(cache, 0, end - start, step)]
+        cache.rotate(start)
+        return slc
+
+    @classmethod
+    def _is_immutable_log_data(cls, data: Any) -> bool:
+        if isinstance(data, (str, bool, int, float, bytes)):
+            return True
+        if isinstance(data, tuple):
+            return all(cls._is_immutable_log_data(x) for x in data)
+        return False
+
+    def call_in_thread(self, call: Callable[[], Any]) -> None:
+        """Submit a call to be run in the logging background thread."""
+        self._event_loop.call_soon_threadsafe(call)
+
+    @override
+    def emit(self, record: logging.LogRecord) -> None:
+        # pylint: disable=too-many-branches
+        # pylint: disable=too-many-locals
+
+        if __debug__:
+            starttime = time.monotonic()
+
+        # Called by logging to send us records.
+
+        # Optimization: if our log args are all simple immutable values,
+        # we can just kick the whole thing over to our background thread
+        # to be formatted there at our leisure. If anything is mutable
+        # and thus could possibly change between now and then or if we
+        # want to do immediate file echoing then we need to bite the
+        # bullet and do that stuff here at the call site.
+        fast_path = self._echofile is None and self._is_immutable_log_data(
+            record.args
+        )
+
+        # Note: just assuming types are correct here, but they'll be
+        # checked properly when the resulting LogEntry gets exported.
+        labels: dict[str, str] | None = getattr(record, 'labels', None)
+        if labels is None:
+            labels = {}
+
+        if fast_path:
+            if __debug__:
+                formattime = echotime = time.monotonic()
+            self._event_loop.call_soon_threadsafe(
+                partial(
+                    self._emit_in_thread,
+                    record.name,
+                    record.levelno,
+                    record.created,
+                    record,
+                    labels,
+                )
+            )
+        else:
+            # Slow case; do formatting and echoing here at the log call
+            # site.
+            msg = self.format(record)
+
+            if __debug__:
+                formattime = time.monotonic()
+
+            # Also immediately print pretty colored output to our echo
+            # file (generally stderr). We do this part here instead of
+            # in our bg thread because the delay can throw off command
+            # line prompts or make tight debugging harder.
+            if self._echofile is not None:
+                if self._echofile_timestamp_format == 'relative':
+                    timestamp = f'{record.created - self._launch_time:.3f}'
+                else:
+                    timestamp = (
+                        datetime.datetime.fromtimestamp(
+                            record.created, tz=datetime.UTC
+                        ).strftime('%H:%M:%S')
+                        + f'.{int(record.msecs):03d}'
+                    )
+
+                # If color printing is disabled, show level through text
+                # instead of color.
+                lvlnameex = (
+                    ''
+                    if color_enabled
+                    else f' {logging.getLevelName(record.levelno)}'
+                )
+
+                preinfo = (
+                    f'{Clr.WHT}{timestamp}{lvlnameex} {record.name}:'
+                    f'{Clr.RST} '
+                )
+                ends = LEVELNO_COLOR_CODES.get(record.levelno)
+                if ends is not None:
+                    self._echofile.write(f'{preinfo}{ends[0]}{msg}{ends[1]}\n')
+                else:
+                    self._echofile.write(f'{preinfo}{msg}\n')
+                self._echofile.flush()
+
+            if __debug__:
+                echotime = time.monotonic()
+
+            self._event_loop.call_soon_threadsafe(
+                partial(
+                    self._emit_in_thread,
+                    record.name,
+                    record.levelno,
+                    record.created,
+                    msg,
+                    labels,
+                )
+            )
+
+        if __debug__:
+            # pylint: disable=used-before-assignment
+            #
+            # Make noise if we're taking a significant amount of time
+            # here. Limit the noise to once every so often though;
+            # otherwise we could get a feedback loop where every log
+            # emit results in a warning log which results in another,
+            # etc.
+            now = time.monotonic()
+            duration = now - starttime
+            format_duration = formattime - starttime
+            echo_duration = echotime - formattime
+            if duration > 0.05 and (
+                self._last_slow_emit_warning_time is None
+                or now > self._last_slow_emit_warning_time + 10.0
+            ):
+                # Logging calls from *within* a logging handler sounds
+                # sketchy, so let's just kick this over to the bg event
+                # loop thread we've already got.
+                self._last_slow_emit_warning_time = now
+                self._event_loop.call_soon_threadsafe(
+                    partial(
+                        logging.warning,
+                        'efro.logging.LogHandler emit took too long'
+                        ' (%.3fs total; %.3fs format, %.3fs echo,'
+                        ' fast_path=%s).',
+                        duration,
+                        format_duration,
+                        echo_duration,
+                        fast_path,
+                    )
+                )
+
+    def _emit_in_thread(
+        self,
+        name: str,
+        levelno: int,
+        created: float,
+        message: str | logging.LogRecord,
+        labels: dict[str, str],
+    ) -> None:
+        # pylint: disable=too-many-positional-arguments
+        try:
+            # If they passed a raw record here, bake it down to a string.
+            if isinstance(message, logging.LogRecord):
+                message = self.format(message)
+
+            self._emit_entry(
+                LogEntry(
+                    name=name,
+                    message=message,
+                    level=LEVELNO_LOG_LEVELS.get(levelno, LogLevel.INFO),
+                    time=datetime.datetime.fromtimestamp(
+                        created, datetime.timezone.utc
+                    ),
+                    labels=labels,
+                )
+            )
+        except Exception:
+            import traceback
+
+            traceback.print_exc(file=self._echofile)
+
+    def file_write(self, name: str, output: str) -> None:
+        """Send raw stdout/stderr output to the logger to be collated."""
+
+        # Note to self: it turns out that things like '^^^^^^^^^^^^^^'
+        # lines in stack traces get written as lots of individual '^'
+        # writes. It feels a bit dirty to be pushing a deferred call to
+        # another thread for each character. Perhaps should do some sort
+        # of basic accumulation here?
+        self._event_loop.call_soon_threadsafe(
+            partial(self._file_write_in_thread, name, output)
+        )
+
+    def _file_write_in_thread(self, name: str, output: str) -> None:
+        try:
+            assert name in ('stdout', 'stderr')
+
+            # Here we try to be somewhat smart about breaking arbitrary
+            # print output into discrete log entries.
+
+            self._file_chunks[name].append(output)
+
+            # Individual parts of a print come across as separate
+            # writes, and the end of a print will be a standalone '\n'
+            # by default. Let's use that as a hint that we're likely at
+            # the end of a full print statement and ship what we've got.
+            if output == '\n':
+                self._ship_file_chunks(name, cancel_ship_task=True)
+            else:
+                # By default just keep adding chunks. However we keep a
+                # timer running anytime we've got unshipped chunks so
+                # that we can ship what we've got after a short bit if
+                # we never get a newline.
+                ship_task = self._file_chunk_ship_task[name]
+                if ship_task is None:
+                    self._file_chunk_ship_task[name] = (
+                        self._event_loop.create_task(
+                            self._ship_chunks_task(name),
+                            name='log ship file chunks',
+                        )
+                    )
+
+        except Exception:
+            import traceback
+
+            traceback.print_exc(file=self._echofile)
+
+    def shutdown(self) -> None:
+        """Block until all pending logs/prints are done."""
+        done = False
+        self.file_flush('stdout')
+        self.file_flush('stderr')
+
+        def _set_done() -> None:
+            nonlocal done
+            done = True
+
+        self._event_loop.call_soon_threadsafe(_set_done)
+
+        starttime = time.monotonic()
+        while not done:
+            if time.monotonic() - starttime > 5.0:
+                print('LogHandler shutdown hung!!!', file=sys.stderr)
+                break
+            time.sleep(0.01)
+
+    def file_flush(self, name: str) -> None:
+        """Send raw stdout/stderr flush to the logger to be collated."""
+
+        self._event_loop.call_soon_threadsafe(
+            partial(self._file_flush_in_thread, name)
+        )
+
+    def _file_flush_in_thread(self, name: str) -> None:
+        try:
+            assert name in ('stdout', 'stderr')
+
+            # Immediately ship whatever chunks we've got.
+            if self._file_chunks[name]:
+                self._ship_file_chunks(name, cancel_ship_task=True)
+
+        except Exception:
+            import traceback
+
+            traceback.print_exc(file=self._echofile)
+
+    async def _ship_chunks_task(self, name: str) -> None:
+        # Note: it's important we sleep here for a moment. Otherwise,
+        # things like '^^^^^^^^^^^^' lines in stack traces, which come
+        # through as lots of individual '^' writes, tend to get broken
+        # into lots of tiny little lines by us.
+        await asyncio.sleep(0.01)
+        self._ship_file_chunks(name, cancel_ship_task=False)
+
+    def _ship_file_chunks(self, name: str, cancel_ship_task: bool) -> None:
+        # Note: Raw print input generally ends in a newline, but that is
+        # redundant when we break things into log entries and results in
+        # extra empty lines. So strip off a single trailing newline if
+        # one is present.
+        text = ''.join(self._file_chunks[name]).removesuffix('\n')
+
+        self._emit_entry(
+            LogEntry(
+                name=name, message=text, level=LogLevel.INFO, time=utc_now()
+            )
+        )
+        self._file_chunks[name] = []
+        ship_task = self._file_chunk_ship_task[name]
+        if cancel_ship_task and ship_task is not None:
+            ship_task.cancel()
+        self._file_chunk_ship_task[name] = None
+
+    def _emit_entry(self, entry: LogEntry) -> None:
+        assert current_thread() is self._thread
+
+        # Store to our cache.
+        if self._cache_size_limit > 0:
+            with self._cache_lock:
+                # Do a rough calc of how many bytes this entry consumes.
+                entry_size = sum(
+                    sys.getsizeof(x)
+                    for x in (
+                        entry,
+                        entry.name,
+                        entry.message,
+                        entry.level,
+                        entry.time,
+                    )
+                )
+                self._cache.append((entry_size, entry))
+                self._cache_size += entry_size
+
+                # Prune old until we are back at or under our limit.
+                while self._cache_size > self._cache_size_limit:
+                    popped = self._cache.popleft()
+                    self._cache_size -= popped[0]
+                    self._cache_index_offset += 1
+
+        # Pass to callbacks.
+        for call in self._callbacks:
+            self._run_callback_on_entry(call, entry)
+
+        # Dump to our structured log file.
+        #
+        # TODO: should set a timer for flushing; don't flush every line.
+        if self._file is not None:
+            entry_s = dataclass_to_json(entry)
+            assert '\n' not in entry_s  # Make sure its a single line.
+            print(entry_s, file=self._file, flush=True)
+
+    def _run_callback_on_entry(
+        self, callback: Callable[[LogEntry], None], entry: LogEntry
+    ) -> None:
+        """Run a callback and handle any errors."""
+        try:
+            callback(entry)
+        except Exception:
+            # Only print the first callback error to avoid insanity.
+            if not self._printed_callback_error:
+                import traceback
+
+                traceback.print_exc(file=self._echofile)
+                self._printed_callback_error = True
+
+
+class FileLogEcho:
+    """A file-like object for forwarding stdout/stderr to a LogHandler."""
+
+    def __init__(
+        self, original: TextIO, name: str, handler: LogHandler
+    ) -> None:
+        assert name in ('stdout', 'stderr')
+        self._original = original
+        self._name = name
+        self._handler = handler
+
+    def write(self, output: Any) -> None:
+        """Override standard write call."""
+        self._original.write(output)
+        self._handler.file_write(self._name, output)
+
+    def flush(self) -> None:
+        """Flush the file."""
+        self._original.flush()
+
+        # We also use this as a hint to ship whatever file chunks
+        # we've accumulated (we have to try and be smart about breaking
+        # our arbitrary file output into discrete entries).
+        self._handler.file_flush(self._name)
+
+    def isatty(self) -> bool:
+        """Are we a terminal?"""
+        return self._original.isatty()
+
+
+def setup_logging(
+    log_path: str | Path | None,
+    level: LogLevel,
+    *,
+    log_stdout_stderr: bool = False,
+    echo_to_stderr: bool = True,
+    cache_size_limit: int = 0,
+    cache_time_limit: datetime.timedelta | None = None,
+    launch_time: float | None = None,
+) -> LogHandler:
+    """Set up our logging environment.
+
+    Returns the custom handler which can be used to fetch information
+    about logs that have passed through it. (worst log-levels, caches, etc.).
+    """
+
+    lmap = {
+        LogLevel.DEBUG: logging.DEBUG,
+        LogLevel.INFO: logging.INFO,
+        LogLevel.WARNING: logging.WARNING,
+        LogLevel.ERROR: logging.ERROR,
+        LogLevel.CRITICAL: logging.CRITICAL,
+    }
+
+    # Wire logger output to go to a structured log file. Also echo it to
+    # stderr IF we're running in a terminal.
+    #
+    # UPDATE: Actually gonna always go to stderr. Is there a reason we
+    # shouldn't? This makes debugging possible if all we have is access
+    # to a non-interactive terminal or file dump. We could add a
+    # '--quiet' arg or whatnot to change this behavior.
+
+    # Note: by passing in the *original* stderr here before we
+    # (potentially) replace it, we ensure that our log echos won't
+    # themselves be intercepted and sent to the logger which would
+    # create an infinite loop.
+    loghandler = LogHandler(
+        path=log_path,
+        echofile=sys.stderr if echo_to_stderr else None,
+        echofile_timestamp_format='relative',
+        cache_size_limit=cache_size_limit,
+        cache_time_limit=cache_time_limit,
+        launch_time=launch_time,
+    )
+
+    # Note: going ahead with force=True here so that we replace any
+    # existing logger. Though we warn if it looks like we are doing that
+    # so we can try to avoid creating the first one.
+    had_previous_handlers = bool(logging.root.handlers)
+    logging.basicConfig(
+        level=lmap[level],
+        # We dump *only* the message here. We pass various log record
+        # bits around so we can write rich logs or format things later.
+        format='%(message)s',
+        handlers=[loghandler],
+        force=True,
+    )
+    if had_previous_handlers:
+        logging.warning(
+            'setup_logging: Replacing existing handlers.'
+            ' Something may have logged before expected.'
+        )
+
+    # Optionally intercept Python's stdout/stderr output and generate
+    # log entries from it.
+    if log_stdout_stderr:
+        sys.stdout = FileLogEcho(sys.stdout, 'stdout', loghandler)
+        sys.stderr = FileLogEcho(sys.stderr, 'stderr', loghandler)
+
+    return loghandler