-
-
Save stephan-hof/c87aefa776779e2bc1dccec649d0d663 to your computer and use it in GitHub Desktop.
Make urllib3 thread safe.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
diff --git a/src/urllib3/connectionpool.py b/src/urllib3/connectionpool.py | |
index 68015357..45f223a7 100644 | |
--- a/src/urllib3/connectionpool.py | |
+++ b/src/urllib3/connectionpool.py | |
@@ -4,6 +4,8 @@ import queue | |
import socket | |
import sys | |
import warnings | |
+import weakref | |
+ | |
from http.client import HTTPResponse as _HttplibHTTPResponse | |
from socket import timeout as SocketTimeout | |
from types import TracebackType | |
@@ -129,6 +131,39 @@ class ConnectionPool: | |
_blocking_errnos = {errno.EAGAIN, errno.EWOULDBLOCK} | |
+# This is basically a backport to make it run under python2: | |
+# https://docs.python.org/3.9/library/weakref.html#finalizer-objects | |
+class finalize(object): | |
+ __slots__ = () | |
+ | |
+ _active = {} | |
+ | |
+ class State(object): | |
+ __slots__ = ('weakref', 'callback') | |
+ | |
+ def __init__(self, ob, callback): | |
+ state = self.State() | |
+ state.weakref = weakref.ref(ob, self) | |
+ state.callback = callback | |
+ | |
+ self._active[self] = state | |
+ | |
+ def __call__(self, _=None): | |
+ state = self._active.pop(self, None) | |
+ if state: | |
+ state.callback() | |
+ | |
+ | |
+def close_connections(pool): | |
+ try: | |
+ while True: | |
+ conn = pool.get(block=False) | |
+ if conn: | |
+ conn.close() | |
+ except queue.Empty: | |
+ pass # Done. | |
+ | |
+ | |
class HTTPConnectionPool(ConnectionPool, RequestMethods): | |
""" | |
Thread-safe connection pool for one host. | |
@@ -236,6 +271,16 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods): | |
self.conn_kw["proxy"] = self.proxy | |
self.conn_kw["proxy_config"] = self.proxy_config | |
+ # Do not pass 'self' as callback to 'finalize'. | |
+ # Then the 'finalize' would keep an endless living (leak) to self. | |
+ # By just passing a reference to the pool allows the garbage collector | |
+ # to free self if nobody else has a reference to it. | |
+ pool = self.pool | |
+ | |
+ # This calls 'close_connections(pools)' just before HTTPConnectionPool | |
+ # gets gargabe collected. | |
+ finalize(self, lambda: close_connections(pool)) | |
+ | |
def _new_conn(self) -> HTTPConnection: | |
""" | |
Return a fresh :class:`HTTPConnection`. | |
@@ -500,14 +545,7 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods): | |
# Disable access to the pool | |
old_pool, self.pool = self.pool, None | |
- try: | |
- while True: | |
- conn = old_pool.get(block=False) | |
- if conn: | |
- conn.close() | |
- | |
- except queue.Empty: | |
- pass # Done. | |
+ close_connections(old_pool) | |
def is_same_host(self, url: str) -> bool: | |
""" | |
diff --git a/src/urllib3/poolmanager.py b/src/urllib3/poolmanager.py | |
index a36f96ae..81ba25d8 100644 | |
--- a/src/urllib3/poolmanager.py | |
+++ b/src/urllib3/poolmanager.py | |
@@ -214,11 +214,8 @@ class PoolManager(RequestMethods): | |
super().__init__(headers) | |
self.connection_pool_kw = connection_pool_kw | |
- def dispose_func(p: Any) -> None: | |
- p.close() | |
- | |
self.pools: RecentlyUsedContainer[PoolKey, HTTPConnectionPool] | |
- self.pools = RecentlyUsedContainer(num_pools, dispose_func=dispose_func) | |
+ self.pools = RecentlyUsedContainer(num_pools) | |
# Locally set the pool classes and keys so other PoolManagers can | |
# override them. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment