refactor: restore thread safety for aio session and simplify exception handling

Author: kevin

dashscope/api_entities/aio_session.py

@@ -8,6 +8,7 @@
 """
 import asyncio
 import ssl
+import threading
 import weakref
 from typing import Optional
 
@@ -16,14 +17,15 @@
 
 _shared_ssl_context: Optional[ssl.SSLContext] = None
 _aio_sessions: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
-
+_lock = threading.RLock()
 
 def get_ssl_context() -> ssl.SSLContext:
     global _shared_ssl_context
-    if _shared_ssl_context is None:
-        _shared_ssl_context = ssl.create_default_context(
-            cafile=certifi.where(),
-        )
+    with _lock:
+        if _shared_ssl_context is None:
+            _shared_ssl_context = ssl.create_default_context(
+                cafile=certifi.where(),
+            )
     return _shared_ssl_context
 
 
@@ -36,19 +38,21 @@ async def get_shared_aio_session() -> aiohttp.ClientSession:
     """
     loop = asyncio.get_running_loop()
 
-    session = _aio_sessions.get(loop)
-    if session is not None and not session.closed:
-        return session
+    with _lock:
+        session = _aio_sessions.get(loop)
+        if session is not None and not session.closed:
+            return session
 
-    connector = aiohttp.TCPConnector(ssl=get_ssl_context())
-    session = aiohttp.ClientSession(connector=connector)
-    _aio_sessions[loop] = session
+        connector = aiohttp.TCPConnector(ssl=get_ssl_context())
+        session = aiohttp.ClientSession(connector=connector)
+        _aio_sessions[loop] = session
     return session
 
 
 async def close_shared_aio_session() -> None:
     """Close the shared session for the current event loop."""
     loop = asyncio.get_running_loop()
-    session = _aio_sessions.pop(loop, None)
+    with _lock:
+        session = _aio_sessions.pop(loop, None)
     if session is not None and not session.closed:
         await session.close()

dashscope/api_entities/aiohttp_request.py

@@ -297,9 +297,6 @@ async def _handle_request(self):
             async with response:
                 async for rsp in self._handle_response(response):
                     yield rsp
-        except aiohttp.ClientConnectorError as e:
-            logger.error(e)
-            raise e
         except Exception as e:
             logger.error(e)
             raise e

dashscope/api_entities/http_request.py

@@ -218,10 +218,7 @@ async def _handle_aio_request(self):  # pylint: disable=too-many-branches
             async with response:
                 async for rsp in self._handle_aio_response(response):
                     yield rsp
-        except aiohttp.ClientConnectorError as e:
-            logger.error(e)
-            raise e
-        except BaseException as e:
+        except Exception as e:
             logger.error(e)
             raise e
 
@@ -507,6 +504,6 @@ def _handle_request(self):  # pylint: disable=too-many-branches
                 # Only close if we created the session
                 if should_close:
                     session.close()
-        except BaseException as e:
+        except Exception as e:
             logger.error(e)
             raise e