From 8c5b414225944bed8a8207977fa89fedf9956d57 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Tue, 2 Jul 2024 15:32:44 -0400 Subject: [PATCH 1/5] add function tags to historical dist metrics --- datadog_lambda/metric.py | 7 +++++-- datadog_lambda/thread_stats_writer.py | 5 ++++- datadog_lambda/wrapper.py | 2 +- tests/test_metric.py | 9 +++++++++ 4 files changed, 19 insertions(+), 4 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 115686af..2aeaac15 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -108,11 +108,14 @@ def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]): ) -def flush_stats(): +def flush_stats(lambda_context=None): lambda_stats.flush() if extension_thread_stats is not None: - extension_thread_stats.flush() + if lambda_context is not None: + tags = get_enhanced_metrics_tags(lambda_context) + tags.append("function_arn:" + lambda_context.invoked_function_arn) + extension_thread_stats.flush(tags) def submit_enhanced_metric(metric_name, lambda_context): diff --git a/datadog_lambda/thread_stats_writer.py b/datadog_lambda/thread_stats_writer.py index bfcf3c99..b631f356 100644 --- a/datadog_lambda/thread_stats_writer.py +++ b/datadog_lambda/thread_stats_writer.py @@ -1,4 +1,5 @@ import logging +import os # Make sure that this package would always be lazy-loaded/outside from the critical path # since underlying packages are quite heavy to load and useless when the extension is present @@ -22,11 +23,13 @@ def distribution(self, metric_name, value, tags=[], timestamp=None): metric_name, value, tags=tags, timestamp=timestamp ) - def flush(self): + def flush(self, tags=None): """ "Flush distributions from ThreadStats to Datadog. Modified based on `datadog.threadstats.base.ThreadStats.flush()`, to gain better control over exception handling. """ + if tags: + self.thread_stats.constant_tags = self.thread_stats.constant_tags + tags _, dists = self.thread_stats._get_aggregate_metrics_and_dists(float("inf")) count_dists = len(dists) if not count_dists: diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index ba31f2be..9025bfef 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -366,7 +366,7 @@ def _after(self, event, context): logger.debug("Failed to create cold start spans. %s", e) if not self.flush_to_log or should_use_extension: - flush_stats() + flush_stats(context) if should_use_extension and self.local_testing_mode: # when testing locally, the extension does not know when an # invocation completes because it does not have access to the diff --git a/tests/test_metric.py b/tests/test_metric.py index f07a4c6a..957378c1 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -84,6 +84,15 @@ def test_retry_on_remote_disconnected(self): lambda_stats.flush() self.assertEqual(self.mock_threadstats_flush_distributions.call_count, 2) + def test_flush_stats_with_tags(self): + lambda_stats = ThreadStatsWriter(True) + tags = ["tag1:value1", "tag2:value2"] + lambda_stats.flush(tags) + self.mock_threadstats_flush_distributions.assert_called_once_with( + lambda_stats.thread_stats._get_aggregate_metrics_and_dists(float("inf"))[1] + ) + self.assertEqual(lambda_stats.thread_stats.constant_tags, tags) + MOCK_FUNCTION_NAME = "myFunction" From 1b07836f314936e003116cac07021faf93f91344 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Tue, 2 Jul 2024 15:54:28 -0400 Subject: [PATCH 2/5] fix --- datadog_lambda/thread_stats_writer.py | 1 - tests/test_metric.py | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/datadog_lambda/thread_stats_writer.py b/datadog_lambda/thread_stats_writer.py index b631f356..367b8b21 100644 --- a/datadog_lambda/thread_stats_writer.py +++ b/datadog_lambda/thread_stats_writer.py @@ -1,5 +1,4 @@ import logging -import os # Make sure that this package would always be lazy-loaded/outside from the critical path # since underlying packages are quite heavy to load and useless when the extension is present diff --git a/tests/test_metric.py b/tests/test_metric.py index 957378c1..fcc5e8e6 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -91,7 +91,8 @@ def test_flush_stats_with_tags(self): self.mock_threadstats_flush_distributions.assert_called_once_with( lambda_stats.thread_stats._get_aggregate_metrics_and_dists(float("inf"))[1] ) - self.assertEqual(lambda_stats.thread_stats.constant_tags, tags) + for tag in tags: + self.assertTrue(tag in lambda_stats.thread_stats.constant_tags) MOCK_FUNCTION_NAME = "myFunction" From f40f0725151924b66f4f3b2afa9db397bd676b63 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Tue, 2 Jul 2024 16:07:41 -0400 Subject: [PATCH 3/5] add timestamp ceiling --- datadog_lambda/metric.py | 11 +++++++++++ tests/test_metric.py | 22 +++++++++++++++++++--- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 2aeaac15..9e988eb7 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -7,6 +7,7 @@ import time import logging import ujson as json +from datetime import datetime, timedelta from datadog_lambda.extension import should_use_extension from datadog_lambda.tags import get_enhanced_metrics_tags, dd_lambda_layer_tag @@ -61,6 +62,16 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal if should_use_extension and timestamp is not None: # The extension does not support timestamps for distributions so we create a # a thread stats writer to submit metrics with timestamps to the API + timestampCeiling = int( + (datetime.now() - timedelta(hours=4)).timestamp() + ) # 4 hours ago + if timestampCeiling > timestamp: + logger.warning( + "Timestamp %s is older than 4 hours, not submitting metric %s", + timestamp, + metric_name, + ) + return global extension_thread_stats if extension_thread_stats is None: from datadog_lambda.thread_stats_writer import ThreadStatsWriter diff --git a/tests/test_metric.py b/tests/test_metric.py index fcc5e8e6..031b1180 100644 --- a/tests/test_metric.py +++ b/tests/test_metric.py @@ -5,7 +5,7 @@ from botocore.exceptions import ClientError as BotocoreClientError from datadog.api.exceptions import ClientError - +from datetime import datetime, timedelta from datadog_lambda.metric import lambda_metric from datadog_lambda.api import decrypt_kms_api_key, KMS_ENCRYPTION_CONTEXT_KEY @@ -49,12 +49,28 @@ def test_lambda_metric_timestamp_with_extension(self): self.mock_metric_extension_thread_stats = patcher.start() self.addCleanup(patcher.stop) - lambda_metric("test_timestamp", 1, 123) + delta = timedelta(minutes=1) + timestamp = int((datetime.now() - delta).timestamp()) + + lambda_metric("test_timestamp", 1, timestamp) self.mock_metric_lambda_stats.distribution.assert_not_called() self.mock_metric_extension_thread_stats.distribution.assert_called_with( - "test_timestamp", 1, timestamp=123, tags=[dd_lambda_layer_tag] + "test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag] ) + @patch("datadog_lambda.metric.should_use_extension", True) + def test_lambda_metric_invalid_timestamp_with_extension(self): + patcher = patch("datadog_lambda.metric.extension_thread_stats") + self.mock_metric_extension_thread_stats = patcher.start() + self.addCleanup(patcher.stop) + + delta = timedelta(hours=5) + timestamp = int((datetime.now() - delta).timestamp()) + + lambda_metric("test_timestamp", 1, timestamp) + self.mock_metric_lambda_stats.distribution.assert_not_called() + self.mock_metric_extension_thread_stats.distribution.assert_not_called() + def test_lambda_metric_flush_to_log(self): os.environ["DD_FLUSH_TO_LOG"] = "True" From f8a30ebb1206b1512129f0fa15b00efe1eb371f7 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Wed, 3 Jul 2024 14:23:38 -0400 Subject: [PATCH 4/5] snake_case --- datadog_lambda/metric.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index 9e988eb7..d598cb10 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -62,10 +62,10 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal if should_use_extension and timestamp is not None: # The extension does not support timestamps for distributions so we create a # a thread stats writer to submit metrics with timestamps to the API - timestampCeiling = int( + timestamp_ceiling = int( (datetime.now() - timedelta(hours=4)).timestamp() ) # 4 hours ago - if timestampCeiling > timestamp: + if timestamp_ceiling > timestamp: logger.warning( "Timestamp %s is older than 4 hours, not submitting metric %s", timestamp, From ac4d4f156bc95e42ae5265e1f2af54bda8d61a69 Mon Sep 17 00:00:00 2001 From: Dylan Yang Date: Mon, 8 Jul 2024 16:30:51 -0400 Subject: [PATCH 5/5] exclude alias when creating arn tag --- datadog_lambda/metric.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/datadog_lambda/metric.py b/datadog_lambda/metric.py index d598cb10..3bc9955f 100644 --- a/datadog_lambda/metric.py +++ b/datadog_lambda/metric.py @@ -125,7 +125,12 @@ def flush_stats(lambda_context=None): if extension_thread_stats is not None: if lambda_context is not None: tags = get_enhanced_metrics_tags(lambda_context) - tags.append("function_arn:" + lambda_context.invoked_function_arn) + split_arn = lambda_context.invoked_function_arn.split(":") + if len(split_arn) > 7: + # Get rid of the alias + split_arn.pop() + arn = ":".join(split_arn) + tags.append("function_arn:" + arn) extension_thread_stats.flush(tags)