Remove connection from pool before waiting for results

Author: vgvoleg

ydb/aio/connection.py

@@ -227,7 +227,7 @@ async def connection_ready(self, ready_timeout=10):
 
         await asyncio.wait_for(self._channel.channel_ready(), timeout=ready_timeout)
 
-    async def close(self, grace: float = None):
+    async def close(self, grace: float = 30):
         """
         Closes the underlying gRPC channel
         :param: grace: If a grace period is specified, this method wait until all active
@@ -239,12 +239,12 @@ async def close(self, grace: float = None):
 
         self.closing = True
 
-        if self.calls:
-            await asyncio.wait(self.calls.values(), timeout=grace)
-
         for callback in self._cleanup_callbacks:
             callback(self)
 
+        if self.calls:
+            await asyncio.wait(self.calls.values(), timeout=grace)
+
         await self.destroy()
 
     async def __aenter__(self):