refactor: Centralize HTTP session management with shared connection pool and token caching

Author: grafuls

src/badfish/helpers/http_client.py

@@ -1,11 +1,11 @@
-import asyncio
 import json
 from typing import Any, Dict, Optional
 
 import aiohttp
 
 from src.badfish.helpers.async_lru import alru_cache
 from src.badfish.helpers.exceptions import BadfishException
+from src.badfish.helpers.session_manager import SessionManager
 
 
 class HTTPClient:
@@ -19,10 +19,15 @@ def __init__(self, host: str, username: str, password: str, logger, retries: int
         self.host_uri = f"https://{host}"
         self.redfish_uri = "/redfish/v1"
         self.root_uri = f"{self.host_uri}{self.redfish_uri}"
-        self.semaphore = asyncio.Semaphore(50)
+        self.session_manager: Optional[SessionManager] = None
         self.token = None
         self.session_id = None
 
+    async def _ensure_session_manager(self):
+        """Lazy initialization of session manager."""
+        if self.session_manager is None:
+            self.session_manager = await SessionManager.get_instance()
+
     async def error_handler(self, response: aiohttp.ClientResponse, message: Optional[str] = None) -> None:
         try:
             raw = await response.text("utf-8", "ignore")
@@ -47,13 +52,13 @@ async def error_handler(self, response: aiohttp.ClientResponse, message: Optiona
     @alru_cache(maxsize=64)
     async def get_request(self, uri: str, _continue: bool = False, _get_token: bool = False):
         return await self.get_raw(uri, _continue, _get_token)
-    
+
     @alru_cache(maxsize=64)
     async def get_json(self, uri: str, _continue: bool = False, _get_token: bool = False):
         response = await self.get_raw(uri, _continue, _get_token)
         if not response:
             return None
-            
+
         # Parse JSON from response
         try:
             raw = await response.text("utf-8", "ignore")
@@ -64,28 +69,30 @@ async def get_json(self, uri: str, _continue: bool = False, _get_token: bool = F
             return None
 
     async def get_raw(self, uri: str, _continue: bool = False, _get_token: bool = False):
+        await self._ensure_session_manager()
+
+        # Check token cache first (unless we're getting a new token)
+        if not _get_token and not self.token:
+            cached_token = self.session_manager.get_token(self.host)
+            if cached_token:
+                self.token = cached_token
+                cached_session_id = self.session_manager.get_session_id(self.host)
+                if cached_session_id:
+                    self.session_id = cached_session_id
+
+        headers = {}
+        auth = None
+
+        if _get_token:
+            auth = aiohttp.BasicAuth(self.username, self.password)
+        elif self.token:
+            headers["X-Auth-Token"] = self.token
+
         try:
-            async with self.semaphore:
-                async with aiohttp.ClientSession() as session:
-                    if not _get_token:
-                        async with session.get(
-                            uri,
-                            headers={"X-Auth-Token": self.token} if self.token else {},
-                            ssl=False,
-                            timeout=60,
-                        ) as _response:
-                            await _response.read()
-                    else:
-                        async with session.get(
-                            uri,
-                            auth=aiohttp.BasicAuth(self.username, self.password),
-                            ssl=False,
-                            timeout=60,
-                        ) as _response:
-                            await _response.read()
+            _response = await self.session_manager.get_request(uri, headers, auth)
         except (Exception, TimeoutError) as ex:
             if _continue:
-                return
+                return None
             else:
                 self.logger.debug(f"HTTPClient get_raw exception: {ex}")
                 self.logger.debug(f"Exception type: {type(ex)}")
@@ -100,39 +107,28 @@ async def post_request(
         headers: Dict[str, str],
         _get_token: bool = False,
     ):
+        await self._ensure_session_manager()
+
+        if not _get_token and self.token:
+            headers = headers.copy()  # Avoid mutating caller's dict
+            headers.update({"X-Auth-Token": self.token})
+
         try:
-            async with self.semaphore:
-                async with aiohttp.ClientSession() as session:
-                    if not _get_token and self.token:
-                        headers.update({"X-Auth-Token": self.token})
-                    async with session.post(
-                        uri,
-                        data=json.dumps(payload),
-                        headers=headers,
-                        ssl=False,
-                    ) as _response:
-                        if _response.status != 204:
-                            await _response.read()
-                        else:
-                            return _response
+            _response = await self.session_manager.post_request(uri, payload, headers)
         except (Exception, TimeoutError):
             raise BadfishException("Failed to communicate with server.")
         return _response
 
     async def patch_request(self, uri: str, payload: Dict[str, Any], headers: Dict[str, str], _continue: bool = False):
+        await self._ensure_session_manager()
+
+        if self.token:
+            headers = headers.copy()  # Avoid mutating caller's dict
+            headers.update({"X-Auth-Token": self.token})
+
         try:
-            async with self.semaphore:
-                async with aiohttp.ClientSession() as session:
-                    if self.token:
-                        headers.update({"X-Auth-Token": self.token})
-                    async with session.patch(
-                        uri,
-                        data=json.dumps(payload),
-                        headers=headers,
-                        ssl=False,
-                    ) as _response:
-                        raw_data = await _response.read()
-                        return _response
+            _response = await self.session_manager.patch_request(uri, payload, headers)
+            return _response
         except Exception as ex:
             if _continue:
                 return None
@@ -141,18 +137,15 @@ async def patch_request(self, uri: str, payload: Dict[str, Any], headers: Dict[s
                 raise BadfishException("Failed to communicate with server.")
 
     async def delete_request(self, uri: str, headers: Dict[str, str]):
+        await self._ensure_session_manager()
+
+        if self.token:
+            headers = headers.copy()  # Avoid mutating caller's dict
+            headers.update({"X-Auth-Token": self.token})
+
         try:
-            async with self.semaphore:
-                async with aiohttp.ClientSession() as session:
-                    if self.token:
-                        headers.update({"X-Auth-Token": self.token})
-                    async with session.delete(
-                        uri,
-                        headers=headers,
-                        ssl=False,
-                    ) as _response:
-                        raw_data = await _response.read()
-                        return _response
+            _response = await self.session_manager.delete_request(uri, headers)
+            return _response
         except (Exception, TimeoutError):
             raise BadfishException("Failed to communicate with server.")
 
@@ -183,6 +176,8 @@ async def find_session_uri(self):
         return session_uri
 
     async def validate_credentials(self):
+        await self._ensure_session_manager()
+
         payload = {"UserName": self.username, "Password": self.password}
         headers = {"content-type": "application/json"}
         session_uri = await self.find_session_uri()
@@ -200,6 +195,10 @@ async def validate_credentials(self):
 
         self.session_id = _response.headers.get("Location")
         token = _response.headers.get("X-Auth-Token")
+
+        # Cache token in session manager
+        self.session_manager.cache_token(self.host, token, self.session_id, ttl_seconds=3600)
+
         return token
 
     async def delete_session(self):
@@ -223,8 +222,13 @@ async def delete_session(self):
                 except Exception as ex:
                     self.logger.warning(f"Failed to delete session for {self.host}: {ex}")
             finally:
+                # Invalidate cached token in session manager
+                if self.session_manager:
+                    self.session_manager.invalidate_token(self.host)
                 self.session_id = None
                 self.token = None
         except Exception:
+            if self.session_manager:
+                self.session_manager.invalidate_token(self.host)
             self.session_id = None
             self.token = None

src/badfish/helpers/parser.py

@@ -1,6 +1,6 @@
 import argparse
 
-from badfish.config import RETRIES
+from src.badfish.config import RETRIES
 
 
 def create_parser():