Detect exceptions raised outside of the handler (#475)

* emit enhanced error metric and create span when an exception is raised outside of the handler function

* rename handler

* move fallback handler to wrapper.py

* fix duration

* respect DD_TRACE_ENABLED

* emit telemetry and raise during init

* Update datadog_lambda/tags.py

Co-authored-by: jordan gonzález <30836115+duncanista@users.noreply.github.com>

---------

Co-authored-by: jordan gonzález <30836115+duncanista@users.noreply.github.com>

datadog_lambda/handler.py

@@ -7,6 +7,9 @@
 from importlib import import_module
 
 import os
+from time import time_ns
+
+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
 
@@ -27,5 +30,17 @@ class HandlerError(Exception):
 
 (mod_name, handler_name) = parts
 modified_mod_name = modify_module_name(mod_name)
-handler_module = import_module(modified_mod_name)
-handler = datadog_lambda_wrapper(getattr(handler_module, handler_name))
+
+try:
+    handler_load_start_time_ns = time_ns()
+    handler_module = import_module(modified_mod_name)
+    handler_func = getattr(handler_module, handler_name)
+except Exception as e:
+    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/metric.py

@@ -100,7 +100,7 @@ def submit_enhanced_metric(metric_name, lambda_context):
 
     Args:
         metric_name (str): metric name w/o enhanced prefix i.e. "invocations" or "errors"
-        lambda_context (dict): Lambda context dict passed to the function by AWS
+        lambda_context (object): Lambda context dict passed to the function by AWS
     """
     if not enhanced_metrics_enabled:
         logger.debug(
@@ -118,7 +118,7 @@ def submit_invocations_metric(lambda_context):
     """Increment aws.lambda.enhanced.invocations by 1, applying runtime, layer, and cold_start tags
 
     Args:
-        lambda_context (dict): Lambda context dict passed to the function by AWS
+        lambda_context (object): Lambda context dict passed to the function by AWS
     """
     submit_enhanced_metric("invocations", lambda_context)
 
@@ -127,6 +127,6 @@ def submit_errors_metric(lambda_context):
     """Increment aws.lambda.enhanced.errors by 1, applying runtime, layer, and cold_start tags
 
     Args:
-        lambda_context (dict): Lambda context dict passed to the function by AWS
+        lambda_context (object): Lambda context dict passed to the function by AWS
     """
     submit_enhanced_metric("errors", lambda_context)

datadog_lambda/tags.py

@@ -55,9 +55,11 @@ 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)
+    tags = []
+    if lambda_context:
+        tags = parse_lambda_tags_from_arn(lambda_context)
+        tags.append(f"memorysize:{lambda_context.memory_limit_in_mb}")
     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

@@ -6,6 +6,7 @@
 import logging
 import os
 import base64
+import traceback
 import ujson as json
 from datetime import datetime, timezone
 from typing import Optional, Dict
@@ -1320,3 +1321,34 @@ def is_async(span: Span) -> bool:
                 e,
             )
             return False
+
+
+def emit_telemetry_on_exception_outside_of_handler(
+    exception, resource_name, handler_load_start_time_ns
+):
+    """
+    Emit an enhanced error metric and create a span for exceptions occurring outside the handler
+    """
+    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 = handler_load_start_time_ns
+
+        tags = {
+            "error.status": 500,
+            "error.type": type(exception).__name__,
+            "error.message": exception,
+            "error.stack": traceback.format_exc(),
+            "resource_names": resource_name,
+            "resource.name": resource_name,
+            "operation_name": "aws.lambda",
+            "status": "error",
+        }
+        span.set_tags(tags)
+        span.error = 1
+        span.finish()

tests/test_handler.py

@@ -0,0 +1,72 @@
+import os
+import sys
+import unittest
+from unittest.mock import patch
+
+from tests.utils import get_mock_context
+
+
+class TestHandler(unittest.TestCase):
+    def tearDown(self):
+        for mod in sys.modules.copy():
+            if mod.startswith("datadog_lambda.handler"):
+                del sys.modules[mod]
+
+    def test_dd_lambda_handler_env_var_none(self):
+        with self.assertRaises(Exception) as context:
+            import datadog_lambda.handler as handler
+
+            assert context.exception == handler.HandlerError(
+                "DD_LAMBDA_HANDLER is not defined. Can't use prebuilt datadog handler"
+            )
+
+    @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "malformed"}, clear=True)
+    def test_dd_lambda_handler_env_var_malformed(self):
+        with self.assertRaises(Exception) as context:
+            import datadog_lambda.handler as handler
+
+            assert context.exception == handler.HandlerError(
+                "Value malformed for DD_LAMBDA_HANDLER has invalid format."
+            )
+
+    @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
+    @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
+
+        mock_emit_telemetry.assert_called_once_with(
+            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.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
+
+        mock_emit_telemetry.assert_called_once_with(
+            test_context.exception, "nonsense", 42
+        )
+
+    @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
+    @patch("importlib.import_module")
+    @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
+    ):
+        def nonsense():
+            pass
+
+        mock_import.nonsense.return_value = nonsense
+
+        import datadog_lambda.handler
+
+        mock_emit_telemetry.assert_not_called()
+        mock_lambda_wrapper.assert_called_once_with(mock_import().nonsense)

tests/test_tracing.py

@@ -1,6 +1,7 @@
 import copy
 import functools
 import json
+import traceback
 import pytest
 import os
 import unittest
@@ -36,6 +37,7 @@
     determine_service_name,
     service_mapping as global_service_mapping,
     propagator,
+    emit_telemetry_on_exception_outside_of_handler,
 )
 from datadog_lambda.trigger import EventTypes
 
@@ -1999,3 +2001,62 @@ def test_deterministic_m5_hash__always_leading_with_zero(self):
             # Leading zeros will be omitted, so only test for full 64 bits present
             if len(result_in_binary) == 66:  # "0b" + 64 bits.
                 self.assertTrue(result_in_binary.startswith("0b0"))
+
+
+class TestExceptionOutsideHandler(unittest.TestCase):
+    @patch("datadog_lambda.tracing.dd_tracing_enabled", True)
+    @patch("datadog_lambda.tracing.submit_errors_metric")
+    @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()
+        with patch(
+            "datadog_lambda.tracing.tracer.trace", return_value=mock_span
+        ) as mock_trace:
+            emit_telemetry_on_exception_outside_of_handler(
+                fake_error, resource_name, 42
+            )
+
+        mock_submit_errors_metric.assert_called_once_with(None)
+
+        mock_trace.assert_called_once_with(
+            span_type,
+            service="aws.lambda",
+            resource=resource_name,
+            span_type="serverless",
+        )
+        mock_span.set_tags.assert_called_once_with(
+            {
+                "error.status": 500,
+                "error.type": "ValueError",
+                "error.message": fake_error,
+                "error.stack": traceback.format_exc(),
+                "resource_names": resource_name,
+                "resource.name": resource_name,
+                "operation_name": span_type,
+                "status": "error",
+            }
+        )
+        mock_span.finish.assert_called_once()
+        assert mock_span.error == 1
+        assert mock_span.start_ns == 42
+
+    @patch("datadog_lambda.tracing.dd_tracing_enabled", False)
+    @patch("datadog_lambda.tracing.submit_errors_metric")
+    @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"
+        with patch("datadog_lambda.tracing.tracer.trace") as mock_trace:
+            emit_telemetry_on_exception_outside_of_handler(
+                fake_error, resource_name, 42
+            )
+
+        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",