Implement PR feedback

sfirrin · sfirrin · commit 6ab22dd164e5 · 2019-10-07T14:47:47.000-04:00
diff --git a/datadog_lambda/cold_start.py b/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())
diff --git a/datadog_lambda/metric.py b/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,59 @@ 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 submit_invocations_metric(lambda_arn):
+    """Increment aws.lambda.enhanced.invocations by 1
+    """
+    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
+    """
+    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)
diff --git a/datadog_lambda/tags.py b/datadog_lambda/tags.py
@@ -18,23 +18,3 @@ def parse_lambda_tags_from_arn(arn):
         "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
diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py
@@ -7,16 +7,19 @@
 import logging
 import traceback
 
-from datadog_lambda.metric import lambda_stats, lambda_metric
+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.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__)
 
@@ -34,10 +37,6 @@ 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):
     """
@@ -59,17 +58,10 @@ def __init__(self, func):
         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
+        set_cold_start()
 
         try:
-            lambda_metric(
-                "{}.invocations".format(ENHANCED_METRICS_NAMESPACE_PREFIX),
-                1,
-                tags=get_tags_from_context(context, cold_start_request_id),
-            )
+            submit_invocations_metric(context.invoked_function_arn)
             # Extract Datadog trace context from incoming requests
             extract_dd_trace_context(event)
 
@@ -90,11 +82,7 @@ def __call__(self, 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),
-            )
+            submit_errors_metric(context.invoked_function_arn)
             raise
         finally:
             self._after(event, context)
diff --git a/tests/test_tags.py b/tests/test_tags.py
@@ -1,7 +1,6 @@
 import unittest
 
-from datadog_lambda.tags import parse_lambda_tags_from_arn, get_tags_from_context
-from tests.test_wrapper import get_mock_context
+from datadog_lambda.tags import parse_lambda_tags_from_arn
 
 
 class TestMetricTags(unittest.TestCase):
@@ -28,33 +27,3 @@ def test_parse_lambda_tags_from_arn(self):
             ],
         )
 
-    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",
-            ],
-        )
-
diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py
@@ -6,7 +6,7 @@
 except ImportError:
     from mock import patch, call, ANY, MagicMock
 
-from datadog_lambda.wrapper import datadog_lambda_wrapper, cold_start_request_id
+from datadog_lambda.wrapper import datadog_lambda_wrapper
 from datadog_lambda.metric import lambda_metric
 
 
@@ -32,7 +32,7 @@ def setUp(self):
         self.mock_wrapper_lambda_stats = patcher.start()
         self.addCleanup(patcher.stop)
 
-        patcher = patch("datadog_lambda.wrapper.lambda_metric")
+        patcher = patch("datadog_lambda.metric.lambda_metric")
         self.mock_wrapper_lambda_metric = patcher.start()
         self.addCleanup(patcher.stop)
 
@@ -52,6 +52,11 @@ def setUp(self):
         self.mock_patch_all = patcher.start()
         self.addCleanup(patcher.stop)
 
+        patcher = patch("datadog_lambda.cold_start.is_cold_start")
+        self.mock_is_cold_start = patcher.start()
+        self.mock_is_cold_start.return_value = True
+        self.addCleanup(patcher.stop)
+
     def test_datadog_lambda_wrapper(self):
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -118,7 +123,6 @@ def lambda_handler(event, context):
                         "region:us-west-1",
                         "account_id:123457598159",
                         "functionname:python-layer-test",
-                        "memorysize:256",
                         "cold_start:true",
                     ],
                 )
@@ -144,7 +148,6 @@ def lambda_handler(event, context):
                         "region:us-west-1",
                         "account_id:123457598159",
                         "functionname:python-layer-test",
-                        "memorysize:256",
                         "cold_start:true",
                     ],
                 ),
@@ -155,22 +158,23 @@ def lambda_handler(event, context):
                         "region:us-west-1",
                         "account_id:123457598159",
                         "functionname:python-layer-test",
-                        "memorysize:256",
                         "cold_start:true",
                     ],
                 ),
             ]
         )
 
-    def test_cold_start_tag(self):
+    def test_enhanced_metrics_cold_start_tag(self):
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
             lambda_metric("test.metric", 100)
 
         lambda_event = {}
-        self.assertIsNone(cold_start_request_id)
+
         lambda_handler(lambda_event, get_mock_context())
 
+        self.mock_is_cold_start.return_value = False
+
         lambda_handler(
             lambda_event, get_mock_context(aws_request_id="second-request-id")
         )
@@ -184,7 +188,6 @@ def lambda_handler(event, context):
                         "region:us-west-1",
                         "account_id:123457598159",
                         "functionname:python-layer-test",
-                        "memorysize:256",
                         "cold_start:true",
                     ],
                 ),
@@ -195,7 +198,6 @@ def lambda_handler(event, context):
                         "region:us-west-1",
                         "account_id:123457598159",
                         "functionname:python-layer-test",
-                        "memorysize:256",
                         "cold_start:false",
                     ],
                 ),
