Merge pull request #328 from DataDog/quentin.pierre/aws-lambda-decorator

Lambda wrapper - start TLS connection on init & RequestClientKeptAlive

Author: Yann

datadog/api/http_client.py

@@ -8,6 +8,7 @@
 # stdlib
 import logging
 import urllib
+from threading import Lock
 
 # 3p
 try:
@@ -59,26 +60,30 @@ def request(cls, method, url, headers, params, data, timeout, proxies, verify, m
 
 class RequestClient(HTTPClient):
     """
-    HTTP client based on 3rd party `requests` module.
+    HTTP client based on 3rd party `requests` module, using a single session.
+    This allows us to keep the session alive to spare some execution time.
     """
+
+    _session = None
+    _session_lock = Lock()
+
     @classmethod
     def request(cls, method, url, headers, params, data, timeout, proxies, verify, max_retries):
-        """
-        """
         try:
-            # Use a session to set a max_retries parameters
-            with requests.Session() as s:
-                http_adapter = requests.adapters.HTTPAdapter(max_retries=max_retries)
-                s.mount('https://', http_adapter)
-
-                # Since stream=False we can close the session after this call
-                result = s.request(
-                    method, url,
-                    headers=headers, params=params, data=data,
-                    timeout=timeout,
-                    proxies=proxies, verify=verify)
-
-                result.raise_for_status()
+
+            with cls._session_lock:
+                if cls._session is None:
+                    cls._session = requests.Session()
+                    http_adapter = requests.adapters.HTTPAdapter(max_retries=max_retries)
+                    cls._session.mount('https://', http_adapter)
+
+            result = cls._session.request(
+                method, url,
+                headers=headers, params=params, data=data,
+                timeout=timeout,
+                proxies=proxies, verify=verify)
+
+            result.raise_for_status()
 
         except requests.ConnectionError as e:
             raise _remove_context(ClientError(method, url, e))

datadog/threadstats/aws_lambda.py

@@ -1,5 +1,5 @@
 from datadog.threadstats import ThreadStats
-from threading import Lock
+from threading import Lock, Thread
 from datadog import api
 import os
 
@@ -34,6 +34,11 @@ def _enter(cls):
                 cls._was_initialized = True
                 api._api_key = os.environ.get('DATADOG_API_KEY')
                 api._api_host = os.environ.get('DATADOG_HOST', 'https://api.datadoghq.com')
+
+                # Async initialization of the TLS connection with our endpoints
+                # This avoids adding execution time at the end of the lambda run
+                t = Thread(target=_init_api_client)
+                t.start()
             cls._counter = cls._counter + 1
 
     @classmethod
@@ -70,3 +75,19 @@ def __call__(self, *args, **kw):
 def lambda_metric(*args, **kw):
     """ Alias to expose only distributions for lambda functions"""
     _lambda_stats.distribution(*args, **kw)
+
+
+def _init_api_client():
+    """ No-op GET to initialize the requests connection with DD's endpoints
+
+    The goal here is to make the final flush faster:
+    we keep alive the Requests session, this means that we can re-use the connection
+    The consequence is that the HTTP Handshake, which can take hundreds of ms,
+    is now made at the beginning of a lambda instead of at the end.
+
+    By making the initial request async, we spare a lot of execution time in the lambdas.
+    """
+    try:
+        api.api_client.APIClient.submit('GET', 'validate')
+    except Exception:
+        pass

tests/unit/api/helper.py

@@ -8,6 +8,7 @@
 
 # datadog
 from datadog import initialize, api
+from datadog.api.http_client import RequestClient
 from datadog.api.exceptions import ApiError
 from datadog.api.resources import (
     CreateableAPIResource,
@@ -33,6 +34,25 @@
 }
 
 
+class MockSession(object):
+    """docstring for MockSession"""
+    _args = None
+    _kwargs = None
+    _count = 0
+
+    def request(self, *args, **kwargs):
+        self._args = args
+        self._kwargs = kwargs
+        self._count += 1
+        return MockResponse()
+
+    def call_args(self):
+        return self._args, self._kwargs
+
+    def call_count(self):
+        return self._count
+
+
 class MockResponse(requests.Response):
 
     def __init__(self, raise_for_status=False):
@@ -98,13 +118,15 @@ class DatadogAPITestCase(unittest.TestCase):
 
     def setUp(self):
         # Mock patch requests
-        self.request_patcher = patch('requests.Session')
-        request_class_mock = self.request_patcher.start()
-        self.request_mock = request_class_mock.return_value.__enter__.return_value
-        self.request_mock.request = Mock(return_value=MockResponse())
+        self.request_mock = MockSession()
+        RequestClient._session = self.request_mock
+        # self.request_patcher = patch('requests.Session')
+        # request_class_mock = self.request_patcher.start()
+        # self.request_mock = request_class_mock.return_value
+        # self.request_mock.request = Mock(return_value=MockResponse())
 
     def tearDown(self):
-        self.request_patcher.stop()
+        del RequestClient._session
 
     def arm_requests_to_raise(self):
         """
@@ -116,11 +138,11 @@ def get_request_data(self):
         """
         Returns JSON formatted data from the submitted `requests`.
         """
-        _, kwargs = self.request_mock.request.call_args
+        _, kwargs = self.request_mock.call_args()
         return json.loads(kwargs['data'])
 
     def request_called_with(self, method, url, data=None, params=None):
-        (req_method, req_url), others = self.request_mock.request.call_args
+        (req_method, req_url), others = self.request_mock.call_args()
         self.assertEquals(method, req_method, req_method)
         self.assertEquals(url, req_url, req_url)
 

tests/unit/api/test_api.py

@@ -73,7 +73,7 @@ def test_no_initialization_fails(self):
         # Finally, initialize with an API key
         initialize(api_key=API_KEY, api_host=API_HOST)
         MyCreatable.create()
-        self.assertEquals(self.request_mock.request.call_count, 1)
+        self.assertEquals(self.request_mock.call_count(), 1)
 
     @mock.patch('datadog.util.config.get_config_path')
     def test_get_hostname(self, mock_config_path):
@@ -105,7 +105,7 @@ def test_request_parameters(self):
         # Make a simple API call
         MyCreatable.create()
 
-        _, options = self.request_mock.request.call_args
+        _, options = self.request_mock.call_args()
 
         # Assert `requests` parameters
         self.assertIn('params', options)
@@ -128,7 +128,7 @@ def test_initialize_options(self):
         # Make a simple API call
         MyCreatable.create()
 
-        _, options = self.request_mock.request.call_args
+        _, options = self.request_mock.call_args()
 
         # Assert `requests` parameters
         self.assertIn('proxies', options)