From 30d07c616a19cc07fa2ac8bb0301fea83696555e Mon Sep 17 00:00:00 2001 From: maxday Date: Wed, 16 Jun 2021 13:38:24 -0400 Subject: [PATCH 1/3] fix hello route hit too early --- datadog_lambda/metric.py | 24 ++++++++++++++---------- datadog_lambda/wrapper.py | 2 ++ tests/test_wrapper.py | 5 +++-- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 5eb1c2ac..4683697e 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -20,6 +20,7 @@ logger = logging.getLogger(__name__) +lambda_stats = None class StatsWriter: def distribution(self, metric_name, value, tags=[], timestamp=None): @@ -110,16 +111,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 +137,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 +166,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..a2b36f5c 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -133,10 +133,10 @@ 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 +157,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 From bfdc1d63e7108c0608ad082b4d29302bb7a7a3cd Mon Sep 17 00:00:00 2001 From: maxday Date: Wed, 16 Jun 2021 13:54:55 -0400 Subject: [PATCH 2/3] lint --- datadog_lambda/metric.py | 1 + 1 file changed, 1 insertion(+) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 4683697e..b696fe3c 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -22,6 +22,7 @@ lambda_stats = None + class StatsWriter: def distribution(self, metric_name, value, tags=[], timestamp=None): raise NotImplementedError() From 1c48fa8c681adc414d8ac98bad81fb030f965508 Mon Sep 17 00:00:00 2001 From: maxday Date: Wed, 16 Jun 2021 13:57:24 -0400 Subject: [PATCH 3/3] lint --- tests/test_wrapper.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index a2b36f5c..cfb04ff1 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -135,6 +135,7 @@ 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.init_lambda_stats()