Flush metrics to log

Author: tianchu

.flake8

@@ -0,0 +1,2 @@
+[flake8]
+max-line-length = 100

Dockerfile

@@ -12,7 +12,7 @@ COPY requirements.txt requirements.txt
 RUN pip install -r requirements.txt -t ./python/lib/$runtime/site-packages
 
 # Install datadog_lambda
-COPY datadog_lambda ./python/lib/$runtime/site-packages
+COPY datadog_lambda ./python/lib/$runtime/site-packages/datadog_lambda
 
 # Remove *.pyc files
 RUN find ./python/lib/$runtime/site-packages -name \*.pyc -delete

datadog_lambda/constants.py

@@ -3,6 +3,7 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+
 # Datadog trace sampling priority
 class SamplingPriority(object):
     USER_REJECT = -1

datadog_lambda/metric.py

@@ -5,6 +5,8 @@
 
 import os
 import sys
+import json
+import time
 
 from datadog import api
 from datadog.threadstats import ThreadStats
@@ -22,26 +24,42 @@ def _format_dd_lambda_layer_tag():
     return "dd_lambda_layer:datadog-{}_{}".format(runtime, __version__)
 
 
-def _tag_dd_lambda_layer(args, kwargs):
+def _tag_dd_lambda_layer(tags):
     """
     Used by lambda_metric to insert the dd_lambda_layer tag
     """
     dd_lambda_layer_tag = _format_dd_lambda_layer_tag()
-    if 'tags' in kwargs:
-        kwargs['tags'].append(dd_lambda_layer_tag)
-    elif len(args) >= 4:
-        args[3].append(dd_lambda_layer_tag)
+    if tags:
+        return tags + [dd_lambda_layer_tag]
     else:
-        kwargs['tags'] = [dd_lambda_layer_tag]
+        return [dd_lambda_layer_tag]
 
 
-def lambda_metric(*args, **kwargs):
+def lambda_metric(metric_name, value, timestamp=None, tags=None):
     """
     Submit a data point to Datadog distribution metrics.
     https://docs.datadoghq.com/graphing/metrics/distributions/
+
+    When DATADOG_LOG_FORWARDER is True, write metric to log, and
+    wait for the Datadog Log Forwarder Lambda function to submit
+    the metrics asynchronously.
+
+    Otherwise, the metrics will be submitted to the Datadog API
+    periodically and at the end of the function execution in a
+    background thread.
     """
-    _tag_dd_lambda_layer(args, kwargs)
-    lambda_stats.distribution(*args, **kwargs)
+    tags = _tag_dd_lambda_layer(tags)
+    if os.environ.get('DATADOG_FLUSH_TO_LOG') == 'True':
+        print(json.dumps({
+            'metric_name': metric_name,
+            'value': value,
+            'timestamp': timestamp or int(time.time()),
+            'tags': tags
+        }))
+    else:
+        lambda_stats.distribution(
+            metric_name, value, timestamp=timestamp, tags=tags
+        )
 
 
 def init_api_client():

datadog_lambda/tracing.py

@@ -55,8 +55,7 @@ def get_dd_trace_context():
     Return the Datadog trace context to be propogated on the outgoing requests.
 
     If the Lambda function is invoked by a Datadog-traced service, a Datadog
-    trace
-    context may already exist, and it should be used. Otherwise, use the
+    trace context may already exist, and it should be used. Otherwise, use the
     current X-Ray trace entity.
 
     Most of widely-used HTTP clients are patched to inject the context

datadog_lambda/wrapper.py

@@ -3,6 +3,7 @@
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
 
+import os
 import traceback
 from threading import Thread
 
@@ -33,12 +34,14 @@ class _LambdaDecorator(object):
 
     def __init__(self, func):
         self.func = func
+        self.flush_to_log = os.environ.get('DATADOG_FLUSH_TO_LOG') == 'True'
 
     def _before(self, event, context):
         try:
             # Async initialization of the TLS connection with Datadog API,
             # and reduces the overhead to the final metric flush at the end.
