Skip to content

Instantly share code, notes, and snippets.

@stephan-hof
Last active October 29, 2021 18:17
Show Gist options
  • Save stephan-hof/c87aefa776779e2bc1dccec649d0d663 to your computer and use it in GitHub Desktop.
Save stephan-hof/c87aefa776779e2bc1dccec649d0d663 to your computer and use it in GitHub Desktop.
Make urllib3 thread safe.
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