feat (diracx): add a randomized connection pooling

Author: chaen

src/DIRAC/FrameworkSystem/Utilities/diracx.py

@@ -1,10 +1,16 @@
+import random
 import requests
 
+
 from cachetools import TTLCache, cached
 from cachetools.keys import hashkey
 from pathlib import Path
+from requests.adapters import HTTPAdapter
 from tempfile import NamedTemporaryFile
 from typing import Any
+from urllib3 import PoolManager
+from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool
+
 from DIRAC import gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 
@@ -26,6 +32,65 @@
 DEFAULT_TOKEN_CACHE_SIZE = 1024
 
 
+class RandomizedPoolManager(PoolManager):
+    """
+    A PoolManager subclass that creates multiple connection pools per host.
+    Each connection request randomly picks one of the available pools.
+    """
+
+    def __init__(self, num_pools=3, **kwargs):
+        self.num_pools = num_pools
+        super().__init__(**kwargs)
+
+    def connection_from_host(self, host, port=None, scheme="http"):
+        # Pick a random index to diversify the pool key.
+        rand_index = random.randint(0, self.num_pools - 1)
+        pool_key = (f"{host}-{rand_index}", port, scheme)
+
+        if pool_key in self.pools:
+            return self.pools[pool_key]
+
+        # Create a new pool if none exists for this key.
+        if scheme == "http":
+            self.pools[pool_key] = HTTPConnectionPool(host, port, **self.connection_pool_kw)
+        elif scheme == "https":
+            self.pools[pool_key] = HTTPSConnectionPool(host, port, **self.connection_pool_kw)
+        else:
+            raise ValueError(f"Unsupported scheme: {scheme}")
+
+        return self.pools[pool_key]
+
+
+class RandomizedHTTPAdapter(HTTPAdapter):
+    """
+    An HTTPAdapter that uses the RandomizedPoolManager.
+    """
+
+    def __init__(self, num_pools=3, maxsize=10, **kwargs):
+        self.num_pools = num_pools
+        self.custom_maxsize = maxsize
+        super().__init__(**kwargs)
+
+    def init_poolmanager(self, connections, maxsize, block=False, **pool_kwargs):
+        """
+        Initialize the pool manager with our custom RandomizedPoolManager.
+        """
+        pool_kwargs.update({"num_pools": self.num_pools, "maxsize": self.custom_maxsize, "block": block})
+        self.poolmanager = RandomizedPoolManager(**pool_kwargs)
+
+
+# Create a requests session.
+dirax_session = requests.Session()
+
+# Create an instance of the custom adapter.
+# For example, 5 sub-pools per host, each allowing up to 10 connections.
+diracx_pool_adapter = RandomizedHTTPAdapter(num_pools=10, maxsize=10)
+
+# Mount the adapter to handle both HTTP and HTTPS.
+dirax_session.mount("http://", diracx_pool_adapter)
+dirax_session.mount("https://", diracx_pool_adapter)
+
+
 def get_token(
     username: str, group: str, dirac_properties: set[str], *, expires_minutes: int | None = None, source: str = ""
 ):
@@ -42,7 +107,7 @@ def get_token(
     vo = Registry.getVOForGroup(group)
     scopes = [f"vo:{vo}", f"group:{group}"] + [f"property:{prop}" for prop in dirac_properties]
 
-    r = requests.get(
+    r = dirax_session.get(
         f"{diracxUrl}/api/auth/legacy-exchange",
         params={
             "preferred_username": username,