Datadog wrapper should never break the wrapped Lambda function

Author: tianchu

datadog_lambda/wrapper.py

@@ -53,37 +53,45 @@ class _LambdaDecorator(object):
     and extracts/injects trace context.
     """
 
-    _force_new = False
+    _force_wrap = False
 
     def __new__(cls, func):
         """
         If the decorator is accidentally applied to the same function multiple times,
-        only the first one takes effect.
+        wrap only once.
 
-        If _force_new, always return a real decorator, useful for unit tests.
+        If _force_wrap, always return a real decorator, useful for unit tests.
         """
-        if cls._force_new or not getattr(func, "_dd_wrapped", False):
-            wrapped = super(_LambdaDecorator, cls).__new__(cls)
-            wrapped._dd_wrapped = True
-            return wrapped
-        else:
-            return _NoopDecorator(func)
+        try:
+            if cls._force_wrap or not isinstance(func, _LambdaDecorator):
+                wrapped = super(_LambdaDecorator, cls).__new__(cls)
+                logger.debug("datadog_lambda_wrapper wrapped")
+                return wrapped
+            else:
+                logger.debug("datadog_lambda_wrapper already wrapped")
+                return _NoopDecorator(func)
+        except Exception:
+            traceback.print_exc()
+            return func
 
     def __init__(self, func):
         """Executes when the wrapped function gets wrapped"""
-        self.func = func
-        self.flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
-        self.logs_injection = (
-            os.environ.get("DD_LOGS_INJECTION", "true").lower() == "true"
-        )
-
-        # Inject trace correlation ids to logs
-        if self.logs_injection:
-            inject_correlation_ids()
-
-        # Patch HTTP clients to propagate Datadog trace context
-        patch_all()
-        logger.debug("datadog_lambda_wrapper initialized")
+        try:
+            self.func = func
+            self.flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
+            self.logs_injection = (
+                os.environ.get("DD_LOGS_INJECTION", "true").lower() == "true"
+            )
+
+            # Inject trace correlation ids to logs
+            if self.logs_injection:
+                inject_correlation_ids()
+
+            # Patch HTTP clients to propagate Datadog trace context
+            patch_all()
+            logger.debug("datadog_lambda_wrapper initialized")
+        except Exception:
+            traceback.print_exc()
 
     def __call__(self, event, context, **kwargs):
         """Executes when the wrapped function gets called"""
@@ -97,21 +105,23 @@ def __call__(self, event, context, **kwargs):
             self._after(event, context)
 
     def _before(self, event, context):
-        set_cold_start()
         try:
+            set_cold_start()
             submit_invocations_metric(context)
             # Extract Datadog trace context from incoming requests
             extract_dd_trace_context(event)
 
             # Set log correlation ids using extracted trace context
             set_correlation_ids()
+            logger.debug("datadog_lambda_wrapper _before() done")
         except Exception:
             traceback.print_exc()
 
     def _after(self, event, context):
         try:
             if not self.flush_to_log:
                 lambda_stats.flush(float("inf"))
+            logger.debug("datadog_lambda_wrapper _after() done")
         except Exception:
             traceback.print_exc()
 

tests/test_wrapper.py

@@ -26,7 +26,7 @@ class TestDatadogLambdaWrapper(unittest.TestCase):
     def setUp(self):
         # Force @datadog_lambda_wrapper to always create a real
         # (not no-op) wrapper.
-        datadog_lambda_wrapper._force_new = True
+        datadog_lambda_wrapper._force_wrap = True
 
         patcher = patch("datadog_lambda.metric.lambda_stats")
         self.mock_metric_lambda_stats = patcher.start()
@@ -265,9 +265,9 @@ def test_only_one_wrapper_in_use(self):
         def lambda_handler(event, context):
             lambda_metric("test.metric", 100)
 
-        # Turn off _force_new to emulate the nested wrapper scenario,
+        # Turn off _force_wrap to emulate the nested wrapper scenario,
         # the second @datadog_lambda_wrapper should actually be no-op.
-        datadog_lambda_wrapper._force_new = False
+        datadog_lambda_wrapper._force_wrap = False
 
         lambda_handler_double_wrapped = datadog_lambda_wrapper(lambda_handler)