diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 5eb1c2ac..b696fe3c 100644 --- a/datadog_lambda/metric.py +++ b/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() diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 352b8f49..702c8259 100644 --- a/datadog_lambda/wrapper.py +++ b/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) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 77d1bde3..cfb04ff1 100644 --- a/tests/test_wrapper.py +++ b/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