Adding more logic

Author: leandrodamascena

aws_lambda_powertools/logging/buffer/cache.py

@@ -20,6 +20,7 @@ def __init__(self, max_size_bytes: int):
         self.max_size_bytes: int = max_size_bytes
         self.cache: dict[str, deque] = {}
         self.current_size: dict[str, int] = {}
+        self.has_evicted: bool = False
 
     def add(self, key: str, item: Any) -> None:
         """
@@ -53,6 +54,7 @@ def add(self, key: str, item: Any) -> None:
         while self.current_size[key] + item_size > self.max_size_bytes and self.cache[key]:
             removed_item = self.cache[key].popleft()
             self.current_size[key] -= len(str(removed_item))
+            self.has_evicted = True
 
         self.cache[key].append(item)
         self.current_size[key] += item_size

aws_lambda_powertools/logging/constants.py

@@ -3,3 +3,5 @@
 # logger.init attribute is set when Logger has been configured
 LOGGER_ATTRIBUTE_PRECONFIGURED = "init"
 LOGGER_ATTRIBUTE_HANDLER = "logger_handler"
+
+LOGGER_BUFFER_FIRST_INVOKE = "FIRST_INVOKE"

aws_lambda_powertools/logging/exceptions.py

@@ -12,3 +12,11 @@ class OrphanedChildLoggerError(Exception):
     """
 
     pass
+
+
+class InvalidBufferItem(Exception):
+    """
+    Raised when a buffer item exceeds the maximum allowed buffer size.
+    """
+
+    pass

aws_lambda_powertools/logging/logger.py

@@ -22,8 +22,13 @@
     LOGGER_ATTRIBUTE_HANDLER,
     LOGGER_ATTRIBUTE_POWERTOOLS_HANDLER,
     LOGGER_ATTRIBUTE_PRECONFIGURED,
+    LOGGER_BUFFER_FIRST_INVOKE,
+)
+from aws_lambda_powertools.logging.exceptions import (
+    InvalidBufferItem,
+    InvalidLoggerSamplingRateError,
+    OrphanedChildLoggerError,
 )
