diff --git a/src/appier/http.py b/src/appier/http.py index 6495affe..8c0bfb2f 100644 --- a/src/appier/http.py +++ b/src/appier/http.py @@ -908,10 +908,20 @@ def _client_netius(level=logging.CRITICAL): if netius_client: return netius_client + # reaps stale clients from dead threads before creating + # a new one, this prevents unbounded growth of the global + # clients map in long-running multi-threaded applications + _reap_netius_clients() + # creates the "new" HTTP client for the current thread and registers - # it under the netius client structure so that it may be re-used + # it under the netius client structure so that it may be re-used, the + # insertion is done inside the lock to avoid TOCTOU race conditions netius_client = netius.clients.HTTPClient(auto_release=False) - _netius_clients[tid] = netius_client + ACCESS_LOCK.acquire() + try: + _netius_clients[tid] = netius_client + finally: + ACCESS_LOCK.release() # in case this is the first registration of the dictionary a new on # exit callback is registered to cleanup the netius infra-structure @@ -983,6 +993,26 @@ def _callback(protocol, parser, message): return extra +def _reap_netius_clients(): + global _netius_clients + if not _netius_clients: + return + live_tids = set(t.ident for t in threading.enumerate()) + ACCESS_LOCK.acquire() + try: + dead_tids = [tid for tid in _netius_clients if tid not in live_tids] + for tid in dead_tids: + client = _netius_clients.pop(tid, None) + if not client: + continue + try: + client.cleanup() + except Exception: + pass + finally: + ACCESS_LOCK.release() + + def _cleanup_netius(): global _netius_clients for netius_client in _netius_clients.values(): diff --git a/src/appier/test/http.py b/src/appier/test/http.py index 56926122..605fd0ac 100644 --- a/src/appier/test/http.py +++ b/src/appier/test/http.py @@ -29,6 +29,7 @@ """ The license for the module """ import unittest +import unittest.mock import threading import appier @@ -262,3 +263,50 @@ def test_invalid(self): self.assertRaises( BaseException, lambda: appier.get("https://invalidlargedomain.org/") ) + + def test_reap_netius_clients(self): + dead_client = unittest.mock.MagicMock() + alive_client = unittest.mock.MagicMock() + + alive_tid = threading.current_thread().ident + dead_tid = 99999999 + + original = getattr(appier.http, "_netius_clients", None) + try: + appier.http._netius_clients = { + dead_tid: dead_client, + alive_tid: alive_client, + } + + appier.http._reap_netius_clients() + + self.assertEqual(dead_tid in appier.http._netius_clients, False) + self.assertEqual(alive_tid in appier.http._netius_clients, True) + dead_client.cleanup.assert_called_once() + alive_client.cleanup.assert_not_called() + finally: + if original is None: + if hasattr(appier.http, "_netius_clients"): + del appier.http._netius_clients + else: + appier.http._netius_clients = original + + def test_client_netius_cleanup(self): + dead_client = unittest.mock.MagicMock() + dead_tid = 99999999 + + original = getattr(appier.http, "_netius_clients", None) + try: + appier.http._netius_clients = {dead_tid: dead_client} + + appier.http._reap_netius_clients() + + self.assertEqual(dead_tid in appier.http._netius_clients, False) + self.assertEqual(len(appier.http._netius_clients), 0) + dead_client.cleanup.assert_called_once() + finally: + if original is None: + if hasattr(appier.http, "_netius_clients"): + del appier.http._netius_clients + else: + appier.http._netius_clients = original