Submit metrics for invocations and errors

Author: swfunc

datadog_lambda/tags.py

@@ -0,0 +1,40 @@
+def parse_lambda_tags_from_arn(arn):
+    """Generate the list of lambda tags based on the data in the arn
+    Args:
+        arn (str): Lambda ARN.
+            ex: arn:aws:lambda:us-east-1:123597598159:function:my-lambda[:optional-version]
+    """
+    # Cap the number of times to split
+    split_arn = arn.split(":")
+
+    # If ARN includes version / alias at the end, drop it
+    if len(split_arn) > 7:
+        split_arn = split_arn[:7]
+
+    _, _, _, region, account_id, _, function_name = split_arn
+
+    return [
+        "region:{}".format(region),
+        "account_id:{}".format(account_id),
+        "functionname:{}".format(function_name),
+    ]
+
+
+def get_tags_from_context(context, cold_start_request_id):
+    """Uses properties of the Lambda context to create the list of tags
+
+    Args:
+        context (dict<str, multiple types>): context this Lambda was invoked with
+        cold_start_request_id (str): the first request ID to execute in this container
+
+    Returns:
+        tag list (str[]): list of string tags in key:value format
+    """
+    tags = parse_lambda_tags_from_arn(context.invoked_function_arn)
+    tags.append("memorysize:{}".format(context.memory_limit_in_mb))
+
+    did_request_cold_start = cold_start_request_id == context.aws_request_id
+    cold_start_tag = "cold_start:{}".format(str(did_request_cold_start).lower())
+    tags.append(cold_start_tag)
+
+    return tags

datadog_lambda/wrapper.py

@@ -7,14 +7,17 @@
 import logging
 import traceback
 
-from datadog_lambda.metric import lambda_stats
+from datadog_lambda.metric import lambda_stats, lambda_metric
 from datadog_lambda.patch import patch_all
+from datadog_lambda.tags import get_tags_from_context
 from datadog_lambda.tracing import (
     extract_dd_trace_context,
     set_correlation_ids,
     inject_correlation_ids,
 )
 
+ENHANCED_METRICS_NAMESPACE_PREFIX = "aws.lambda.enhanced"
+
 logger = logging.getLogger(__name__)
 
 
@@ -31,6 +34,10 @@ def my_lambda_handle(event, context):
     requests.get("https://www.datadoghq.com")
 """
 
+# On the first run of this Lambda container this variable is set
+# to the str value of the request ID, taken from the Lambda context
+cold_start_request_id = None
+
 
 class _LambdaDecorator(object):
     """
@@ -40,19 +47,29 @@ class _LambdaDecorator(object):
 
     def __init__(self, func):
         self.func = func
-        self.flush_to_log = os.environ.get('DD_FLUSH_TO_LOG', '').lower() == 'true'
-        self.logs_injection = os.environ.get('DD_LOGS_INJECTION', '').lower() == 'true'
+        self.flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
+        self.logs_injection = os.environ.get("DD_LOGS_INJECTION", "").lower() == "true"
 
         # Inject trace correlation ids to logs
         if self.logs_injection:
             inject_correlation_ids()
 
         # Patch HTTP clients to propagate Datadog trace context
         patch_all()
-        logger.debug('datadog_lambda_wrapper initialized')
+        logger.debug("datadog_lambda_wrapper initialized")
 
     def _before(self, event, context):
+        global cold_start_request_id
+        # Assign this request ID as the cold start if there is no value yet
+        if cold_start_request_id is None:
+            cold_start_request_id = context.aws_request_id
+
         try:
+            lambda_metric(
+                "{}.invocations".format(ENHANCED_METRICS_NAMESPACE_PREFIX),
+                1,
+                tags=get_tags_from_context(context, cold_start_request_id),
+            )
             # Extract Datadog trace context from incoming requests
             extract_dd_trace_context(event)
 
@@ -72,6 +89,13 @@ def __call__(self, event, context):
         self._before(event, context)
         try:
             return self.func(event, context)
+        except Exception:
+            lambda_metric(
+                "{}.errors".format(ENHANCED_METRICS_NAMESPACE_PREFIX),
+                1,
+                tags=get_tags_from_context(context, cold_start_request_id),
+            )
+            raise
         finally:
             self._after(event, context)
 

tests/test_tags.py

@@ -0,0 +1,60 @@
+import unittest
+
+from datadog_lambda.tags import parse_lambda_tags_from_arn, get_tags_from_context
+from tests.test_wrapper import get_mock_context
+
+
+class TestMetricTags(unittest.TestCase):
+    def test_parse_lambda_tags_from_arn(self):
+        self.assertListEqual(
+            parse_lambda_tags_from_arn(
+                "arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test"
+            ),
+            [
+                "region:us-east-1",
+                "account_id:1234597598159",
+                "functionname:swf-hello-test",
+            ],
+        )
+
+        self.assertListEqual(
+            parse_lambda_tags_from_arn(
+                "arn:aws:lambda:us-west-1:1234597598159:function:other-function:function-alias"
+            ),
+            [
+                "region:us-west-1",
+                "account_id:1234597598159",
+                "functionname:other-function",
+            ],
+        )
+
+    def test_get_tags_from_context(self):
+        cold_start_request_id = "first-request-id"
+        self.assertListEqual(
+            get_tags_from_context(
+                get_mock_context(aws_request_id=cold_start_request_id),
+                cold_start_request_id,
+            ),
+            [
+                "region:us-west-1",
+                "account_id:123457598159",
+                "functionname:python-layer-test",
+                "memorysize:256",
+                "cold_start:true",
+            ],
+        )
+
+        self.assertListEqual(
+            get_tags_from_context(
+                get_mock_context(aws_request_id="non-cold-start-request-id"),
+                cold_start_request_id,
+            ),
+            [
+                "region:us-west-1",
+                "account_id:123457598159",
+                "functionname:python-layer-test",
+                "memorysize:256",
+                "cold_start:false",
+            ],
+        )
+