Merge pull request #19 from DataDog/swf/submit-invocations-errors

Submit metrics for invocations and errors

Author: swfunc

CHANGELOG.md

@@ -1,5 +1,9 @@
 # CHANGELOG
 
+# Version: 7 / 2019-10-24
+
+- Increment `aws.lambda.enhanced.invocations` and `aws.lambda.enhanced.errors` metrics for each invocation if `DD_ENHANCED_METRICS` env var is set to true.
+
 # Version: 6 / 2019-09-16
 
 - Support `DD_LOGS_INJECTION` for trace and log correlation

datadog_lambda/__init__.py

@@ -1,6 +1,6 @@
 # The minor version corresponds to the Lambda layer version.
 # E.g.,, version 0.5.0 gets packaged into layer version 5.
-__version__ = '0.6.0'
+__version__ = '0.7.0'
 
 
 import os

datadog_lambda/cold_start.py

@@ -0,0 +1,25 @@
+_cold_start = True
+_lambda_container_initialized = False
+
+
+def set_cold_start():
+    """Set the value of the cold start global
+
+    This should be executed once per Lambda execution before the execution
+    """
+    global _cold_start
+    global _lambda_container_initialized
+    _cold_start = not _lambda_container_initialized
+    _lambda_container_initialized = True
+
+
+def is_cold_start():
+    """Returns the value of the global cold_start
+    """
+    return _cold_start
+
+
+def get_cold_start_tag():
+    """Returns the cold start tag to be used in metrics
+    """
+    return "cold_start:{}".format(str(is_cold_start()).lower())

datadog_lambda/metric.py

@@ -14,6 +14,11 @@
 from datadog import api
 from datadog.threadstats import ThreadStats
 from datadog_lambda import __version__
+from datadog_lambda.cold_start import get_cold_start_tag
+from datadog_lambda.tags import parse_lambda_tags_from_arn
+
+
+ENHANCED_METRICS_NAMESPACE_PREFIX = "aws.lambda.enhanced"
 
 logger = logging.getLogger(__name__)
 
@@ -54,39 +59,71 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None):
     background thread.
     """
     tags = _tag_dd_lambda_layer(tags)
-    if os.environ.get('DD_FLUSH_TO_LOG', '').lower() == 'true':
-        logger.debug('Sending metric %s to Datadog via log forwarder', metric_name)
-        print(json.dumps({
-            'm': metric_name,
-            'v': value,
-            'e': timestamp or int(time.time()),
-            't': tags
-        }))
-    else:
-        logger.debug('Sending metric %s to Datadog via lambda layer', metric_name)
-        lambda_stats.distribution(
-            metric_name, value, timestamp=timestamp, tags=tags
+    if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
+        logger.debug("Sending metric %s to Datadog via log forwarder", metric_name)
+        print(
+            json.dumps(
+                {
+                    "m": metric_name,
+                    "v": value,
+                    "e": timestamp or int(time.time()),
+                    "t": tags,
+                }
+            )
         )
+    else:
+        logger.debug("Sending metric %s to Datadog via lambda layer", metric_name)
+        lambda_stats.distribution(metric_name, value, timestamp=timestamp, tags=tags)
+
+
+def are_enhanced_metrics_enabled():
+    """Check env var to find if enhanced metrics should be submitted
+    """
+    return os.environ.get("DD_ENHANCED_METRICS", "false").lower() == "true"
+
+
+def submit_invocations_metric(lambda_arn):
+    """Increment aws.lambda.enhanced.invocations by 1
+    """
+    if not are_enhanced_metrics_enabled():
+        return
+
+    lambda_metric(
+        "{}.invocations".format(ENHANCED_METRICS_NAMESPACE_PREFIX),
+        1,
+        tags=parse_lambda_tags_from_arn(lambda_arn) + [get_cold_start_tag()],
+    )
+
+
+def submit_errors_metric(lambda_arn):
+    """Increment aws.lambda.enhanced.errors by 1
+    """
+    if not are_enhanced_metrics_enabled():
+        return
+
+    lambda_metric(
+        "{}.errors".format(ENHANCED_METRICS_NAMESPACE_PREFIX),
+        1,
+        tags=parse_lambda_tags_from_arn(lambda_arn) + [get_cold_start_tag()],
+    )
 
 
 # Decrypt code should run once and variables stored outside of the function
 # handler so that these are decrypted once per container
-DD_KMS_API_KEY = os.environ.get('DD_KMS_API_KEY', '')
+DD_KMS_API_KEY = os.environ.get("DD_KMS_API_KEY", "")
 if DD_KMS_API_KEY:
-    DD_KMS_API_KEY = boto3.client('kms').decrypt(
+    DD_KMS_API_KEY = boto3.client("kms").decrypt(
         CiphertextBlob=base64.b64decode(DD_KMS_API_KEY)
-    )['Plaintext']
+    )["Plaintext"]
 
 # Set API Key and Host in the module, so they only set once per container
 api._api_key = os.environ.get(
-    'DATADOG_API_KEY',
-    os.environ.get('DD_API_KEY', DD_KMS_API_KEY),
+    "DATADOG_API_KEY", os.environ.get("DD_API_KEY", DD_KMS_API_KEY)
 )
-logger.debug('Setting DATADOG_API_KEY of length %d', len(api._api_key))
+logger.debug("Setting DATADOG_API_KEY of length %d", len(api._api_key))
 
 # Set DATADOG_HOST, to send data to a non-default Datadog datacenter
 api._api_host = os.environ.get(
-    'DATADOG_HOST',
-    'https://api.' + os.environ.get('DD_SITE', 'datadoghq.com')
+    "DATADOG_HOST", "https://api." + os.environ.get("DD_SITE", "datadoghq.com")
 )
-logger.debug('Setting DATADOG_HOST to %s', api._api_host)
+logger.debug("Setting DATADOG_HOST to %s", api._api_host)

datadog_lambda/tags.py

@@ -0,0 +1,20 @@
+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),
+    ]

datadog_lambda/wrapper.py

@@ -7,14 +7,20 @@
 import logging
 import traceback
 
-from datadog_lambda.metric import lambda_stats
+from datadog_lambda.cold_start import set_cold_start
+from datadog_lambda.metric import (
+    lambda_stats,
+    submit_invocations_metric,
+    submit_errors_metric,
+)
 from datadog_lambda.patch import patch_all
 from datadog_lambda.tracing import (
     extract_dd_trace_context,
     set_correlation_ids,
     inject_correlation_ids,
 )
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -40,19 +46,22 @@ 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):
+        set_cold_start()
+
         try:
+            submit_invocations_metric(context.invoked_function_arn)
             # Extract Datadog trace context from incoming requests
             extract_dd_trace_context(event)
 
@@ -72,6 +81,9 @@ def __call__(self, event, context):
         self._before(event, context)
         try:
             return self.func(event, context)
+        except Exception:
+            submit_errors_metric(context.invoked_function_arn)
+            raise
         finally:
             self._after(event, context)
 

tests/test_tags.py

@@ -0,0 +1,29 @@
+import unittest
+
+from datadog_lambda.tags import parse_lambda_tags_from_arn
+
+
+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",
+            ],
+        )
+