bugfixes and speedups

Author: Artem Krivonos

awslambdaric/bootstrap.py

@@ -14,16 +14,21 @@
 from .lambda_runtime_client import LambdaRuntimeClient
 from .lambda_runtime_exception import FaultException
 from .lambda_runtime_log_utils import (
-    INVOCATION_LOGGING_CONTEXT,
     JsonFormatter,
     DATETIME_FORMAT,
-    LambdaLogFormat,
+    JSON_FORMAT,
+    TEXT_FORMAT,
+    get_log_format_from_str,
 )
 from .lambda_runtime_marshaller import to_json
 
 ERROR_LOG_LINE_TERMINATE = "\r"
 ERROR_LOG_IDENT = "\u00a0"  # NO-BREAK SPACE U+00A0
 RUNTIME_ERROR_LOGGER_NAME = "system"
+AWS_LAMBDA_LOG_FORMAT = get_log_format_from_str(
+    os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT").upper()
+)
+AWS_LAMBDA_LOG_LEVEL = os.environ.get("AWS_LAMBDA_LOG_LEVEL")
 
 
 def _get_handler(handler):
@@ -106,18 +111,11 @@ def replace_line_indentation(line, indent_char, new_indent_char):
     return (new_indent_char * ident_chars_count) + line[ident_chars_count:]
 
 
-def log_error(error_result, log_sink):
-    if AWS_LAMBDA_LOG_FORMAT == LambdaLogFormat.JSON:
-        log_error_json(error_result, log_sink)
-    else:
-        log_error_text(error_result, log_sink)
-
-
 def log_error_json(error_result, log_sink):
     log_level = error_result.pop("log_level", logging.ERROR)
     error_result["level"] = logging.getLevelName(log_level)
     log_sink.log_error(
-        [to_json(error_result)],
+        [to_json(error_result, ensure_ascii=False)],
         log_level=log_level,
     )
 
@@ -154,6 +152,12 @@ def log_error_text(error_result, log_sink):
     )
 
 
+if AWS_LAMBDA_LOG_FORMAT == JSON_FORMAT:
+    log_error = log_error_json
+else:
+    log_error = log_error_text
+
+
 def handle_event_request(
     lambda_runtime_client,
     request_handler,
@@ -292,13 +296,13 @@ def __init__(self, log_sink):
         logging.Handler.__init__(self)
         self.log_sink = log_sink
 
-    def emit(self, record: logging.LogRecord):
+    def emit(self, record):
         msg = self.format(record)
 
         self.log_sink.log(
             msg,
             log_level=record.levelno,
-            log_format=getattr(record, "log_format", LambdaLogFormat.TEXT),
+            log_format=getattr(record, "log_format", TEXT_FORMAT),
         )
 
 
@@ -340,14 +344,31 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_tb):
         pass
 
-    def log(self, msg, log_level=logging.NOTSET, log_format=LambdaLogFormat.TEXT):
+    def log(self, msg, log_level=None, log_format=None):
         sys.stdout.write(msg)
 
     def log_error(self, message_lines, log_level=logging.ERROR):
         error_message = ERROR_LOG_LINE_TERMINATE.join(message_lines) + "\n"
         sys.stdout.write(error_message)
 
 
+FRAME_TYPES = {
+    (JSON_FORMAT, logging.NOTSET): 0xA55A0002.to_bytes(4, "big"),
+    (JSON_FORMAT, logging.DEBUG): 0xA55A000A.to_bytes(4, "big"),
+    (JSON_FORMAT, logging.INFO): 0xA55A000E.to_bytes(4, "big"),
+    (JSON_FORMAT, logging.WARNING): 0xA55A0012.to_bytes(4, "big"),
+    (JSON_FORMAT, logging.ERROR): 0xA55A0016.to_bytes(4, "big"),
+    (JSON_FORMAT, logging.CRITICAL): 0xA55A001A.to_bytes(4, "big"),
+    (TEXT_FORMAT, logging.NOTSET): 0xA55A0003.to_bytes(4, "big"),
+    (TEXT_FORMAT, logging.DEBUG): 0xA55A000B.to_bytes(4, "big"),
+    (TEXT_FORMAT, logging.INFO): 0xA55A000F.to_bytes(4, "big"),
+    (TEXT_FORMAT, logging.WARNING): 0xA55A0013.to_bytes(4, "big"),
+    (TEXT_FORMAT, logging.ERROR): 0xA55A0017.to_bytes(4, "big"),
+    (TEXT_FORMAT, logging.CRITICAL): 0xA55A001B.to_bytes(4, "big"),
+}
+DEFAULT_FRAME_TYPE = 0xA55A0003.to_bytes(4, "big")
+
+
 class FramedTelemetryLogSink(object):
     """
     FramedTelemetryLogSink implements the logging contract between runtimes and the platform. It implements a simple
@@ -364,19 +385,8 @@ class FramedTelemetryLogSink(object):
     The next 'len' bytes contain the message. The byte order is big-endian.
     """
 
