Use one connection for all communication with extension.

Author: purple4reina

datadog_lambda/extension.py

@@ -1,28 +1,29 @@
 import logging
 from os import path
 
-try:
-    # only available in python 3
-    # not an issue since the extension is not compatible with python 2.x runtime
-    # https://docs.aws.amazon.com/lambda/latest/dg/using-extensions.html
-    import urllib.request
-except ImportError:
-    # safe since both calls to urllib are protected with try/expect and will return false
-    urllib = None
-
-AGENT_URL = "http://127.0.0.1:8124"
 HELLO_PATH = "/lambda/hello"
 FLUSH_PATH = "/lambda/flush"
 EXTENSION_PATH = "/opt/extensions/datadog-agent"
 
 logger = logging.getLogger(__name__)
 
 
+try:
+    import http.client
+
+    conn = http.client.HTTPConnection("127.0.0.1", 8124)
+except Exception as e:
+    logger.debug("unable to create http connection to extension: ", e)
+    conn = None
+
+
 def is_extension_running():
     if not path.exists(EXTENSION_PATH):
         return False
     try:
-        urllib.request.urlopen(AGENT_URL + HELLO_PATH)
+        conn.request("GET", HELLO_PATH)
+        resp = conn.getresponse()
+        return resp.status == 200
     except Exception as e:
         logger.debug("Extension is not running, returned with error %s", e)
         return False
@@ -31,8 +32,9 @@ def is_extension_running():
 
 def flush_extension():
     try:
-        req = urllib.request.Request(AGENT_URL + FLUSH_PATH, "".encode("ascii"))
-        urllib.request.urlopen(req)
+        conn.request("POST", FLUSH_PATH, b"")
+        resp = conn.getresponse()
+        return resp.status == 200
     except Exception as e:
         logger.debug("Failed to flush extension, returned with error %s", e)
         return False