-            Thread(target=init_api_client).start()
+            if not self.flush_to_log:
+                Thread(target=init_api_client).start()
 
             # Extract Datadog trace context from incoming requests
             extract_dd_trace_context(event)
@@ -50,7 +53,8 @@ def _before(self, event, context):
 
     def _after(self, event, context):
         try:
-            lambda_stats.flush(float("inf"))
+            if not self.flush_to_log:
+                lambda_stats.flush(float("inf"))
         except Exception:
             traceback.print_exc()
 

requirements.txt

@@ -1,3 +1,4 @@
 aws-xray-sdk==2.4.2
 datadog==0.28.0
-wrapt==1.11.1
+wrapt==1.11.1
+setuptools==40.8.0

tests/Dockerfile

@@ -3,6 +3,7 @@ FROM python:$python_version
 
 ENV PYTHONDONTWRITEBYTECODE True
 
+COPY .flake8 .flake8
 COPY requirements.txt requirements.txt
 COPY requirements-dev.txt requirements-dev.txt
 RUN pip install -r requirements.txt

tests/test_metric.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 try:
     from unittest.mock import patch, call
@@ -18,12 +19,20 @@ def setUp(self):
         self.addCleanup(patcher.stop)
 
     def test_lambda_metric_tagged_with_dd_lambda_layer(self):
-        lambda_metric('test.metric', 1)
-        lambda_metric('test.metric', 1, 123, ['tag1:test'])
-        lambda_metric('test.metric', 1, tags=['tag1:test'])
+        lambda_metric('test', 1)
+        lambda_metric('test', 1, 123, [])
+        lambda_metric('test', 1, tags=['tag1:test'])
         expected_tag = _format_dd_lambda_layer_tag()
         self.mock_metric_lambda_stats.distribution.assert_has_calls([
-            call('test.metric', 1, tags=[expected_tag]),
-            call('test.metric', 1, 123, ['tag1:test', expected_tag]),
-            call('test.metric', 1, tags=['tag1:test', expected_tag]),
+            call('test', 1, timestamp=None, tags=[expected_tag]),
+            call('test', 1, timestamp=123, tags=[expected_tag]),
+            call('test', 1, timestamp=None, tags=['tag1:test', expected_tag]),
         ])
+
+    def test_lambda_metric_flush_to_log(self):
+        os.environ["DATADOG_FLUSH_TO_LOG"] = 'True'
+
+        lambda_metric('test', 1)
+        self.mock_metric_lambda_stats.distribution.assert_not_called()
+
+        del os.environ["DATADOG_FLUSH_TO_LOG"]

tests/test_wrapper.py

@@ -1,3 +1,4 @@
+import os
 import unittest
 try:
     from unittest.mock import patch, call, ANY
@@ -35,9 +36,28 @@ def lambda_handler(event, context):
         lambda_event = {}
         lambda_context = {}
         lambda_handler(lambda_event, lambda_context)
+
         self.mock_metric_lambda_stats.distribution.assert_has_calls([
-            call('test.metric', 100, tags=ANY)
+            call('test.metric', 100, timestamp=None, tags=ANY)
         ])
         self.mock_wrapper_lambda_stats.flush.assert_called()
         self.mock_extract_dd_trace_context.assert_called_with(lambda_event)
         self.mock_patch_all.assert_called()
+
+    def test_datadog_lambda_wrapper_flush_to_log(self):
+        os.environ["DATADOG_FLUSH_TO_LOG"] = 'True'
+
+        @datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            lambda_metric("test.metric", 100)
+
+        lambda_event = {}
+        lambda_context = {}
+        lambda_handler(lambda_event, lambda_context)
+
+        self.mock_metric_lambda_stats.distribution.assert_not_called()
+        self.mock_wrapper_lambda_stats.flush.assert_not_called()
+        self.mock_extract_dd_trace_context.assert_called_with(lambda_event)
+        self.mock_patch_all.assert_called()
+
+        del os.environ["DATADOG_FLUSH_TO_LOG"]