aio query session pool redesign

Author: vgvoleg

tests/aio/query/test_query_session_pool.py

@@ -1,3 +1,4 @@
+import asyncio
 import pytest
 import ydb
 from ydb.aio.query.pool import QuerySessionPoolAsync
@@ -61,17 +62,17 @@ async def test_pool_size_limit_logic(self, pool: QuerySessionPoolAsync):
         ids = set()
 
         for i in range(1, target_size + 1):
-            session = await pool.acquire_wih_timeout(timeout=0.5)
+            session = await pool.acquire()
             assert pool._current_size == i
             assert session._state.session_id not in ids
             ids.add(session._state.session_id)
 
-        with pytest.raises(ydb.SessionPoolEmpty):
-            await pool.acquire_wih_timeout(timeout=0.5)
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(pool.acquire(), timeout=0.5)
 
         await pool.release(session)
 
-        session = await pool.acquire_wih_timeout(timeout=0.5)
+        session = await pool.acquire()
         assert pool._current_size == target_size
         assert session._state.session_id in ids
 
@@ -105,14 +106,14 @@ async def test_acquire_from_closed_pool_raises(self, pool: QuerySessionPoolAsync
     async def test_acquire_with_timeout_from_closed_pool_raises(self, pool: QuerySessionPoolAsync):
         await pool.stop()
         with pytest.raises(RuntimeError):
-            await pool.acquire_wih_timeout(timeout=0.5)
+            await asyncio.wait_for(pool.acquire(), timeout=0.5)
 
     @pytest.mark.asyncio
     async def test_no_session_leak(self, driver, docker_project):
         pool = ydb.aio.QuerySessionPoolAsync(driver, 1)
         docker_project.stop()
         try:
-            await pool.acquire_wih_timeout(timeout=0.5)
+            await asyncio.wait_for(pool.acquire(), timeout=0.5)
         except ydb.Error:
             pass
         assert pool._current_size == 0

ydb/aio/query/pool.py

@@ -55,7 +55,13 @@ async def acquire(self) -> QuerySessionAsync:
             pass
 
         if session is None and self._current_size == self._size:
-            _, session = await self._queue.get()
+            queue_get = asyncio.ensure_future(self._queue.get())
+            task_stop = asyncio.ensure_future(asyncio.ensure_future(self._should_stop.wait()))
+            done, _ = await asyncio.wait((queue_get, task_stop), return_when=asyncio.FIRST_COMPLETED)
+            if task_stop in done:
+                queue_get.cancel()
+                return await self._create_new_session()  # TODO: not sure why
+            _, session = queue_get.result()
 
         if session is not None:
             if session._state.attached:
@@ -70,23 +76,6 @@ async def acquire(self) -> QuerySessionAsync:
         self._current_size += 1
         return session
 
-    async def acquire_wih_timeout(self, timeout: float):
-        if self._should_stop.is_set():
-            logger.error("An attempt to take session from closed session pool.")
-            raise RuntimeError("An attempt to take session from closed session pool.")
-
-        try:
-            task_wait = asyncio.ensure_future(asyncio.wait_for(self.acquire(), timeout=timeout))
-            task_stop = asyncio.ensure_future(asyncio.ensure_future(self._should_stop.wait()))
-            done, _ = await asyncio.wait((task_wait, task_stop), return_when=asyncio.FIRST_COMPLETED)
-            if task_stop in done:
-                task_wait.cancel()
-                return await self._create_new_session()  # TODO: not sure why
-            session = task_wait.result()
-            return session
-        except asyncio.TimeoutError:
-            raise issues.SessionPoolEmpty("Timeout on acquire session")
-
     async def release(self, session: QuerySessionAsync) -> None:
         self._queue.put_nowait((1, session))
         logger.debug("Session returned to queue: %s", session._state.session_id)
@@ -170,14 +159,14 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
 
 class SimpleQuerySessionCheckoutAsync:
-    def __init__(self, pool: QuerySessionPoolAsync, timeout: Optional[float]):
+    def __init__(self, pool: QuerySessionPoolAsync, timeout: Optional[float] = None):
         self._pool = pool
         self._timeout = timeout
         self._session = None
 
     async def __aenter__(self) -> QuerySessionAsync:
         if self._timeout and self._timeout > 0:
-            self._session = await self._pool.acquire_wih_timeout(self._timeout)
+            self._session = await self._pool.acquire_with_timeout(self._timeout)
         else:
             self._session = await self._pool.acquire()
         return self._session

ydb/query/session.py

@@ -1,7 +1,6 @@
 import abc
 import enum
 import logging
-import time
 import threading
 from typing import (
     Iterable,