-    LEVEL_TO_MASK = {
-        logging.NOTSET: 0b00000,
-        logging.DEBUG: 0b01000,
-        logging.INFO: 0b01100,
-        logging.WARNING: 0b10000,
-        logging.ERROR: 0b10100,
-        logging.FATAL: 0b11000,
-    }
-    DEFAULT_LEVEL_MASK = 0b00000
-
     def __init__(self, fd):
         self.fd = int(fd)
-        self.frame_type = 0xA55A0002
 
     def __enter__(self):
         self.file = os.fdopen(self.fd, "wb", 0)
@@ -385,16 +395,13 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, exc_tb):
         self.file.close()
 
-    def log(self, msg, log_level=logging.NOTSET, log_format=LambdaLogFormat.TEXT):
-        frame_type = self.frame_type | self.LEVEL_TO_MASK.get(
-            log_level, self.DEFAULT_LEVEL_MASK
-        )
-        frame_type = frame_type | log_format
-
+    def log(self, msg, log_level=logging.NOTSET, log_format: int = TEXT_FORMAT):
+        frame_type = FRAME_TYPES.get((log_format, log_level), DEFAULT_FRAME_TYPE)
         encoded_msg = msg.encode("utf8")
+
         timestamp = int(time.time_ns() / 1000)  # UNIX timestamp in microseconds
         log_msg = (
-            frame_type.to_bytes(4, "big")
+            frame_type
             + len(encoded_msg).to_bytes(4, "big")
             + timestamp.to_bytes(8, "big")
             + encoded_msg
@@ -424,10 +431,6 @@ def create_log_sink():
         return StandardLogSink()
 
 
-AWS_LAMBDA_LOG_FORMAT = LambdaLogFormat.from_str(
-    os.environ.get("AWS_LAMBDA_LOG_FORMAT", "TEXT")
-)
-AWS_LAMBDA_LOG_LEVEL = os.environ.get("AWS_LAMBDA_LOG_LEVEL")
 _GLOBAL_AWS_REQUEST_ID = None
 
 
@@ -436,7 +439,7 @@ def setup_logging(log_format, log_level, log_sink):
     logging.Formatter.converter = time.gmtime
     logger = logging.getLogger()
     logger_handler = LambdaLoggerHandler(log_sink)
-    if log_format == LambdaLogFormat.JSON:
+    if log_format == JSON_FORMAT:
         logger_handler.setFormatter(JsonFormatter())
     else:
         logger_handler.setFormatter(
@@ -475,7 +478,6 @@ def run(app_root, handler, lambda_runtime_api_addr):
 
         while True:
             event_request = lambda_runtime_client.wait_next_invocation()
-            INVOCATION_LOGGING_CONTEXT.clear()
 
             _GLOBAL_AWS_REQUEST_ID = event_request.invoke_id
 

awslambdaric/lambda_runtime_log_utils.py

@@ -4,7 +4,6 @@
 
 import logging
 import traceback
-from enum import IntFlag
 
 from .lambda_runtime_marshaller import to_json
 
@@ -33,17 +32,14 @@
     "aws_request_id",
     "log_format",
 }
+JSON_FORMAT = 0b0
+TEXT_FORMAT = 0b1
 
 
-class LambdaLogFormat(IntFlag):
-    JSON = 0b0
-    TEXT = 0b1
-
-    @classmethod
-    def from_str(cls, value):
-        if value == "JSON":
-            return cls.JSON
-        return cls.TEXT
+def get_log_format_from_str(value: str):
+    if value == "JSON":
+        return JSON_FORMAT
+    return TEXT_FORMAT
 
 
 class JsonFormatter(logging.Formatter):
@@ -79,12 +75,12 @@ def format_location(record: logging.LogRecord):
 
     @staticmethod
     def format_log_level(record: logging.LogRecord):
-        level = record.levelno % 51 // 10 * 10
+        level = min(50, max(0, record.levelno)) // 10 * 10
         record.levelname = logging.getLevelName(level)
 
     def format(self, record: logging.LogRecord) -> str:
         self.format_log_level(record)
-        record.log_format = LambdaLogFormat.JSON
+        record.log_format = JSON_FORMAT
 
         result = {
             "timestamp": self.formatTime(record, self.datefmt),
@@ -102,21 +98,7 @@ def format(self, record: logging.LogRecord) -> str:
             for key, value in record.__dict__.items()
             if key not in RESERVED_FIELDS and key not in result
         )
-        result.update(INSTANCE_LOGGING_CONTEXT)
-        result.update(INVOCATION_LOGGING_CONTEXT)
 
         result = {k: v for k, v in result.items() if v is not None}
 
-        return to_json(result) + "\n"
-
-
-INSTANCE_LOGGING_CONTEXT = dict()
-INVOCATION_LOGGING_CONTEXT = dict()
-
-
-def lambda_instance_logs_update(key: str, value):
-    INSTANCE_LOGGING_CONTEXT[str(key)] = value
-
-
-def lambda_invocation_logs_update(key: str, value):
-    INVOCATION_LOGGING_CONTEXT[str(key)] = value
+        return to_json(result, ensure_ascii=False) + "\n"

awslambdaric/lambda_runtime_marshaller.py

@@ -13,8 +13,8 @@
 # simplejson's Decimal encoding allows '-NaN' as an output, which is a parse error for json.loads
 # to get the good parts of Decimal support, we'll special-case NaN decimals and otherwise duplicate the encoding for decimals the same way simplejson does
 class Encoder(json.JSONEncoder):
-    def __init__(self):
-        super().__init__(use_decimal=False)
+    def __init__(self, ensure_ascii=True):
+        super().__init__(use_decimal=False, ensure_ascii=ensure_ascii)
 
     def default(self, obj):
         if isinstance(obj, decimal.Decimal):
@@ -24,8 +24,8 @@ def default(self, obj):
         return super().default(obj)
 
 
-def to_json(obj):
-    return Encoder().encode(obj)
+def to_json(obj, ensure_ascii=True):
+    return Encoder(ensure_ascii=ensure_ascii).encode(obj)
 
 
 class LambdaMarshaller:

tests/test_bootstrap.py

@@ -20,11 +20,9 @@
 from awslambdaric.lambda_runtime_exception import FaultException
 from awslambdaric.lambda_runtime_marshaller import LambdaMarshaller
 from awslambdaric.lambda_runtime_log_utils import (
-    lambda_instance_logs_update,
-    lambda_invocation_logs_update,
-    INVOCATION_LOGGING_CONTEXT,
-    INSTANCE_LOGGING_CONTEXT,
-    LambdaLogFormat,
+    JSON_FORMAT,
+    TEXT_FORMAT,
+    get_log_format_from_str,
 )
 
 
@@ -1144,12 +1142,12 @@ def test_log_level_frame_type(self):
                     content = f.read()
 
                     frame_type = int.from_bytes(content[:4], "big")
-                    self.assertEqual(frame_type, expected_frame_type, hex(frame_type))
+                    self.assertEqual(bin(frame_type), bin(expected_frame_type))
 
     def test_log_format_frame_type(self):
         test_cases = [
-            (LambdaLogFormat.TEXT, 0xA55A0003),
-            (LambdaLogFormat.JSON, 0xA55A0002),
+            (TEXT_FORMAT, 0xA55A0003),
+            (JSON_FORMAT, 0xA55A0002),
         ]
 
         for fmt, expected_frame_type in test_cases:
@@ -1162,7 +1160,7 @@ def test_log_format_frame_type(self):
                     content = f.read()
 
                     frame_type = int.from_bytes(content[:4], "big")
-                    self.assertEqual(frame_type, expected_frame_type, hex(frame_type))
+                    self.assertEqual(hex(frame_type), hex(expected_frame_type))
 
     def test_single_frame(self):
         with NamedTemporaryFile() as temp_file:
@@ -1230,13 +1228,9 @@ class TestLogging(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
         bootstrap.setup_logging(
-            LambdaLogFormat.from_str("JSON"), "INFO", bootstrap.StandardLogSink()
+            get_log_format_from_str("JSON"), "INFO", bootstrap.StandardLogSink()
         )
 
-    def tearDown(self) -> None:
-        INVOCATION_LOGGING_CONTEXT.clear()
-        INSTANCE_LOGGING_CONTEXT.clear()
-
     @patch("sys.stderr", new_callable=StringIO)
     def test_json_formatter(self, mock_stderr):
         logger = logging.getLogger("a.b")
@@ -1317,36 +1311,6 @@ def test_exception(self, mock_stderr, mock_stdout):
 
         self.assertEqual(mock_stderr.getvalue(), "")
 
-    @patch("sys.stdout", new_callable=StringIO)
-    @patch("sys.stderr", new_callable=StringIO)
-    def test_log_with_extra_params(self, mock_stderr, mock_stdout):
-        lambda_instance_logs_update("instance_key", "instance_value")
-        lambda_invocation_logs_update("invocation_key", "invocation_value")
-        lambda_invocation_logs_update("int_param", 42)
-        lambda_invocation_logs_update("list_param", ["1", 2, {}])
-        lambda_invocation_logs_update("dict_param", {"a": "b"})
-
-        logging.getLogger("test.logger").error("test extra params")
-
-        data = json.loads(mock_stdout.getvalue())
-        data.pop("timestamp")
-
-        self.assertEqual(
-            data,
-            {
-                "instance_key": "instance_value",
-                "invocation_key": "invocation_value",
-                "int_param": 42,
-                "list_param": ["1", 2, {}],
-                "dict_param": {"a": "b"},
-                "level": "ERROR",
-                "logger": "test.logger",
-                "message": "test extra params",
-                "requestId": "",
-            },
-        )
-        self.assertEqual(mock_stderr.getvalue(), "")
-
     @patch("sys.stdout", new_callable=StringIO)
     @patch("sys.stderr", new_callable=StringIO)
     def test_log_level(self, mock_stderr, mock_stdout):