emit telemetry and raise during init

Author: TalUsvyatsky

datadog_lambda/handler.py

@@ -9,7 +9,8 @@
 import os
 from time import time_ns
 
-from datadog_lambda.wrapper import datadog_lambda_wrapper, error_fallback_handler
+from datadog_lambda.tracing import emit_telemetry_on_exception_outside_of_handler
+from datadog_lambda.wrapper import datadog_lambda_wrapper
 from datadog_lambda.module_name import modify_module_name
 
 
@@ -34,8 +35,12 @@ class HandlerError(Exception):
     handler_load_start_time_ns = time_ns()
     handler_module = import_module(modified_mod_name)
     handler_func = getattr(handler_module, handler_name)
-    handler = datadog_lambda_wrapper(handler_func)
 except Exception as e:
-    handler = error_fallback_handler(
-        e, modified_mod_name, time_ns() - handler_load_start_time_ns
+    emit_telemetry_on_exception_outside_of_handler(
+        e,
+        modified_mod_name,
+        handler_load_start_time_ns,
     )
+    raise
+
+handler = datadog_lambda_wrapper(handler_func)

datadog_lambda/tags.py

@@ -55,9 +55,12 @@ def parse_lambda_tags_from_arn(lambda_context):
 
 def get_enhanced_metrics_tags(lambda_context):
     """Get the list of tags to apply to enhanced metrics"""
-    tags = parse_lambda_tags_from_arn(lambda_context)
+    if lambda_context:
+        tags = parse_lambda_tags_from_arn(lambda_context)
+        tags.append(f"memorysize:{lambda_context.memory_limit_in_mb}")
+    else:
+        tags = []
     tags.append(get_cold_start_tag())
-    tags.append(f"memorysize:{lambda_context.memory_limit_in_mb}")
     tags.append(runtime_tag)
     tags.append(library_version_tag)
     return tags

datadog_lambda/tracing.py

@@ -9,7 +9,6 @@
 import traceback
 import ujson as json
 from datetime import datetime, timezone
-from time import time_ns
 from typing import Optional, Dict
 
 from datadog_lambda.metric import submit_errors_metric
@@ -1325,35 +1324,30 @@ def is_async(span: Span) -> bool:
 
 
 def emit_telemetry_on_exception_outside_of_handler(
-    context, exception, resource_name, handler_load_duration
+    exception, resource_name, handler_load_start_time_ns
 ):
     """
-    Emit an enhanced error metric and create a span for exceptions occuring outside of the handler
+    Emit an enhanced error metric and create a span for exceptions occurring outside the handler
     """
-    submit_errors_metric(context)
+    submit_errors_metric(None)
     if dd_tracing_enabled:
         span = tracer.trace(
             "aws.lambda",
             service="aws.lambda",
             resource=resource_name,
             span_type="serverless",
         )
-        span.start_ns = time_ns() - handler_load_duration
+        span.start_ns = handler_load_start_time_ns
 
         tags = {
             "error.status": 500,
             "error.type": type(exception).__name__,
             "error.message": exception,
-            "error.stack": "".join(
-                traceback.format_exception(
-                    type(exception), exception, exception.__traceback__
-                )
-            ),
+            "error.stack": traceback.format_exc(),
             "resource_names": resource_name,
             "resource.name": resource_name,
             "operation_name": "aws.lambda",
             "status": "error",
-            "request_id": context.aws_request_id,
         }
         span.set_tags(tags)
         span.error = 1

datadog_lambda/wrapper.py

@@ -44,7 +44,6 @@
     is_authorizer_response,
     tracer,
     propagator,
-    emit_telemetry_on_exception_outside_of_handler,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -383,31 +382,9 @@ def _after(self, event, context):
             logger.error(format_err_with_traceback(e))
 
 
-class _ErrorFallbackHandler(object):
-    """
-    Fallback handler for when an exception occurs outside of the handler function.
-    Emits telemetry and re-raises the exception.
-    """
-
-    def __init__(self, exception, modified_mod_name, handler_load_duration_ns):
-        self.exception = exception
-        self.modified_mod_name = modified_mod_name
-        self.handler_load_duration_ns = handler_load_duration_ns
-
-    def __call__(self, event, context, **kwargs):
-        emit_telemetry_on_exception_outside_of_handler(
-            context,
-            self.exception,
-            self.modified_mod_name,
-            self.handler_load_duration_ns,
-        )
-        raise self.exception
-
-
 def format_err_with_traceback(e):
     tb = traceback.format_exc().replace("\n", "\r")
     return f"Error {e}. Traceback: {tb}"
 
 
 datadog_lambda_wrapper = _LambdaDecorator
-error_fallback_handler = _ErrorFallbackHandler

tests/test_handler.py

@@ -30,37 +30,33 @@ def test_dd_lambda_handler_env_var_malformed(self):
             )
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
-    @patch("datadog_lambda.wrapper.emit_telemetry_on_exception_outside_of_handler")
+    @patch("datadog_lambda.tracing.emit_telemetry_on_exception_outside_of_handler")
     @patch("time.time_ns", return_value=42)
     def test_exception_importing_module(self, mock_time, mock_emit_telemetry):
         with self.assertRaises(ModuleNotFoundError) as test_context:
             import datadog_lambda.handler
 
-            lambda_context = get_mock_context()
-            datadog_lambda.handler.handler.__call__(None, lambda_context)
         mock_emit_telemetry.assert_called_once_with(
-            lambda_context, test_context.exception, "nonsense", 0
+            test_context.exception, "nonsense", 42
         )
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
     @patch("importlib.import_module", return_value=None)
