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
 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[
 95        dict[str, str], IOAttrs('la', store_default=False)
 96    ] = field(default_factory=dict)
 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 emit(self, record: logging.LogRecord) -> None:
306        # pylint: disable=too-many-branches
307        if __debug__:
308            starttime = time.monotonic()
309
310        # Called by logging to send us records.
311
312        # TODO - kill this.
313        if (
314            self._suppress_non_root_debug
315            and record.name != 'root'
316            and record.levelname == 'DEBUG'
317        ):
318            return
319
320        # Optimization: if our log args are all simple immutable values,
321        # we can just kick the whole thing over to our background thread to
322        # be formatted there at our leisure. If anything is mutable and
323        # thus could possibly change between now and then or if we want
324        # to do immediate file echoing then we need to bite the bullet
325        # and do that stuff here at the call site.
326        fast_path = self._echofile is None and self._is_immutable_log_data(
327            record.args
328        )
329
330        # Note: just assuming types are correct here, but they'll be
331        # checked properly when the resulting LogEntry gets exported.
332        labels: dict[str, str] | None = getattr(record, 'labels', None)
333        if labels is None:
334            labels = {}
335
336        if fast_path:
337            if __debug__:
338                formattime = echotime = time.monotonic()
339            self._event_loop.call_soon_threadsafe(
340                tpartial(
341                    self._emit_in_thread,
342                    record.name,
343                    record.levelno,
344                    record.created,
345                    record,
346                    labels,
347                )
348            )
349        else:
350            # Slow case; do formatting and echoing here at the log call
351            # site.
352            msg = self.format(record)
353
354            if __debug__:
355                formattime = time.monotonic()
356
357            # Also immediately print pretty colored output to our echo file
358            # (generally stderr). We do this part here instead of in our bg
359            # thread because the delay can throw off command line prompts or
360            # make tight debugging harder.
361            if self._echofile is not None:
362                ends = LEVELNO_COLOR_CODES.get(record.levelno)
363                namepre = f'{Clr.WHT}{record.name}:{Clr.RST} '
364                if ends is not None:
365                    self._echofile.write(f'{namepre}{ends[0]}{msg}{ends[1]}\n')
366                else:
367                    self._echofile.write(f'{namepre}{msg}\n')
368                self._echofile.flush()
369
370            if __debug__:
371                echotime = time.monotonic()
372
373            self._event_loop.call_soon_threadsafe(
374                tpartial(
375                    self._emit_in_thread,
376                    record.name,
377                    record.levelno,
378                    record.created,
379                    msg,
380                    labels,
381                )
382            )
383
384        if __debug__:
385            # Make noise if we're taking a significant amount of time here.
386            # Limit the noise to once every so often though; otherwise we
387            # could get a feedback loop where every log emit results in a
388            # warning log which results in another, etc.
389            now = time.monotonic()
390            # noinspection PyUnboundLocalVariable
391            duration = now - starttime  # pyright: ignore
392            # noinspection PyUnboundLocalVariable
393            format_duration = formattime - starttime  # pyright: ignore
394            # noinspection PyUnboundLocalVariable
395            echo_duration = echotime - formattime  # pyright: ignore
396            if duration > 0.05 and (
397                self._last_slow_emit_warning_time is None
398                or now > self._last_slow_emit_warning_time + 10.0
399            ):
400                # Logging calls from *within* a logging handler
401                # sounds sketchy, so let's just kick this over to
402                # the bg event loop thread we've already got.
403                self._last_slow_emit_warning_time = now
404                self._event_loop.call_soon_threadsafe(
405                    tpartial(
406                        logging.warning,
407                        'efro.log.LogHandler emit took too long'
408                        ' (%.2fs total; %.2fs format, %.2fs echo,'
409                        ' fast_path=%s).',
410                        duration,
411                        format_duration,
412                        echo_duration,
413                        fast_path,
414                    )
415                )
416
417    def _emit_in_thread(
418        self,
419        name: str,
420        levelno: int,
421        created: float,
422        message: str | logging.LogRecord,
423        labels: dict[str, str],
424    ) -> None:
425        try:
426            # If they passed a raw record here, bake it down to a string.
427            if isinstance(message, logging.LogRecord):
428                message = self.format(message)
429
430            self._emit_entry(
431                LogEntry(
432                    name=name,
433                    message=message,
434                    level=LEVELNO_LOG_LEVELS.get(levelno, LogLevel.INFO),
435                    time=datetime.datetime.fromtimestamp(
436                        created, datetime.timezone.utc
437                    ),
438                    labels=labels,
439                )
440            )
441        except Exception:
442            import traceback
443
444            traceback.print_exc(file=self._echofile)
445
446    def file_write(self, name: str, output: str) -> None:
447        """Send raw stdout/stderr output to the logger to be collated."""
448
449        self._event_loop.call_soon_threadsafe(
450            tpartial(self._file_write_in_thread, name, output)
451        )
452
453    def _file_write_in_thread(self, name: str, output: str) -> None:
454        try:
455            assert name in ('stdout', 'stderr')
456
457            # Here we try to be somewhat smart about breaking arbitrary
458            # print output into discrete log entries.
459
460            self._file_chunks[name].append(output)
461
462            # Individual parts of a print come across as separate writes,
463            # and the end of a print will be a standalone '\n' by default.
464            # Let's use that as a hint that we're likely at the end of
465            # a full print statement and ship what we've got.
466            if output == '\n':
467                self._ship_file_chunks(name, cancel_ship_task=True)
468            else:
469                # By default just keep adding chunks.
470                # However we keep a timer running anytime we've got
471                # unshipped chunks so that we can ship what we've got
472                # after a short bit if we never get a newline.
473                ship_task = self._file_chunk_ship_task[name]
474                if ship_task is None:
475                    self._file_chunk_ship_task[
476                        name
477                    ] = self._event_loop.create_task(
478                        self._ship_chunks_task(name),
479                        name='log ship file chunks',
480                    )
481
482        except Exception:
483            import traceback
484
485            traceback.print_exc(file=self._echofile)
486
487    def file_flush(self, name: str) -> None:
488        """Send raw stdout/stderr flush to the logger to be collated."""
489
490        self._event_loop.call_soon_threadsafe(
491            tpartial(self._file_flush_in_thread, name)
492        )
493
494    def _file_flush_in_thread(self, name: str) -> None:
495        try:
496            assert name in ('stdout', 'stderr')
497
498            # Immediately ship whatever chunks we've got.
499            if self._file_chunks[name]:
500                self._ship_file_chunks(name, cancel_ship_task=True)
501
502        except Exception:
503            import traceback
504
505            traceback.print_exc(file=self._echofile)
506
507    async def _ship_chunks_task(self, name: str) -> None:
508        self._ship_file_chunks(name, cancel_ship_task=False)
509
510    def _ship_file_chunks(self, name: str, cancel_ship_task: bool) -> None:
511        # Note: Raw print input generally ends in a newline, but that is
512        # redundant when we break things into log entries and results
513        # in extra empty lines. So strip off a single trailing newline.
514        text = ''.join(self._file_chunks[name]).removesuffix('\n')
515
516        self._emit_entry(
517            LogEntry(
518                name=name, message=text, level=LogLevel.INFO, time=utc_now()
519            )
520        )
521        self._file_chunks[name] = []
522        ship_task = self._file_chunk_ship_task[name]
523        if cancel_ship_task and ship_task is not None:
524            ship_task.cancel()
525        self._file_chunk_ship_task[name] = None
526
527    def _emit_entry(self, entry: LogEntry) -> None:
528        assert current_thread() is self._thread
529
530        # Store to our cache.
531        if self._cache_size_limit > 0:
532            with self._cache_lock:
533                # Do a rough calc of how many bytes this entry consumes.
534                entry_size = sum(
535                    sys.getsizeof(x)
536                    for x in (
537                        entry,
538                        entry.name,
539                        entry.message,
540                        entry.level,
541                        entry.time,
542                    )
543                )
544                self._cache.append((entry_size, entry))
545                self._cache_size += entry_size
546
547                # Prune old until we are back at or under our limit.
548                while self._cache_size > self._cache_size_limit:
549                    popped = self._cache.popleft()
550                    self._cache_size -= popped[0]
551                    self._cache_index_offset += 1
552
553        # Pass to callbacks.
554        for call in self._callbacks:
555            self._run_callback_on_entry(call, entry)
556
557        # Dump to our structured log file.
558        # TODO: should set a timer for flushing; don't flush every line.
559        if self._file is not None:
560            entry_s = dataclass_to_json(entry)
561            assert '\n' not in entry_s  # Make sure its a single line.
562            print(entry_s, file=self._file, flush=True)
563
564    def _run_callback_on_entry(
565        self, callback: Callable[[LogEntry], None], entry: LogEntry
566    ) -> None:
567        """Run a callback and handle any errors."""
568        try:
569            callback(entry)
570        except Exception:
571            # Only print the first callback error to avoid insanity.
572            if not self._printed_callback_error:
573                import traceback
574
575                traceback.print_exc(file=self._echofile)
576                self._printed_callback_error = True
577
578
579class FileLogEcho:
580    """A file-like object for forwarding stdout/stderr to a LogHandler."""
581
582    def __init__(
583        self, original: TextIO, name: str, handler: LogHandler
584    ) -> None:
585        assert name in ('stdout', 'stderr')
586        self._original = original
587        self._name = name
588        self._handler = handler
589
590    def write(self, output: Any) -> None:
591        """Override standard write call."""
592        self._original.write(output)
593        self._handler.file_write(self._name, output)
594
595    def flush(self) -> None:
596        """Flush the file."""
597        self._original.flush()
598
599        # We also use this as a hint to ship whatever file chunks
600        # we've accumulated (we have to try and be smart about breaking
601        # our arbitrary file output into discrete entries).
602        self._handler.file_flush(self._name)
603
604    def isatty(self) -> bool:
605        """Are we a terminal?"""
606        return self._original.isatty()
607
608
609def setup_logging(
610    log_path: str | Path | None,
611    level: LogLevel,
612    suppress_non_root_debug: bool = False,
613    log_stdout_stderr: bool = False,
614    echo_to_stderr: bool = True,
615    cache_size_limit: int = 0,
616    cache_time_limit: datetime.timedelta | None = None,
617) -> LogHandler:
618    """Set up our logging environment.
619
620    Returns the custom handler which can be used to fetch information
621    about logs that have passed through it. (worst log-levels, caches, etc.).
622    """
623
624    lmap = {
625        LogLevel.DEBUG: logging.DEBUG,
626        LogLevel.INFO: logging.INFO,
627        LogLevel.WARNING: logging.WARNING,
628        LogLevel.ERROR: logging.ERROR,
629        LogLevel.CRITICAL: logging.CRITICAL,
630    }
631
632    # Wire logger output to go to a structured log file.
633    # Also echo it to stderr IF we're running in a terminal.
634    # UPDATE: Actually gonna always go to stderr. Is there a
635    # reason we shouldn't? This makes debugging possible if all
636    # we have is access to a non-interactive terminal or file dump.
637    # We could add a '--quiet' arg or whatnot to change this behavior.
638
639    # Note: by passing in the *original* stderr here before we
640    # (potentially) replace it, we ensure that our log echos
641    # won't themselves be intercepted and sent to the logger
642    # which would create an infinite loop.
643    loghandler = LogHandler(
644        log_path,
645        echofile=sys.stderr if echo_to_stderr else None,
646        suppress_non_root_debug=suppress_non_root_debug,
647        cache_size_limit=cache_size_limit,
648        cache_time_limit=cache_time_limit,
649    )
650
651    # Note: going ahead with force=True here so that we replace any
652    # existing logger. Though we warn if it looks like we are doing
653    # that so we can try to avoid creating the first one.
654    had_previous_handlers = bool(logging.root.handlers)
655    logging.basicConfig(
656        level=lmap[level],
657        # format='%(name)s: %(message)s',
658        # We dump *only* the message here. We pass various log record bits
659        # around and format things fancier where they end up.
660        format='%(message)s',
661        handlers=[loghandler],
662        force=True,
663    )
664    if had_previous_handlers:
665        logging.warning(
666            'setup_logging: Replacing existing handlers.'
667            ' Something may have logged before expected.'
668        )
669
670    # Optionally intercept Python's stdout/stderr output and generate
671    # log entries from it.
672    if log_stdout_stderr:
673        sys.stdout = FileLogEcho(  # type: ignore
674            sys.stdout, 'stdout', loghandler
675        )
676        sys.stderr = FileLogEcho(  # type: ignore
677            sys.stderr, 'stderr', loghandler
678        )
679
680    return loghandler