[SVLS-4999] Add Lambda tags to metrics sent via the API (#501)

* add function tags to historical dist metrics * fix * add timestamp ceiling * snake_case * exclude alias when creating arn tag

Author: DylanLovesCoffee

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
+ timestamp_ceiling = int(
+ (datetime.now() - timedelta(hours=4)).timestamp()
+ ) # 4 hours ago
+ if timestamp_ceiling > 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
@@ -108,11 +119,19 @@ 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)
+ 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)
 
 
 def submit_enhanced_metric(metric_name, lambda_context):

datadog_lambda/thread_stats_writer.py

@@ -22,11 +22,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:

datadog_lambda/wrapper.py

@@ -374,7 +374,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

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"
 
@@ -84,6 +100,16 @@ 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]
+ )
+ for tag in tags:
+ self.assertTrue(tag in lambda_stats.thread_stats.constant_tags)
+
 
 MOCK_FUNCTION_NAME = "myFunction"