-from aws_lambda_powertools.logging.exceptions import InvalidLoggerSamplingRateError, OrphanedChildLoggerError
 from aws_lambda_powertools.logging.filters import SuppressFilter
 from aws_lambda_powertools.logging.formatter import (
     RESERVED_FORMATTER_CUSTOM_KEYS,
@@ -34,10 +39,12 @@
 from aws_lambda_powertools.shared import constants
 from aws_lambda_powertools.shared.functions import (
     extract_event_from_common_models,
+    get_tracer_id,
     resolve_env_var_choice,
     resolve_truthy_env_var_choice,
 )
 from aws_lambda_powertools.utilities import jmespath_utils
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
 if TYPE_CHECKING:
     from aws_lambda_powertools.logging.buffer.config import LoggerBufferConfig
@@ -518,14 +525,62 @@ def _add_log_record_to_buffer(
         level: int,
         msg: object,
         args: object,
-        exc_info: logging._ExcInfoType,
-        stack_info: bool,
-        extra: Mapping[str, object],
+        exc_info: logging._ExcInfoType = None,
+        stack_info: bool = False,
+        extra: Mapping[str, object] | None = None,
     ):
-        tracer_id = os.getenv(constants.XRAY_TRACE_ID_ENV, None)
-        if tracer_id:
+        """
+        Add log record to buffer with intelligent tracer ID handling.
+
+        Parameters
+        ----------
+        level : int
+            Logging level of the record.
+        msg : object
+            Log message to be recorded.
+        args : object
+            Additional arguments for the log message.
+        exc_info : logging._ExcInfoType, optional
+            Exception information for the log record.
+        stack_info : bool, optional
+            Whether to include stack information.
+        extra : Mapping[str, object], optional
+            Additional contextual information for the log record.
+
+        Raises
+        ------
+        InvalidBufferItem
+            If the log record cannot be added to the buffer.
+
+        Notes
+        -----
+        Handles special first invocation buffering and migration of log records
+        between different tracer contexts.
+        """
+        # Determine tracer ID, defaulting to first invoke marker
+        tracer_id = get_tracer_id() or LOGGER_BUFFER_FIRST_INVOKE
+
+        try:
+            # Create log record for buffering
             log_record: dict[str, Any] = _create_buffer_record(level=level, msg=msg, args=args, extra=extra)
+
+            # Migrate log records from first invoke to current tracer context
+            if tracer_id != LOGGER_BUFFER_FIRST_INVOKE and self._buffer_cache.get(LOGGER_BUFFER_FIRST_INVOKE):
+                # Retrieve first invoke log records
+                first_invoke_items = self._buffer_cache.get(LOGGER_BUFFER_FIRST_INVOKE)
+
+                # Transfer log records to current tracer context
+                for item in first_invoke_items:
+                    self._buffer_cache.add(tracer_id, item)
+
+                # Clear first invoke buffer
+                self._buffer_cache.clear(LOGGER_BUFFER_FIRST_INVOKE)
+
+            # Add current log record to buffer
             self._buffer_cache.add(tracer_id, log_record)
+        except InvalidBufferItem as exc:
+            # Wrap and re-raise buffer addition error
+            raise InvalidBufferItem("Cannot add item to the buffer") from exc
 
     def flush_buffer(self):
         """
@@ -535,17 +590,27 @@ def flush_buffer(self):
         -----
         Retrieves log records for current trace from buffer
         Immediately processes and logs each record
+        Warning if some cache was evicted in that execution
         Clears buffer after complete processing
 
         Raises
         ------
         Any exceptions from underlying logging or buffer mechanisms
         will be propagated to caller
         """
-        tracer_id = os.getenv(constants.XRAY_TRACE_ID_ENV, None)
+        tracer_id = get_tracer_id()
         for log_line in self._buffer_cache.get(tracer_id):
             self._create_and_flush_log_record(log_line)
 
+        if self._buffer_cache.has_evicted:
+            warnings.warn(
+                message="Some logs are not displayed because they were evicted from the buffer. "
+                "Increase buffer size to store more logs in the buffer",
+                category=PowertoolsUserWarning,
+                stacklevel=2,
+            )
+            self._buffer_cache.has_evicted = False
+
         self._buffer_cache.clear()
 
     def debug(

aws_lambda_powertools/shared/functions.py

@@ -283,3 +283,8 @@ def abs_lambda_path(relative_path: str = "") -> str:
 
 def sanitize_xray_segment_name(name: str) -> str:
     return re.sub(constants.INVALID_XRAY_NAME_CHARACTERS, "", name)
+
+
+def get_tracer_id() -> str | None:
+    xray_trace_id = os.getenv(constants.XRAY_TRACE_ID_ENV)
+    return xray_trace_id.split(";")[0].replace("Root=", "") if xray_trace_id else None

tests/functional/logger/required_dependencies/test_powertools_logger_buffer.py

@@ -4,13 +4,15 @@
 import json
 import random
 import string
+import warnings
 from collections import namedtuple
 
 import pytest
 
 from aws_lambda_powertools import Logger
 from aws_lambda_powertools.logging.buffer import LoggerBufferConfig
 from aws_lambda_powertools.shared import constants
+from aws_lambda_powertools.warnings import PowertoolsUserWarning
 
 
 @pytest.fixture
@@ -158,3 +160,67 @@ def test_create_and_flush_logs(stdout, service_name, monkeypatch):
     # THEN: We expect the log record is not buffered
     log = capture_multiple_logging_statements_output(stdout)
     assert "this log line will be flushed" == log[0]["message"]
+
+
+def test_create_buffer_with_item_overflow(stdout, service_name, monkeypatch):
+    monkeypatch.setenv(constants.XRAY_TRACE_ID_ENV, "1234")
+
+    # GIVEN: A logger configured with 2 bytes
+    logger_buffer_config = LoggerBufferConfig(max_size=2, minimum_log_level="DEBUG")
+
+    logger = Logger(level="DEBUG", service=service_name, stream=stdout, logger_buffer=logger_buffer_config)
+
+    # WHEN logging a line with a size higher than buffer
+    # THEN must raise a warning
+    with pytest.warns(PowertoolsUserWarning, match="Item size*"):
+        logger.debug("this log line will be flushed")
+
+
+def test_create_buffer_with_items_evicted(stdout, service_name, monkeypatch):
+    monkeypatch.setenv(constants.XRAY_TRACE_ID_ENV, "1234")
+
+    # GIVEN: A logger configured with 1024 bytes
+    logger_buffer_config = LoggerBufferConfig(max_size=1024, minimum_log_level="DEBUG")
+
+    logger = Logger(level="DEBUG", service=service_name, stream=stdout, logger_buffer=logger_buffer_config)
+
+    # WHEN we add 3 lines that exceeds than 1024 bytes
+    logger.debug("this log line will be flushed")
+    logger.debug("this log line will be flushed")
+    logger.debug("this log line will be flushed")
+    logger.debug("this log line will be flushed")
+    logger.debug("this log line will be flushed")
+
+    # THEN must raise a warning when trying to flush the lugs
+    with pytest.warns(PowertoolsUserWarning, match="Some logs are not displayed because*"):
+        logger.flush_buffer()
+
+
+def test_create_buffer_with_items_evicted_next_invocation(stdout, service_name, monkeypatch):
+    monkeypatch.setenv(constants.XRAY_TRACE_ID_ENV, "1234")
+
+    # GIVEN: A logger configured with 1024 bytes
+    logger_buffer_config = LoggerBufferConfig(max_size=1024, minimum_log_level="DEBUG")
+
+    logger = Logger(level="DEBUG", service=service_name, stream=stdout, logger_buffer=logger_buffer_config)
+
+    # WHEN Add multiple log entries that exceed buffer size
+    message = "this log line will be flushed"
+    logger.debug(message)
+    logger.debug(message)
+    logger.debug(message)
+    logger.debug(message)
+    logger.debug(message)
+
+    # THEN First buffer flush triggers warning about log eviction
+    with pytest.warns(PowertoolsUserWarning, match="Some logs are not displayed because*"):
+        logger.flush_buffer()
+
+    # WHEN Add another log entry
+    logger.debug("new log entry after buffer flush")
+
+    # THEN Subsequent buffer flush should not trigger warning
+    with warnings.catch_warnings(record=True) as warning_list:
+        warnings.simplefilter("always")
+        logger.flush_buffer()
+        assert len(warning_list) == 0, "No warnings should be raised"