efro.log

  1# Released under the MIT License. See LICENSE for details.
  2#
  3"""Logging functionality."""
  4from __future__ import annotations
  5
  6import sys
  7import time
  8import asyncio
  9import logging
 10import datetime
 11import itertools
 12from enum import Enum
 13from collections import deque
 14from dataclasses import dataclass, field
 15from typing import TYPE_CHECKING, Annotated, override
 16from threading import Thread, current_thread, Lock
 17
 18from efro.util import utc_now
 19from efro.call import tpartial
 20from efro.terminal import Clr
 21from efro.dataclassio import ioprepped, IOAttrs, dataclass_to_json
 22
 23if TYPE_CHECKING:
 24    from pathlib import Path
 25    from typing import Any, Callable, TextIO
 26
 27
 28class LogLevel(Enum):
 29    """Severity level for a log entry.
 30
 31    These enums have numeric values so they can be compared in severity.
 32    Note that these values are not currently interchangeable with the
 33    logging.ERROR, logging.DEBUG, etc. values.
 34    """
 35
 36    DEBUG = 0
 37    INFO = 1
 38    WARNING = 2
 39    ERROR = 3
 40    CRITICAL = 4
 41
 42    @property
 43    def python_logging_level(self) -> int:
 44        """Give the corresponding logging level."""
 45        return LOG_LEVEL_LEVELNOS[self]
 46
 47    @classmethod
 48    def from_python_logging_level(cls, levelno: int) -> LogLevel:
 49        """Given a Python logging level, return a LogLevel."""
 50        return LEVELNO_LOG_LEVELS[levelno]
 51
 52
 53# Python logging levels from LogLevels
 54LOG_LEVEL_LEVELNOS = {
 55    LogLevel.DEBUG: logging.DEBUG,
 56    LogLevel.INFO: logging.INFO,
 57    LogLevel.WARNING: logging.WARNING,
 58    LogLevel.ERROR: logging.ERROR,
 59    LogLevel.CRITICAL: logging.CRITICAL,
 60}
 61
 62# LogLevels from Python logging levels
 63LEVELNO_LOG_LEVELS = {
 64    logging.DEBUG: LogLevel.DEBUG,
 65    logging.INFO: LogLevel.INFO,
 66    logging.WARNING: LogLevel.WARNING,
 67    logging.ERROR: LogLevel.ERROR,
 68    logging.CRITICAL: LogLevel.CRITICAL,
 69}
 70
 71LEVELNO_COLOR_CODES: dict[int, tuple[str, str]] = {
 72    logging.DEBUG: (Clr.CYN, Clr.RST),
 73    logging.INFO: ('', ''),
 74    logging.WARNING: (Clr.YLW, Clr.RST),
 75    logging.ERROR: (Clr.RED, Clr.RST),
 76    logging.CRITICAL: (Clr.SMAG + Clr.BLD + Clr.BLK, Clr.RST),
 77}
 78
 79
 80@ioprepped
 81@dataclass
 82class LogEntry:
 83    """Single logged message."""
 84
 85    name: Annotated[str, IOAttrs('n', soft_default='root', store_default=False)]
 86    message: Annotated[str, IOAttrs('m')]
 87    level: Annotated[LogLevel, IOAttrs('l')]
 88    time: Annotated[datetime.datetime, IOAttrs('t')]
 89
 90    # We support arbitrary string labels per log entry which can be
 91    # incorporated into custom log processing. To populate this, our
 92    # LogHandler class looks for a 'labels' dict passed in the optional
 93    # 'extra' dict arg to standard Python log calls.
 94    labels: Annotated[dict[str, str], IOAttrs('la', store_default=False)] = (
 95        field(default_factory=dict)
 96    )
 97
 98
 99@ioprepped
