Update app.py

app.py

@@ -25,87 +25,72 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 # --- Enhanced Session Management ---
+import asyncio
+import time
+from datetime import datetime
+import aiohttp
+import logging
+
+logger = logging.getLogger(__name__)
+
 class SessionManager:
-    """Thread-safe session manager with proper event loop handling."""
+    """
+    Thread-safe session manager with graceful rotation to prevent race conditions.
+    """
     
     def __init__(self):
-        self._session = None
-        self._session_lock = None  # Will be created per event loop
+        self._session: aiohttp.ClientSession = None
+        self._session_lock = asyncio.Lock() # This lock is tied to the event loop it's created in.
         self._creation_time = None
         self._request_count = 0
-        self._loop_id = None  # Track which event loop owns the session
         self.max_session_age = 300  # 5 minutes
         self.max_requests_per_session = 1000
         
-    def _get_or_create_lock(self):
-        """Get or create lock for current event loop."""
-        try:
-            current_loop = asyncio.get_running_loop()
-            current_loop_id = id(current_loop)
-            
-            # If we're in a different event loop, reset session
-            if self._loop_id and self._loop_id != current_loop_id:
-                logger.info(f"🔄 Event loop changed, resetting session")
-                self._session = None
-                self._session_lock = None
-                self._creation_time = None
-                self._request_count = 0
-            
-            self._loop_id = current_loop_id
-            
-            if self._session_lock is None:
-                self._session_lock = asyncio.Lock()
-                
-            return self._session_lock
-        except RuntimeError:
-            # No event loop running
-            return None
-        
+    async def _graceful_close(self, session_to_close: aiohttp.ClientSession, delay: int = 5):
+        """Waits for a delay before closing a stale session to allow in-flight requests to complete."""
+        if session_to_close and not session_to_close.closed:
+            logger.info(f"⏳ Waiting {delay}s before closing stale session (ID: {id(session_to_close)})...")
+            await asyncio.sleep(delay)
+            try
+                logger.info(f"🧹 Gracefully closing stale session (ID: {id(session_to_close)}).")
+                await session_to_close.close()
+            except Exception as e:
+                logger.warning(f"⚠️ Error during graceful close of stale session: {e}")
+
     async def get_session(self) -> aiohttp.ClientSession:
-        """Get or create a session with proper event loop handling."""
-        lock = self._get_or_create_lock()
-        if lock is None:
-            raise RuntimeError("No event loop is running")
-            
-        async with lock:
+        """
+        Get or create a session with graceful rotation to prevent race conditions.
+        """
+        async with self._session_lock:
             now = time.time()
             
-            # Check if session needs renewal
+            # Check if session needs renewal. Use >= for precision on max requests.
             needs_renewal = (
                 self._session is None or 
                 self._session.closed or
                 (self._creation_time and now - self._creation_time > self.max_session_age) or
-                self._request_count > self.max_requests_per_session
+                self._request_count >= self.max_requests_per_session
             )
             
             if needs_renewal:
-                # Close old session if exists
-                if self._session and not self._session.closed:
-                    logger.info(f"🔄 Closing old session after {self._request_count} requests")
-                    try:
-                        await self._session.close()
-                        await asyncio.sleep(0.1)  # Allow cleanup
-                    except Exception as e:
-                        logger.warning(f"Error closing old session: {e}")
+                # Keep a reference to the old session
+                old_session = self._session
+                
+                # *** THE FIX: Schedule the old session's closure instead of awaiting it. ***
+                if old_session and not old_session.closed:
+                    logger.info(f"🔄 Scheduling closure of old session after {self._request_count} requests.")
+                    loop = asyncio.get_running_loop()
+                    # This creates a background task that will run independently.
+                    loop.create_task(self._graceful_close(old_session))
                 
-                # Create new optimized session
+                # --- Create the new session immediately ---
                 try:
                     connector = aiohttp.TCPConnector(
-                        limit=100,
-                        limit_per_host=50,
-                        ttl_dns_cache=300,
-                        use_dns_cache=True,
-                        keepalive_timeout=30,
+                        limit=100, limit_per_host=50, ttl_dns_cache=300,
+                        use_dns_cache=True, keepalive_timeout=30,
                         enable_cleanup_closed=True,
-                        force_close=False,
                     )
-                    
-                    timeout = aiohttp.ClientTimeout(
-                        total=60,
-                        connect=10,
-                        sock_read=30
-                    )
-                    
+                    timeout = aiohttp.ClientTimeout(total=60, connect=10, sock_read=30)
                     headers = {
                         'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
                         'Accept': 'application/json',
@@ -115,38 +100,35 @@ class SessionManager:
                     }
                     
                     self._session = aiohttp.ClientSession(
-                        connector=connector,
-                        timeout=timeout,
-                        headers=headers,
-                        raise_for_status=False
+                        connector=connector, timeout=timeout, headers=headers
                     )
                     
                     self._creation_time = now
                     self._request_count = 0
-                    logger.info(f"🚀 Created new session at {datetime.fromtimestamp(now)}")
+                    logger.info(f"🚀 Created new session (ID: {id(self._session)}) at {datetime.fromtimestamp(now)}")
                     
                 except Exception as e:
-                    logger.error(f"Failed to create session: {e}")
+                    logger.error(f"Failed to create new session: {e}")
+                    # If creation fails, restore the old session if it's still valid
+                    if old_session and not old_session.closed:
+                        self._session = old_session
                     raise
             
             return self._session
     
     def increment_request_count(self):
-        """Thread-safe request count increment."""
+        """Increment the request count. No lock needed as it's called after getting the session."""
         self._request_count += 1
     
     async def close(self):
-        """Clean shutdown of session."""
+        """Clean shutdown of the current active session."""
         if self._session and not self._session.closed:
             try:
-                logger.info(f"🔒 Closing session after {self._request_count} requests")
+                logger.info(f"🔒 Closing active session after {self._request_count} requests.")
                 await self._session.close()
-                await asyncio.sleep(0.1)
             except Exception as e:
-                logger.warning(f"Error during session cleanup: {e}")
+                logger.warning(f"Error during final session cleanup: {e}")
         self._session = None
-        self._session_lock = None
-
 # --- Enhanced Data Models ---
 
 class StockDataRequest(BaseModel):
@@ -252,21 +234,6 @@ def get_time_range_in_unix_ms(start_date_str: str, end_date_str: str, timezone:
     return {"start_timestamp_ms": start_ts, "end_timestamp_ms": end_ts}
 
 
-# Global session manager per event loop to handle multiple threads/loops
-_session_managers = weakref.WeakKeyDictionary()
-
-def get_session_manager():
-    """Get or create session manager for current event loop."""
-    try:
-        loop = asyncio.get_running_loop()
-        if loop not in _session_managers:
-            _session_managers[loop] = SessionManager()
-            logger.debug(f"📝 Created new SessionManager for event loop {id(loop)}")
-        return _session_managers[loop]
-    except RuntimeError:
-        # No event loop running
-        logger.warning("⚠️ No event loop running")
-        return None
 
 # --- Optimized API Functions ---