Statsd integration (DataDog#89)

* Implement statsd protocol

* Bump version to 2.22.0

* Remove redundant whitespace

* Clean up from PR review

Author: DarcyRaynerDD

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__ = "2.21.0"
+__version__ = "2.22.0"
 
 
 import os

datadog_lambda/extension.py

@@ -0,0 +1,29 @@
+import logging
+import requests
+
+AGENT_URL = "http://127.0.0.1:8124"
+HELLO_PATH = "/lambda/hello"
+FLUSH_PATH = "/lambda/flush"
+
+logger = logging.getLogger(__name__)
+
+
+def is_extension_running():
+    try:
+        requests.get(AGENT_URL + HELLO_PATH)
+    except Exception as e:
+        logger.debug("Extension is not running, returned with error %s", e)
+        return False
+    return True
+
+
+def flush_extension():
+    try:
+        requests.post(AGENT_URL + FLUSH_PATH, data={})
+    except Exception as e:
+        logger.debug("Failed to flush extension, returned with error %s", e)
+        return False
+    return True
+
+
+should_use_extension = is_extension_running()

datadog_lambda/metric.py

@@ -10,20 +10,42 @@
 import logging
 
 import boto3
-from datadog import api
+from datadog import api, initialize, statsd
 from datadog.threadstats import ThreadStats
+from datadog_lambda.extension import should_use_extension
 from datadog_lambda.tags import get_enhanced_metrics_tags, tag_dd_lambda_layer
 
 
 ENHANCED_METRICS_NAMESPACE_PREFIX = "aws.lambda.enhanced"
 
 logger = logging.getLogger(__name__)
 
-lambda_stats = ThreadStats()
-lambda_stats.start()
+
+class StatsDWrapper:
+    """
+    Wraps StatsD calls, to give an identical interface to ThreadStats
+    """
+
+    def __init__(self):
+        options = {"statsd_host": "127.0.0.1", "statsd_port": 8125}
+        initialize(**options)
+
+    def distribution(self, metric_name, value, tags=[], timestamp=None):
+        statsd.distribution(metric_name, value, tags=tags)
+
+    def flush(self, value):
+        pass
 
 
-def lambda_metric(metric_name, value, timestamp=None, tags=None):
+lambda_stats = None
+if should_use_extension:
+    lambda_stats = StatsDWrapper()
+else:
+    lambda_stats = ThreadStats()
+    lambda_stats.start()
+
+
+def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=False):
     """
     Submit a data point to Datadog distribution metrics.
     https://docs.datadoghq.com/graphing/metrics/distributions/
@@ -36,12 +58,14 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None):
     periodically and at the end of the function execution in a
     background thread.
     """
+    flush_to_logs = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true"
     tags = tag_dd_lambda_layer(tags)
-    if os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true":
-        write_metric_point_to_stdout(metric_name, value, timestamp, tags)
+
+    if flush_to_logs or (force_async and not should_use_extension):
+        write_metric_point_to_stdout(metric_name, value, timestamp=timestamp, tags=tags)
     else:
         logger.debug("Sending metric %s to Datadog via lambda layer", metric_name)
-        lambda_stats.distribution(metric_name, value, timestamp=timestamp, tags=tags)
+        lambda_stats.distribution(metric_name, value, tags=tags, timestamp=timestamp)
 
 
 def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]):
@@ -85,13 +109,10 @@ def submit_enhanced_metric(metric_name, lambda_context):
             metric_name,
         )
         return
-
-    # Enhanced metrics are always written to logs
-    write_metric_point_to_stdout(
-        "{}.{}".format(ENHANCED_METRICS_NAMESPACE_PREFIX, metric_name),
-        1,
-        tags=get_enhanced_metrics_tags(lambda_context),
-    )
+    tags = get_enhanced_metrics_tags(lambda_context)
+    metric_name = "aws.lambda.enhanced." + metric_name
+    # Enhanced metrics always use an async submission method, (eg logs or extension).
+    lambda_metric(metric_name, 1, timestamp=None, tags=tags, force_async=True)
 
 
 def submit_invocations_metric(lambda_context):

datadog_lambda/tags.py

@@ -85,7 +85,6 @@ def get_enhanced_metrics_tags(lambda_context):
         get_cold_start_tag(),
         "memorysize:{}".format(lambda_context.memory_limit_in_mb),
         get_runtime_tag(),
-        _format_dd_lambda_layer_tag(),
     ]
 
 

datadog_lambda/wrapper.py

@@ -7,6 +7,7 @@
 import logging
 import traceback
 
+from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import set_cold_start, is_cold_start
 from datadog_lambda.metric import (
     lambda_stats,
@@ -137,8 +138,11 @@ def _before(self, event, context):
 
     def _after(self, event, context):
         try:
-            if not self.flush_to_log:
+            if not self.flush_to_log or should_use_extension:
                 lambda_stats.flush(float("inf"))
+            if should_use_extension:
+                flush_extension()
+
             if self.span:
                 self.span.finish()
             logger.debug("datadog_lambda_wrapper _after() done")

scripts/publish_staging.sh

@@ -1,6 +1,5 @@
 #!/bin/bash
 set -e
 
-./scripts/list_layers.sh
 ./scripts/build_layers.sh
 ./scripts/publish_layers.sh us-east-1

tests/test_wrapper.py

@@ -38,10 +38,6 @@ def setUp(self):
         self.mock_wrapper_lambda_stats = patcher.start()
         self.addCleanup(patcher.stop)
 
-        patcher = patch("datadog_lambda.metric.lambda_metric")
-        self.mock_lambda_metric = patcher.start()
-        self.addCleanup(patcher.stop)
-
         patcher = patch("datadog_lambda.wrapper.extract_dd_trace_context")
         self.mock_extract_dd_trace_context = patcher.start()
         self.addCleanup(patcher.stop)
@@ -152,6 +148,7 @@ def lambda_handler(event, context):
                         "runtime:python2.7",
                         "dd_lambda_layer:datadog-python27_0.1.0",
                     ],
+                    timestamp=None,
                 )
             ]
         )
@@ -181,6 +178,7 @@ def lambda_handler(event, context):
                         "runtime:python2.7",
                         "dd_lambda_layer:datadog-python27_0.1.0",
                     ],
+                    timestamp=None,
                 ),
                 call(
                     "aws.lambda.enhanced.errors",
@@ -195,6 +193,7 @@ def lambda_handler(event, context):
                         "runtime:python2.7",
                         "dd_lambda_layer:datadog-python27_0.1.0",
                     ],
+                    timestamp=None,
                 ),
             ]
         )
@@ -229,6 +228,7 @@ def lambda_handler(event, context):
                         "runtime:python2.7",
                         "dd_lambda_layer:datadog-python27_0.1.0",
                     ],
+                    timestamp=None,
                 ),
                 call(
                     "aws.lambda.enhanced.invocations",
@@ -243,6 +243,7 @@ def lambda_handler(event, context):
                         "runtime:python2.7",
                         "dd_lambda_layer:datadog-python27_0.1.0",
                     ],
+                    timestamp=None,
                 ),
             ]
         )
@@ -275,6 +276,7 @@ def lambda_handler(event, context):
                         "runtime:python2.7",
                         "dd_lambda_layer:datadog-python27_0.1.0",
                     ],
+                    timestamp=None,
                 )
             ]
         )
@@ -307,6 +309,7 @@ def lambda_handler(event, context):
                         "runtime:python2.7",
                         "dd_lambda_layer:datadog-python27_0.1.0",
                     ],
+                    timestamp=None,
                 )
             ]
         )