100@dataclass
101class LogArchive:
102    """Info and data for a log."""
103
104    # Total number of entries submitted to the log.
105    log_size: Annotated[int, IOAttrs('t')]
106
107    # Offset for the entries contained here.
108    # (10 means our first entry is the 10th in the log, etc.)
109    start_index: Annotated[int, IOAttrs('c')]
110
111    entries: Annotated[list[LogEntry], IOAttrs('e')]
112
113
114class LogHandler(logging.Handler):
115    """Fancy-pants handler for logging output.
116
117    Writes logs to disk in structured json format and echoes them
118    to stdout/stderr with pretty colors.
119    """
120
121    _event_loop: asyncio.AbstractEventLoop
122
123    # IMPORTANT: Any debug prints we do here should ONLY go to echofile.
124    # Otherwise we can get infinite loops as those prints come back to us
125    # as new log entries.
126
127    def __init__(
128        self,
129        path: str | Path | None,
130        echofile: TextIO | None,
131        suppress_non_root_debug: bool,
132        cache_size_limit: int,
133        cache_time_limit: datetime.timedelta | None,
134    ):
135        super().__init__()
136        # pylint: disable=consider-using-with
137        self._file = None if path is None else open(path, 'w', encoding='utf-8')
138        self._echofile = echofile
139        self._callbacks: list[Callable[[LogEntry], None]] = []
140        self._suppress_non_root_debug = suppress_non_root_debug
141        self._file_chunks: dict[str, list[str]] = {'stdout': [], 'stderr': []}
142        self._file_chunk_ship_task: dict[str, asyncio.Task | None] = {
143            'stdout': None,
144            'stderr': None,
145        }
146        self._cache_size = 0
147        assert cache_size_limit >= 0
148        self._cache_size_limit = cache_size_limit
149        self._cache_time_limit = cache_time_limit
150        self._cache = deque[tuple[int, LogEntry]]()
151        self._cache_index_offset = 0
152        self._cache_lock = Lock()
153        self._printed_callback_error = False
154        self._thread_bootstrapped = False
155        self._thread = Thread(target=self._log_thread_main, daemon=True)
156        if __debug__:
157            self._last_slow_emit_warning_time: float | None = None
158        self._thread.start()
159
160        # Spin until our thread is up and running; otherwise we could
161        # wind up trying to push stuff to our event loop before the
162        # loop exists.
163        while not self._thread_bootstrapped:
164            time.sleep(0.001)
165
166    def add_callback(
167        self, call: Callable[[LogEntry], None], feed_existing_logs: bool = False
168    ) -> None:
169        """Add a callback to be run for each LogEntry.
170
171        Note that this callback will always run in a background thread.
172        Passing True for feed_existing_logs will cause all cached logs
173        in the handler to be fed to the callback (still in the
174        background thread though).
175        """
176
177        # Kick this over to our bg thread to add the callback and
178        # process cached entries at the same time to ensure there are no
179        # race conditions that could cause entries to be skipped/etc.
180        self._event_loop.call_soon_threadsafe(
181            tpartial(self._add_callback_in_thread, call, feed_existing_logs)
182        )
183
184    def _add_callback_in_thread(
185        self, call: Callable[[LogEntry], None], feed_existing_logs: bool
186    ) -> None:
187        """Add a callback to be run for each LogEntry.
188
189        Note that this callback will always run in a background thread.
190        Passing True for feed_existing_logs will cause all cached logs
191        in the handler to be fed to the callback (still in the
192        background thread though).
193        """
194        assert current_thread() is self._thread
195        self._callbacks.append(call)
196
197        # Run all of our cached entries through the new callback if desired.
198        if feed_existing_logs and self._cache_size_limit > 0:
199            with self._cache_lock:
200                for _id, entry in self._cache:
201                    self._run_callback_on_entry(call, entry)
202
203    def _log_thread_main(self) -> None:
204        self._event_loop = asyncio.new_event_loop()
205
206        # In our background thread event loop we do a fair amount of
207        # slow synchronous stuff such as mucking with the log cache.
208        # Let's avoid getting tons of warnings about this in debug mode.
209        self._event_loop.slow_callback_duration = 2.0  # Default is 0.1
210
211        # NOTE: if we ever use default threadpool at all we should allow
212        # setting it for our loop.
213        asyncio.set_event_loop(self._event_loop)
214        self._thread_bootstrapped = True
215        try:
216            if self._cache_time_limit is not None:
217                _prunetask = self._event_loop.create_task(
218                    self._time_prune_cache()
219                )
220            self._event_loop.run_forever()
221        except BaseException:
222            # If this ever goes down we're in trouble; we won't be able
223            # to log about it though. Try to make some noise however we
224            # can.
225            print('LogHandler died!!!', file=sys.stderr)
226            import traceback
227
228            traceback.print_exc()
229            raise
230
231    async def _time_prune_cache(self) -> None:
232        assert self._cache_time_limit is not None
233        while bool(True):
234            await asyncio.sleep(61.27)
235            now = utc_now()
236            with self._cache_lock:
237                # Prune the oldest entry as long as there is a first one that
238                # is too old.
239                while (
240                    self._cache
241                    and (now - self._cache[0][1].time) >= self._cache_time_limit
242                ):
243                    popped = self._cache.popleft()
244                    self._cache_size -= popped[0]
245                    self._cache_index_offset += 1
246
247    def get_cached(
248        self, start_index: int = 0, max_entries: int | None = None
249    ) -> LogArchive:
250        """Build and return an archive of cached log entries.
251
252        This will only include entries that have been processed by the
253        background thread, so may not include just-submitted logs or
254        entries for partially written stdout/stderr lines.
255        Entries from the range [start_index:start_index+max_entries]
256        which are still present in the cache will be returned.
257        """
258
259        assert start_index >= 0
260        if max_entries is not None:
261            assert max_entries >= 0
262        with self._cache_lock:
263            # Transform start_index to our present cache space.
264            start_index -= self._cache_index_offset
265            # Calc end-index in our present cache space.
266            end_index = (
267                len(self._cache)
268                if max_entries is None
269                else start_index + max_entries
270            )
271
272            # Clamp both indexes to both ends of our present space.
273            start_index = max(0, min(start_index, len(self._cache)))
274            end_index = max(0, min(end_index, len(self._cache)))
275
276            return LogArchive(
277                log_size=self._cache_index_offset + len(self._cache),
278                start_index=start_index + self._cache_index_offset,
279                entries=self._cache_slice(start_index, end_index),
280            )
281
282    def _cache_slice(
283        self, start: int, end: int, step: int = 1
284    ) -> list[LogEntry]:
285        # Deque doesn't natively support slicing but we can do it manually.
286        # It sounds like rotating the deque and pulling from the beginning
287        # is the most efficient way to do this. The downside is the deque
288        # gets temporarily modified in the process so we need to make sure
289        # we're holding the lock.
290        assert self._cache_lock.locked()
291        cache = self._cache
292        cache.rotate(-start)
293        slc = [e[1] for e in itertools.islice(cache, 0, end - start, step)]
294        cache.rotate(start)
295        return slc
296
297    @classmethod
298    def _is_immutable_log_data(cls, data: Any) -> bool:
299        if isinstance(data, (str, bool, int, float, bytes)):
300            return True
301        if isinstance(data, tuple):
302            return all(cls._is_immutable_log_data(x) for x in data)
303        return False
304
305    def call_in_thread(self, call: Callable[[], Any]) -> None:
306        """Submit a call to be run in the logging background thread."""
307        self._event_loop.call_soon_threadsafe(call)
308
309    @override
310    def emit(self, record: logging.LogRecord) -> None:
311        # pylint: disable=too-many-branches
312        if __debug__:
313            starttime = time.monotonic()
314
315        # Called by logging to send us records.
316
317        # TODO - kill this.
318        if (
319            self._suppress_non_root_debug
320            and record.name != 'root'
321            and record.levelname == 'DEBUG'
322        ):
323            return
324
325        # Optimization: if our log args are all simple immutable values,
326        # we can just kick the whole thing over to our background thread to
327        # be formatted there at our leisure. If anything is mutable and
328        # thus could possibly change between now and then or if we want
329        # to do immediate file echoing then we need to bite the bullet
330        # and do that stuff here at the call site.
331        fast_path = self._echofile is None and self._is_immutable_log_data(
332            record.args
333        )
334
335        # Note: just assuming types are correct here, but they'll be
336        # checked properly when the resulting LogEntry gets exported.
337        labels: dict[str, str] | None = getattr(record, 'labels', None)
338        if labels is None:
339            labels = {}
340
341        if fast_path:
342            if __debug__:
343                formattime = echotime = time.monotonic()
344            self._event_loop.call_soon_threadsafe(
345                tpartial(
346                    self._emit_in_thread,
347                    record.name,
348                    record.levelno,
349                    record.created,
350                    record,
351                    labels,
352                )
353            )
354        else:
355            # Slow case; do formatting and echoing here at the log call
356            # site.
357            msg = self.format(record)
358
359            if __debug__:
360                formattime = time.monotonic()
361
362            # Also immediately print pretty colored output to our echo file
363            # (generally stderr). We do this part here instead of in our bg
364            # thread because the delay can throw off command line prompts or
365            # make tight debugging harder.
366            if self._echofile is not None:
367                ends = LEVELNO_COLOR_CODES.get(record.levelno)
368                namepre = f'{Clr.WHT}{record.name}:{Clr.RST} '
369                if ends is not None:
370                    self._echofile.write(f'{namepre}{ends[0]}{msg}{ends[1]}\n')
371                else:
372                    self._echofile.write(f'{namepre}{msg}\n')
373                self._echofile.flush()
374
375            if __debug__:
376                echotime = time.monotonic()
377
378            self._event_loop.call_soon_threadsafe(
379                tpartial(
380                    self._emit_in_thread,
381                    record.name,
382                    record.levelno,
383                    record.created,
384                    msg,
385                    labels,
386                )
387            )
388
389        if __debug__:
390            # Make noise if we're taking a significant amount of time here.
391            # Limit the noise to once every so often though; otherwise we
392            # could get a feedback loop where every log emit results in a
393            # warning log which results in another, etc.
394            now = time.monotonic()
395            # noinspection PyUnboundLocalVariable
396            duration = now - starttime  # pyright: ignore
397            # noinspection PyUnboundLocalVariable
398            format_duration = formattime - starttime  # pyright: ignore
399            # noinspection PyUnboundLocalVariable
400            echo_duration = echotime - formattime  # pyright: ignore
401            if duration > 0.05 and (
402                self._last_slow_emit_warning_time is None
403                or now > self._last_slow_emit_warning_time + 10.0
404            ):
405                # Logging calls from *within* a logging handler
406                # sounds sketchy, so let's just kick this over to
407                # the bg event loop thread we've already got.
408                self._last_slow_emit_warning_time = now
409                self._event_loop.call_soon_threadsafe(
410                    tpartial(
411                        logging.warning,
412                        'efro.log.LogHandler emit took too long'
413                        ' (%.2fs total; %.2fs format, %.2fs echo,'
414                        ' fast_path=%s).',
415                        duration,
416                        format_duration,
417                        echo_duration,
418                        fast_path,
419                    )
420                )
421
422    def _emit_in_thread(
423        self,
424        name: str,
425        levelno: int,
426        created: float,
427        message: str | logging.LogRecord,
428        labels: dict[str, str],
429    ) -> None:
430        try:
431            # If they passed a raw record here, bake it down to a string.
432            if isinstance(message, logging.LogRecord):
433                message = self.format(message)
434
435            self._emit_entry(
436                LogEntry(
437                    name=name,
438                    message=message,
439                    level=LEVELNO_LOG_LEVELS.get(levelno, LogLevel.INFO),
440                    time=datetime.datetime.fromtimestamp(
441                        created, datetime.timezone.utc
442                    ),
443                    labels=labels,
444                )
445            )
446        except Exception:
447            import traceback
448
449            traceback.print_exc(file=self._echofile)
450
451    def file_write(self, name: str, output: str) -> None:
452        """Send raw stdout/stderr output to the logger to be collated."""
453
454        # Note to self: it turns out that things like '^^^^^^^^^^^^^^'
455        # lines in stack traces get written as lots of individual '^'
456        # writes. It feels a bit dirty to be pushing a deferred call to
457        # another thread for each character. Perhaps should do some sort
458        # of basic accumulation here?
459        self._event_loop.call_soon_threadsafe(
460            tpartial(self._file_write_in_thread, name, output)
461        )
462
463    def _file_write_in_thread(self, name: str, output: str) -> None:
464        try:
465            assert name in ('stdout', 'stderr')
466
467            # Here we try to be somewhat smart about breaking arbitrary
468            # print output into discrete log entries.
469
470            self._file_chunks[name].append(output)
471
472            # Individual parts of a print come across as separate writes,
473            # and the end of a print will be a standalone '\n' by default.
474            # Let's use that as a hint that we're likely at the end of
475            # a full print statement and ship what we've got.
476            if output == '\n':
477                self._ship_file_chunks(name, cancel_ship_task=True)
478            else:
479                # By default just keep adding chunks.
480                # However we keep a timer running anytime we've got
481                # unshipped chunks so that we can ship what we've got
482                # after a short bit if we never get a newline.
483                ship_task = self._file_chunk_ship_task[name]
484                if ship_task is None:
485                    self._file_chunk_ship_task[name] = (
486                        self._event_loop.create_task(
487                            self._ship_chunks_task(name),
488                            name='log ship file chunks',
489                        )
490                    )
491
492        except Exception:
493            import traceback
494
495            traceback.print_exc(file=self._echofile)
496
497    def file_flush(self, name: str) -> None:
498        """Send raw stdout/stderr flush to the logger to be collated."""
499
500        self._event_loop.call_soon_threadsafe(
501            tpartial(self._file_flush_in_thread, name)
502        )
503
504    def _file_flush_in_thread(self, name: str) -> None:
505        try:
506            assert name in ('stdout', 'stderr')
507
508            # Immediately ship whatever chunks we've got.
509            if self._file_chunks[name]:
510                self._ship_file_chunks(name, cancel_ship_task=True)
511
512        except Exception:
513            import traceback
514
515            traceback.print_exc(file=self._echofile)
516
517    async def _ship_chunks_task(self, name: str) -> None:
518        # Note: it's important we sleep here for a moment. Otherwise,
519        # things like '^^^^^^^^^^^^' lines in stack traces, which come
520        # through as lots of individual '^' writes, tend to get broken
521        # into lots of tiny little lines by us.
522        await asyncio.sleep(0.01)
523        self._ship_file_chunks(name, cancel_ship_task=False)
524
525    def _ship_file_chunks(self, name: str, cancel_ship_task: bool) -> None:
526        # Note: Raw print input generally ends in a newline, but that is
527        # redundant when we break things into log entries and results in
528        # extra empty lines. So strip off a single trailing newline if
529        # one is present.
530        text = ''.join(self._file_chunks[name]).removesuffix('\n')
531
532        self._emit_entry(
533            LogEntry(
534                name=name, message=text, level=LogLevel.INFO, time=utc_now()
535            )
536        )
537        self._file_chunks[name] = []
538        ship_task = self._file_chunk_ship_task[name]
539        if cancel_ship_task and ship_task is not None:
540            ship_task.cancel()
541        self._file_chunk_ship_task[name] = None
542
543    def _emit_entry(self, entry: LogEntry) -> None:
544        assert current_thread() is self._thread
545
546        # Store to our cache.
547        if self._cache_size_limit > 0:
548            with self._cache_lock:
549                # Do a rough calc of how many bytes this entry consumes.
550                entry_size = sum(
551                    sys.getsizeof(x)
552                    for x in (
553                        entry,
554                        entry.name,
555                        entry.message,
556                        entry.level,
557                        entry.time,
558                    )
559                )
560                self._cache.append((entry_size, entry))
561                self._cache_size += entry_size
562
563                # Prune old until we are back at or under our limit.
564                while self._cache_size > self._cache_size_limit:
565                    popped = self._cache.popleft()
566                    self._cache_size -= popped[0]
567                    self._cache_index_offset += 1
568
569        # Pass to callbacks.
570        for call in self._callbacks:
571            self._run_callback_on_entry(call, entry)
572
573        # Dump to our structured log file.
574        # TODO: should set a timer for flushing; don't flush every line.
575        if self._file is not None:
576            entry_s = dataclass_to_json(entry)
577            assert '\n' not in entry_s  # Make sure its a single line.
578            print(entry_s, file=self._file, flush=True)
579
580    def _run_callback_on_entry(
581        self, callback: Callable[[LogEntry], None], entry: LogEntry
582    ) -> None:
583        """Run a callback and handle any errors."""
584        try:
585            callback(entry)
586        except Exception:
587            # Only print the first callback error to avoid insanity.
588            if not self._printed_callback_error:
589                import traceback
590
591                traceback.print_exc(file=self._echofile)
592                self._printed_callback_error = True
593
594
595class FileLogEcho:
596    """A file-like object for forwarding stdout/stderr to a LogHandler."""
597
598    def __init__(
599        self, original: TextIO, name: str, handler: LogHandler
600    ) -> None:
601        assert name in ('stdout', 'stderr')
602        self._original = original
603        self._name = name
604        self._handler = handler
605
606    def write(self, output: Any) -> None:
607        """Override standard write call."""
608        self._original.write(output)
609        self._handler.file_write(self._name, output)
610
611    def flush(self) -> None:
612        """Flush the file."""
613        self._original.flush()
614
615        # We also use this as a hint to ship whatever file chunks
616        # we've accumulated (we have to try and be smart about breaking
617        # our arbitrary file output into discrete entries).
618        self._handler.file_flush(self._name)
619
620    def isatty(self) -> bool:
621        """Are we a terminal?"""
622        return self._original.isatty()
623
624
625def setup_logging(
626    log_path: str | Path | None,
627    level: LogLevel,
628    suppress_non_root_debug: bool = False,
629    log_stdout_stderr: bool = False,
630    echo_to_stderr: bool = True,
631    cache_size_limit: int = 0,
632    cache_time_limit: datetime.timedelta | None = None,
633) -> LogHandler:
634    """Set up our logging environment.
635
636    Returns the custom handler which can be used to fetch information
637    about logs that have passed through it. (worst log-levels, caches, etc.).
638    """
639
640    lmap = {
641        LogLevel.DEBUG: logging.DEBUG,
642        LogLevel.INFO: logging.INFO,
643        LogLevel.WARNING: logging.WARNING,
644        LogLevel.ERROR: logging.ERROR,
645        LogLevel.CRITICAL: logging.CRITICAL,
646    }
647
648    # Wire logger output to go to a structured log file.
649    # Also echo it to stderr IF we're running in a terminal.
650    # UPDATE: Actually gonna always go to stderr. Is there a
651    # reason we shouldn't? This makes debugging possible if all
652    # we have is access to a non-interactive terminal or file dump.
653    # We could add a '--quiet' arg or whatnot to change this behavior.
654
655    # Note: by passing in the *original* stderr here before we
656    # (potentially) replace it, we ensure that our log echos
657    # won't themselves be intercepted and sent to the logger
658    # which would create an infinite loop.
659    loghandler = LogHandler(
660        log_path,
661        echofile=sys.stderr if echo_to_stderr else None,
662        suppress_non_root_debug=suppress_non_root_debug,
663        cache_size_limit=cache_size_limit,
664        cache_time_limit=cache_time_limit,
665    )
666
667    # Note: going ahead with force=True here so that we replace any
668    # existing logger. Though we warn if it looks like we are doing
669    # that so we can try to avoid creating the first one.
670    had_previous_handlers = bool(logging.root.handlers)
671    logging.basicConfig(
672        level=lmap[level],
673        # format='%(name)s: %(message)s',
674        # We dump *only* the message here. We pass various log record bits
675        # around and format things fancier where they end up.
676        format='%(message)s',
677        handlers=[loghandler],
678        force=True,
679    )
680    if had_previous_handlers:
681        logging.warning(
682            'setup_logging: Replacing existing handlers.'
683            ' Something may have logged before expected.'
684        )
685
686    # Optionally intercept Python's stdout/stderr output and generate
687    # log entries from it.
688    if log_stdout_stderr:
689        sys.stdout = FileLogEcho(  # type: ignore
690            sys.stdout, 'stdout', loghandler
691        )
692        sys.stderr = FileLogEcho(  # type: ignore
693            sys.stderr, 'stderr', loghandler
694        )
695
696    return loghandler