-    @patch("datadog_lambda.wrapper.emit_telemetry_on_exception_outside_of_handler")
+    @patch("datadog_lambda.tracing.emit_telemetry_on_exception_outside_of_handler")
     @patch("time.time_ns", return_value=42)
     def test_exception_getting_handler_func(
         self, mock_time, mock_emit_telemetry, mock_import
     ):
         with self.assertRaises(AttributeError) as test_context:
             import datadog_lambda.handler
 
-            lambda_context = get_mock_context()
-            datadog_lambda.handler.handler.__call__(None, lambda_context)
         mock_emit_telemetry.assert_called_once_with(
-            lambda_context, test_context.exception, "nonsense", 0
+            test_context.exception, "nonsense", 42
         )
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
     @patch("importlib.import_module")
-    @patch("datadog_lambda.wrapper.emit_telemetry_on_exception_outside_of_handler")
+    @patch("datadog_lambda.tracing.emit_telemetry_on_exception_outside_of_handler")
     @patch("datadog_lambda.wrapper.datadog_lambda_wrapper")
     def test_handler_success(
         self, mock_lambda_wrapper, mock_emit_telemetry, mock_import
@@ -72,8 +68,5 @@ def nonsense():
 
         import datadog_lambda.handler
 
-        lambda_context = get_mock_context()
-        datadog_lambda.handler.handler.__call__(None, lambda_context)
-
         mock_emit_telemetry.assert_not_called()
         mock_lambda_wrapper.assert_called_once_with(mock_import().nonsense)

tests/test_tracing.py

@@ -2006,23 +2006,22 @@ def test_deterministic_m5_hash__always_leading_with_zero(self):
 class TestExceptionOutsideHandler(unittest.TestCase):
     @patch("datadog_lambda.tracing.dd_tracing_enabled", True)
     @patch("datadog_lambda.tracing.submit_errors_metric")
-    @patch("datadog_lambda.tracing.time_ns", return_value=100)
+    @patch("time.time_ns", return_value=42)
     def test_exception_outside_handler_tracing_enabled(
         self, mock_time, mock_submit_errors_metric
     ):
         fake_error = ValueError("Some error message")
         resource_name = "my_handler"
         span_type = "aws.lambda"
         mock_span = Mock()
-        context = get_mock_context()
         with patch(
             "datadog_lambda.tracing.tracer.trace", return_value=mock_span
         ) as mock_trace:
             emit_telemetry_on_exception_outside_of_handler(
-                context, fake_error, resource_name, 42
+                fake_error, resource_name, 42
             )
 
-        mock_submit_errors_metric.assert_called_once_with(context)
+        mock_submit_errors_metric.assert_called_once_with(None)
 
         mock_trace.assert_called_once_with(
             span_type,
@@ -2035,35 +2034,29 @@ def test_exception_outside_handler_tracing_enabled(
                 "error.status": 500,
                 "error.type": "ValueError",
                 "error.message": fake_error,
-                "error.stack": "".join(
-                    traceback.format_exception(
-                        type(fake_error), fake_error, fake_error.__traceback__
-                    )
-                ),
+                "error.stack": traceback.format_exc(),
                 "resource_names": resource_name,
                 "resource.name": resource_name,
                 "operation_name": span_type,
                 "status": "error",
-                "request_id": context.aws_request_id,
             }
         )
         mock_span.finish.assert_called_once()
         assert mock_span.error == 1
-        assert mock_span.start_ns == 58
+        assert mock_span.start_ns == 42
 
     @patch("datadog_lambda.tracing.dd_tracing_enabled", False)
     @patch("datadog_lambda.tracing.submit_errors_metric")
-    @patch("datadog_lambda.tracing.time_ns", return_value=100)
+    @patch("time.time_ns", return_value=42)
     def test_exception_outside_handler_tracing_disabled(
         self, mock_time, mock_submit_errors_metric
     ):
         fake_error = ValueError("Some error message")
         resource_name = "my_handler"
-        context = get_mock_context()
         with patch("datadog_lambda.tracing.tracer.trace") as mock_trace:
             emit_telemetry_on_exception_outside_of_handler(
-                context, fake_error, resource_name, 42
+                fake_error, resource_name, 42
             )
 
-        mock_submit_errors_metric.assert_called_once_with(context)
+        mock_submit_errors_metric.assert_called_once_with(None)
         mock_trace.assert_not_called()

tests/test_wrapper.py

@@ -220,8 +220,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -251,8 +251,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -267,8 +267,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -306,8 +306,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -322,8 +322,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -358,8 +358,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -374,8 +374,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test",
-                        "cold_start:false",
                         "memorysize:256",
+                        "cold_start:false",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -408,8 +408,8 @@ def lambda_handler(event, context):
                         "account_id:123457598159",
                         "functionname:python-layer-test",
                         "resource:python-layer-test:Latest",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",
@@ -442,8 +442,8 @@ def lambda_handler(event, context):
                         "functionname:python-layer-test",
                         "executedversion:1",
                         "resource:python-layer-test:My_alias-1",
-                        "cold_start:true",
                         "memorysize:256",
+                        "cold_start:true",
                         "runtime:python3.9",
                         "datadog_lambda:v6.6.6",
                         "dd_lambda_layer:datadog-python39_X.X.X",