Add better handling for Keep-Alive connections

Itay Neeman · Itay Neeman · commit 836ce8a28cf6 · 2015-12-09T16:20:06.000-08:00
Python's httplib generally has bad support for Keep-Alive connections,
which is why we never used it. We didn't explicitly ask for "Connection: Close",
but it was implicit. Additionally, splunkd would always return it.

However, an issue arises when talking to a Load Balancer which issues
a Keep-Alive request to splunkd, and returns that header to us. Python
then goes into a particular path which the code didn't handle gracefully.

What this change does is two-fold:

1. Always request to close the connection via the "Connection: Close" header.

2. If we get back a "Connection: Keep-Alive" header, we properly handle
it by not immediately closing our reference to the connection and only
doing so when the response has been completed.
diff --git a/splunklib/binding.py b/splunklib/binding.py
@@ -1235,8 +1235,9 @@ class ResponseReader(io.RawIOBase):
     # For testing, you can use a StringIO as the argument to
     # ``ResponseReader`` instead of an ``httplib.HTTPResponse``. It
     # will work equally well.
-    def __init__(self, response):
+    def __init__(self, response, connection=None):
         self._response = response
+        self._connection = connection
         self._buffer = ''
 
     def __str__(self):
@@ -1262,6 +1263,8 @@ def peek(self, size):
 
     def close(self):
         """Closes this response."""
+        if _connection:
+            _connection.close()
         self._response.close()
 
     def read(self, size = None):
@@ -1332,25 +1335,29 @@ def request(url, message, **kwargs):
             "Host": host,
             "User-Agent": "splunk-sdk-python/1.5.0",
             "Accept": "*/*",
+            "Connection": "Close",
         } # defaults
         for key, value in message["headers"]:
             head[key] = value
         method = message.get("method", "GET")
 
         connection = connect(scheme, host, port)
+        is_keepalive = False
         try:
             connection.request(method, path, body, head)
             if timeout is not None:
                 connection.sock.settimeout(timeout)
             response = connection.getresponse()
+            is_keepalive = "keep-alive" in response.getheader("connection", default="close").lower()
         finally:
-            connection.close()
+            if not is_keepalive:
+                connection.close()
 
         return {
             "status": response.status,
             "reason": response.reason,
             "headers": response.getheaders(),
-            "body": ResponseReader(response),
+            "body": ResponseReader(response, connection if is_keepalive else None),
         }
 
     return request
