Implement PR feedback

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,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)

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

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)

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",
- ],
- )
-

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",
  ],
  ),