AioHTTP and Async HTTPX teardown improvments pt.1

Author: Sebastian Molenda

pubnub/pubnub_core.py

@@ -250,7 +250,7 @@ def list_push_channels(self, device_id: str = None, push_type: PNPushType = None
                            environment: PNPushEnvironment = None) -> ListPushProvisions:
         return ListPushProvisions(self, device_id=device_id, push_type=push_type, topic=topic, environment=environment)
 
-    def add_channels_to_push(self, channels: Union[str, List[str]], device_id: str = None, push_type: PNPushType = None,
+    def add_channels_to_push(self, channels: Union[str, List[str]] = None, device_id: str = None, push_type: PNPushType = None,
                              topic: str = None, environment: PNPushEnvironment = None) -> AddChannelsToPush:
         return AddChannelsToPush(self, channels=channels, device_id=device_id, push_type=push_type, topic=topic,
                                  environment=environment)

pubnub/request_handlers/async_aiohttp.py

@@ -28,25 +28,55 @@ class AsyncAiohttpRequestHandler(BaseRequestHandler):
     ENDPOINT_THREAD_COUNTER: int = 0
     _connector: aiohttp.TCPConnector = None
     _session: aiohttp.ClientSession = None
+    _is_closing: bool = False
+    _max_connections: int = 100
+    _connection_timeout: float = 30.0
 
     def __init__(self, pubnub):
         self.pubnub = pubnub
-        self._connector = aiohttp.TCPConnector(verify_ssl=True, loop=self.pubnub.event_loop)
+        self._connector = aiohttp.TCPConnector(
+            verify_ssl=True,
+            loop=self.pubnub.event_loop,
+            limit=self._max_connections,
+            ttl_dns_cache=300
+        )
+        self._is_closing = False
 
     async def create_session(self):
-        if not self._session:
+        if not self._session and not self._is_closing:
             self._session = aiohttp.ClientSession(
                 loop=self.pubnub.event_loop,
-                timeout=aiohttp.ClientTimeout(connect=self.pubnub.config.connect_timeout),
+                timeout=aiohttp.ClientTimeout(
+                    total=self._connection_timeout,
+                    connect=self.pubnub.config.connect_timeout,
+                    sock_read=self.pubnub.config.connect_timeout,
+                    sock_connect=self.pubnub.config.connect_timeout
+                ),
                 connector=self._connector
             )
 
     async def close_session(self):
-        if self._session is not None:
-            await self._session.close()
+        if self._session is not None and not self._is_closing:
+            self._is_closing = True
+            try:
+                # Cancel any pending requests
+                if hasattr(self._session, '_connector'):
+                    for task in asyncio.all_tasks():
+                        if not task.done() and task is not asyncio.current_task():
+                            task.cancel()
+
+                # Close session and connector
+                await self._session.close()
+                await self._connector.close()
+            except Exception as e:
+                logger.error(f"Error during session cleanup: {str(e)}")
+            finally:
+                self._session = None
+                self._is_closing = False
 
     async def set_connector(self, connector):
-        await self._session.aclose()
+        if self._session is not None:
+            await self.close_session()
         self._connector = connector
         await self.create_session()
 

pubnub/request_handlers/async_httpx.py

@@ -24,33 +24,85 @@ def close(self):
         self.is_closed = True
         asyncio.create_task(super().aclose())
 
+    async def aclose(self):
+        self.is_closed = True
+        await super().aclose()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.aclose()
+
 
 class AsyncHttpxRequestHandler(BaseRequestHandler):
     """ PubNub Python SDK asychronous requests handler based on the `httpx` HTTP library. """
     ENDPOINT_THREAD_COUNTER: int = 0
     _connector: httpx.AsyncHTTPTransport = None
     _session: httpx.AsyncClient = None
+    _is_closing: bool = False
+    _max_connections: int = 100
+    _connection_timeout: float = 30.0
 
     def __init__(self, pubnub):
         self.pubnub = pubnub
-        self._connector = PubNubAsyncHTTPTransport(verify=True, http2=True)
+        self._connector = PubNubAsyncHTTPTransport(
+            verify=True,
+            http2=True,
+            limits=httpx.Limits(max_connections=self._max_connections)
+        )
+        self._is_closing = False
 
     async def create_session(self):
-        self._session = httpx.AsyncClient(
-            timeout=httpx.Timeout(self.pubnub.config.connect_timeout),
-            transport=self._connector
-        )
+        if self._session is None and not self._is_closing:
+            self._session = httpx.AsyncClient(
+                timeout=httpx.Timeout(
+                    connect=self._connection_timeout,
+                    read=self.pubnub.config.connect_timeout,
+                    write=self.pubnub.config.connect_timeout,
+                    pool=self._connection_timeout
+                ),
+                transport=self._connector,
+                http2=True
+            )
 
     async def close_session(self):
-        if self._session is not None:
-            self._connector.close()
-            await self._session.aclose()
+        if self._session is not None and not self._is_closing:
+            self._is_closing = True
+            try:
+                # Cancel any pending requests
+                if hasattr(self._session, '_transport'):
+                    for task in asyncio.all_tasks():
+                        if not task.done() and task is not asyncio.current_task():
+                            task.cancel()
+
+                # Close transport and session
+                await self._connector.aclose()
+                await self._session.aclose()
+            except Exception as e:
+                logger.error(f"Error during session cleanup: {str(e)}")
+            finally:
+                self._session = None
+                self._is_closing = False
 
     async def set_connector(self, connector):
-        await self._session.aclose()
+        if self._session is not None:
+            await self.close_session()
         self._connector = connector
         await self.create_session()
 
+    async def __aenter__(self):
+        await self.create_session()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close_session()
+
+    def __del__(self):
+        ...
+        # if self._session is not None and not self._is_closing:
+        #     asyncio.create_task(self.close_session())
+
     def sync_request(self, **_):
         raise NotImplementedError("sync_request is not implemented for asyncio handler")
 
@@ -123,13 +175,17 @@ async def async_request(self, options_func, cancellation_event):
         except Exception as e:
             logger.error("session.request exception: %s" % str(e))
             raise
+        try:
+            response_body = await response.aread()
+        except Exception as e:
+            logger.error(f"Error reading response body: {str(e)}")
+            response_body = None
 
-        response_body = response.read()
         if not options.non_json_response:
             body = response_body
         else:
             if isinstance(response.content, bytes):
-                body = response.content  # TODO: simplify this logic within the v5 release
+                body = response.content
             else:
                 body = response_body
 
@@ -192,7 +248,6 @@ async def async_request(self, options_func, cancellation_event):
         logger.debug(data)
 
         if response.status_code not in (200, 307, 204):
-
             if response.status_code >= 500:
                 err = PNERR_SERVER_ERROR
             else:

pubnub/request_handlers/httpx.py

@@ -28,7 +28,7 @@ class HttpxRequestHandler(BaseRequestHandler):
     ENDPOINT_THREAD_COUNTER: int = 0
 
     def __init__(self, pubnub):
-        self.session = httpx.Client()
+        self.session = httpx.Client(http2=True)
 
         self.pubnub = pubnub