fix hello route hit too early

datadog_lambda/metric.py

@@ -20,6 +20,8 @@
 
 logger = logging.getLogger(__name__)
 
+lambda_stats = None
+
 
 class StatsWriter:
     def distribution(self, metric_name, value, tags=[], timestamp=None):
@@ -110,16 +112,17 @@ def stop(self):
         self.thread_stats.stop()
 
 
-lambda_stats = None
-if should_use_extension:
-    lambda_stats = StatsDWriter()
-else:
-    # Periodical flushing in a background thread is NOT guaranteed to succeed
-    # and leads to data loss. When disabled, metrics are only flushed at the
-    # end of invocation. To make metrics submitted from a long-running Lambda
-    # function available sooner, consider using the Datadog Lambda extension.
-    flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
-    lambda_stats = ThreadStatsWriter(flush_in_thread)
+def init_lambda_stats():
+    global lambda_stats
+    if should_use_extension:
+        lambda_stats = StatsDWriter()
+    else:
+        # Periodical flushing in a background thread is NOT guaranteed to succeed
+        # and leads to data loss. When disabled, metrics are only flushed at the
+        # end of invocation. To make metrics submitted from a long-running Lambda
+        # function available sooner, consider using the Datadog Lambda extension.
+        flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
+        lambda_stats = ThreadStatsWriter(flush_in_thread)
 
 
 def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=False):
@@ -135,6 +138,7 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
     periodically and at the end of the function execution in a
     background thread.
     """
+    global lambda_stats
     flush_to_logs = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
     tags = tag_dd_lambda_layer(tags)
 
@@ -163,6 +167,7 @@ def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]):
 
 
 def flush_stats():
+    global lambda_stats
     lambda_stats.flush()
 
 

datadog_lambda/wrapper.py

@@ -12,6 +12,7 @@
 from datadog_lambda.cold_start import set_cold_start, is_cold_start
 from datadog_lambda.constants import XraySubsegment, TraceContextSource
 from datadog_lambda.metric import (
+    init_lambda_stats,
     flush_stats,
     submit_invocations_metric,
     submit_errors_metric,
@@ -119,6 +120,7 @@ def __call__(self, event, context, **kwargs):
         """Executes when the wrapped function gets called"""
         self.trigger_tags = extract_trigger_tags(event, context)
         self.response = None
+        init_lambda_stats()
         self._before(event, context)
         try:
             self.response = self.func(event, context, **kwargs)

tests/test_wrapper.py

@@ -133,10 +133,11 @@ def lambda_handler(event, context):
 
     def test_datadog_lambda_wrapper_flush_in_thread(self):
         # force ThreadStats to flush in thread
+        os.environ["DD_FLUSH_IN_THREAD"] = "True"
         import datadog_lambda.metric as metric_module
 
         metric_module.lambda_stats.stop()
-        metric_module.lambda_stats = ThreadStatsWriter(True)
+        metric_module.init_lambda_stats()
 
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -157,6 +158,7 @@ def lambda_handler(event, context):
         # reset ThreadStats
         metric_module.lambda_stats.stop()
         metric_module.lambda_stats = ThreadStatsWriter(False)
+        del os.environ["DD_FLUSH_IN_THREAD"]
 
     def test_datadog_lambda_wrapper_not_flush_in_thread(self):
         # force ThreadStats to not flush in thread