Merge pull request #46 from DataDog/tian.chu/fix-dup-enhanced-metrics

Fix duplicate enhanced metrics

Author: tianchu

datadog_lambda/wrapper.py

@@ -38,13 +38,39 @@ def my_lambda_handle(event, context):
 """
 
 
+class _NoopDecorator(object):
+
+    def __init__(self, func):
+        self.func = func
+
+    def __call__(self, *args, **kwargs):
+        return self.func(*args, **kwargs)
+
+
 class _LambdaDecorator(object):
     """
     Decorator to automatically initialize Datadog API client, flush metrics,
     and extracts/injects trace context.
     """
 
+    _instance = None
+    _force_new = False
+
+    def __new__(cls, func):
+        """
+        If the decorator is accidentally applied to multiple functions or
+        the same function multiple times, only the first one takes effect.
+
+        If _force_new, always return a real decorator, useful for unit tests.
+        """
+        if cls._force_new or cls._instance is None:
+            cls._instance = super(_LambdaDecorator, cls).__new__(cls)
+            return cls._instance
+        else:
+            return _NoopDecorator(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 = (
@@ -59,6 +85,17 @@ def __init__(self, func):
         patch_all()
         logger.debug("datadog_lambda_wrapper initialized")
 
+    def __call__(self, event, context, **kwargs):
+        """Executes when the wrapped function gets called"""
+        self._before(event, context)
+        try:
+            return self.func(event, context, **kwargs)
+        except Exception:
+            submit_errors_metric(context)
+            raise
+        finally:
+            self._after(event, context)
+
     def _before(self, event, context):
         set_cold_start()
         try:
@@ -78,15 +115,5 @@ def _after(self, event, context):
         except Exception:
             traceback.print_exc()
 
-    def __call__(self, event, context, **kwargs):
-        self._before(event, context)
-        try:
-            return self.func(event, context, **kwargs)
-        except Exception:
-            submit_errors_metric(context)
-            raise
-        finally:
-            self._after(event, context)
-
 
 datadog_lambda_wrapper = _LambdaDecorator

tests/test_wrapper.py

@@ -24,6 +24,10 @@ def get_mock_context(
 
 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
+
         patcher = patch("datadog_lambda.metric.lambda_stats")
         self.mock_metric_lambda_stats = patcher.start()
         self.addCleanup(patcher.stop)
@@ -251,3 +255,28 @@ def lambda_handler(event, context):
         self.mock_write_metric_point_to_stdout.assert_not_called()
 
         del os.environ["DD_ENHANCED_METRICS"]
+
+    def test_only_one_wrapper_in_use(self):
+        patcher = patch("datadog_lambda.wrapper.submit_invocations_metric")
+        self.mock_submit_invocations_metric = patcher.start()
+        self.addCleanup(patcher.stop)
+
+        @datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            lambda_metric("test.metric", 100)
+
+        # Turn off _force_new 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
+        def lambda_handler_wrapper(event, context):
+            lambda_handler(event, context)
+
+        lambda_event = {}
+
+        lambda_handler_wrapper(lambda_event, get_mock_context())
+
+        self.mock_patch_all.assert_called_once()
+        self.mock_wrapper_lambda_stats.flush.assert_called_once()
+        self.mock_submit_invocations_metric.assert_called_once()