diff --git a/requests_unixsocket/adapters.py b/requests_unixsocket/adapters.py index 0936742..37351e5 100644 --- a/requests_unixsocket/adapters.py +++ b/requests_unixsocket/adapters.py @@ -2,17 +2,21 @@ import socket from requests.adapters import HTTPAdapter from requests.compat import urlparse, unquote + try: - from requests.packages.urllib3.connection import HTTPConnection - from requests.packages.urllib3.connectionpool import HTTPConnectionPool + import http.client as httplib except ImportError: - from urllib3.connection import HTTPConnection - from urllib3.connectionpool import HTTPConnectionPool + import httplib + +try: + from requests.packages import urllib3 +except ImportError: + import urllib3 # The following was adapted from some code from docker-py -# https://github.com/docker/docker-py/blob/master/docker/unixconn/unixconn.py -class UnixHTTPConnection(HTTPConnection): +# https://github.com/docker/docker-py/blob/master/docker/transport/unixconn.py +class UnixHTTPConnection(httplib.HTTPConnection, object): def __init__(self, unix_socket_url, timeout=60): """Create an HTTP connection to a unix domain socket @@ -21,7 +25,7 @@ class UnixHTTPConnection(HTTPConnection): netloc is a percent-encoded path to a unix domain socket. E.g.: 'http+unix://%2Ftmp%2Fprofilesvc.sock/status/pid' """ - HTTPConnection.__init__(self, 'localhost', timeout=timeout) + super(UnixHTTPConnection, self).__init__('localhost', timeout=timeout) self.unix_socket_url = unix_socket_url self.timeout = timeout @@ -33,10 +37,11 @@ class UnixHTTPConnection(HTTPConnection): self.sock = sock -class UnixHTTPConnectionPool(HTTPConnectionPool): +class UnixHTTPConnectionPool(urllib3.connectionpool.HTTPConnectionPool): def __init__(self, socket_path, timeout=60): - HTTPConnectionPool.__init__(self, 'localhost', timeout=timeout) + super(UnixHTTPConnectionPool, self).__init__( + 'localhost', timeout=timeout) self.socket_path = socket_path self.timeout = timeout @@ -46,18 +51,34 @@ class UnixHTTPConnectionPool(HTTPConnectionPool): class UnixAdapter(HTTPAdapter): - def __init__(self, timeout=60): + def __init__(self, timeout=60, pool_connections=25): super(UnixAdapter, self).__init__() self.timeout = timeout + self.pools = urllib3._collections.RecentlyUsedContainer( + pool_connections, dispose_func=lambda p: p.close() + ) + super(UnixAdapter, self).__init__() - def get_connection(self, socket_path, proxies=None): + def get_connection(self, url, proxies=None): proxies = proxies or {} - proxy = proxies.get(urlparse(socket_path.lower()).scheme) + proxy = proxies.get(urlparse(url.lower()).scheme) if proxy: raise ValueError('%s does not support specifying proxies' % self.__class__.__name__) - return UnixHTTPConnectionPool(socket_path, self.timeout) + + with self.pools.lock: + pool = self.pools.get(url) + if pool: + return pool + + pool = UnixHTTPConnectionPool(url, self.timeout) + self.pools[url] = pool + + return pool def request_url(self, request, proxies): return request.path_url + + def close(self): + self.pools.clear()