safe shutdown client

aiosonic/__init__.py

@@ -437,6 +437,17 @@ def __init__(self, connector: TCPConnector = None):
         """
         self.connector = connector or TCPConnector()
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *_args, **_kwargs):
+        await self.shutdown()
+
+    async def shutdown(self):
+        """Cleanup connections, this method makes client unusable."""
+        await self.connector.cleanup()
+
+
     async def _request_with_body(self,
                                  url: str,
                                  method: str,

aiosonic/connection.py

@@ -72,8 +72,7 @@ def is_closing():
                 return True  # noqa
 
         if not (self.key and key == self.key and not is_closing()):
-            if self.writer:
-                self.writer.close()
+            self.close()
 
             if urlparsed.scheme == 'https':
                 ssl_context = ssl_context or ssl.create_default_context(
@@ -125,7 +124,7 @@ async def __aexit__(self, exc_type, exc, tb):
             self.key = None
             self.h2conn = None
             if self.writer:
-                self.writer.close()
+                self.close()
 
         if not self.blocked:
             await self.release()
@@ -142,10 +141,14 @@ def timeouts(self) -> Timeouts:
 
     def __del__(self):
         """Cleanup."""
+        self.close(True)
+
+    def close(self, check_closing=False):
+        """Close connection if opened."""
         if self.writer:
             is_closing = getattr(self.writer, 'is_closing',
                                  self.writer._transport.is_closing)
-            if is_closing():
+            if not check_closing or is_closing():
                 self.writer.close()
 
     async def http2_request(self, headers: Dict[str, str],

aiosonic/connectors.py

@@ -59,3 +59,7 @@ async def wait_free_pool(self):
             if self.pool.is_all_free():
                 return True
             asyncio_sleep(0.02)
+
+    async def cleanup(self):
+        """Cleanup connector connections."""
+        await self.pool.cleanup()

aiosonic/pools.py

@@ -25,6 +25,12 @@ def is_all_free(self):
         """Indicates if all pool is free."""
         return self.pool_size == self.pool.qsize()
 
+    async def cleanup(self):
+        """Get all conn and close them, this method let this pool unusable."""
+        for _ in range(self.pool_size):
+            conn = self.pool.get()
+            conn.close()
+
 
 class SmartPool:
     """Pool which utilizes alive connections."""
@@ -55,3 +61,9 @@ def release(self, conn):
     def is_all_free(self):
         """Indicates if all pool is free."""
         return self.pool_size == self.sem._value
+
+    async def cleanup(self):
+        """Get all conn and close them, this method let this pool unusable."""
+        for count in range(self.pool_size):
+            conn = await self.acquire()
+            conn.close()