review fixes

Author: vgvoleg

tests/aio/query/test_query_session_pool.py

@@ -68,13 +68,14 @@ async def test_pool_size_limit_logic(self, pool: QuerySessionPoolAsync):
             ids.add(session._state.session_id)
 
         with pytest.raises(asyncio.TimeoutError):
-            await asyncio.wait_for(pool.acquire(), timeout=0.5)
+            await asyncio.wait_for(pool.acquire(), timeout=0.1)
 
+        last_id = session._state.session_id
         await pool.release(session)
 
         session = await pool.acquire()
+        assert session._state.session_id == last_id
         assert pool._current_size == target_size
-        assert session._state.session_id in ids
 
     @pytest.mark.asyncio
     async def test_checkout_do_not_increase_size(self, pool: QuerySessionPoolAsync):
@@ -106,14 +107,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 asyncio.wait_for(pool.acquire(), timeout=0.5)
+            await asyncio.wait_for(pool.acquire(), timeout=0.1)
 
     @pytest.mark.asyncio
     async def test_no_session_leak(self, driver, docker_project):
         pool = ydb.aio.QuerySessionPoolAsync(driver, 1)
         docker_project.stop()
         try:
-            await asyncio.wait_for(pool.acquire(), timeout=0.5)
+            await asyncio.wait_for(pool.acquire(), timeout=0.1)
         except ydb.Error:
             pass
         assert pool._current_size == 0

tests/query/test_query_session_pool.py

@@ -52,19 +52,20 @@ def test_pool_size_limit_logic(self, pool: QuerySessionPool):
         ids = set()
 
         for i in range(1, target_size + 1):
-            session = pool.acquire(timeout=0.5)
+            session = pool.acquire(timeout=0.1)
             assert pool._current_size == i
             assert session._state.session_id not in ids
             ids.add(session._state.session_id)
 
         with pytest.raises(ydb.SessionPoolEmpty):
-            pool.acquire(timeout=0.5)
+            pool.acquire(timeout=0.1)
 
+        last_id = session._state.session_id
         pool.release(session)
 
-        session = pool.acquire(timeout=0.5)
+        session = pool.acquire(timeout=0.1)
+        assert session._state.session_id == last_id
         assert pool._current_size == target_size
-        assert session._state.session_id in ids
 
     def test_checkout_do_not_increase_size(self, pool: QuerySessionPool):
         session_id = None
@@ -93,7 +94,7 @@ def test_no_session_leak(self, driver_sync, docker_project):
         pool = ydb.QuerySessionPool(driver_sync, 1)
         docker_project.stop()
         try:
-            pool.acquire(timeout=0.5)
+            pool.acquire(timeout=0.1)
         except ydb.Error:
             pass
         assert pool._current_size == 0

ydb/aio/query/pool.py

@@ -32,7 +32,7 @@ def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
         self._driver = driver
         self._size = size
         self._should_stop = asyncio.Event()
-        self._queue = asyncio.PriorityQueue()
+        self._queue = asyncio.Queue()
         self._current_size = 0
         self._waiters = 0
 
@@ -49,7 +49,7 @@ async def acquire(self) -> QuerySessionAsync:
 
         session = None
         try:
-            _, session = self._queue.get_nowait()
+            session = self._queue.get_nowait()
         except asyncio.QueueEmpty:
             pass
 
@@ -59,8 +59,9 @@ async def acquire(self) -> QuerySessionAsync:
             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()
+                return await self._create_new_session()
+            task_stop.cancel()
+            session = queue_get.result()
 
         if session is not None:
             if session._state.attached:
@@ -76,15 +77,15 @@ async def acquire(self) -> QuerySessionAsync:
         return session
 
     async def release(self, session: QuerySessionAsync) -> None:
-        self._queue.put_nowait((1, session))
+        self._queue.put_nowait(session)
         logger.debug("Session returned to queue: %s", session._state.session_id)
 
-    def checkout(self, timeout: Optional[float] = None) -> "SimpleQuerySessionCheckoutAsync":
+    def checkout(self) -> "SimpleQuerySessionCheckoutAsync":
         """WARNING: This API is experimental and could be changed.
         Return a Session context manager, that opens session on enter and closes session on exit.
         """
 
-        return SimpleQuerySessionCheckoutAsync(self, timeout)
+        return SimpleQuerySessionCheckoutAsync(self)
 
     async def retry_operation_async(
         self, callee: Callable, retry_settings: Optional[RetrySettings] = None, *args, **kwargs
@@ -135,13 +136,13 @@ async def wrapped_callee():
 
         return await retry_operation_async(wrapped_callee, retry_settings)
 
-    async def stop(self, timeout=None):
+    async def stop(self):
         self._should_stop.set()
 
         tasks = []
         while True:
             try:
-                _, session = self._queue.get_nowait()
+                session = self._queue.get_nowait()
                 tasks.append(session.delete())
             except asyncio.QueueEmpty:
                 break
@@ -158,16 +159,12 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
 
 class SimpleQuerySessionCheckoutAsync:
-    def __init__(self, pool: QuerySessionPoolAsync, timeout: Optional[float] = None):
+    def __init__(self, pool: QuerySessionPoolAsync):
         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_with_timeout(self._timeout)
-        else:
-            self._session = await self._pool.acquire()
+        self._session = await self._pool.acquire()
         return self._session
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):

ydb/query/pool.py

@@ -34,7 +34,7 @@ def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
 
         logger.warning("QuerySessionPool is an experimental API, which could be changed.")
         self._driver = driver
-        self._queue = queue.PriorityQueue()
+        self._queue = queue.Queue()
         self._current_size = 0
         self._size = size
         self._should_stop = threading.Event()
@@ -54,14 +54,14 @@ def acquire(self, timeout: float) -> QuerySessionSync:
 
             session = None
             try:
-                _, session = self._queue.get_nowait()
+                session = self._queue.get_nowait()
             except queue.Empty:
                 pass
 
             start = time.monotonic()
             if session is None and self._current_size == self._size:
                 try:
-                    _, session = self._queue.get(block=True, timeout=timeout)
+                    session = self._queue.get(block=True, timeout=timeout)
                 except queue.Empty:
                     raise issues.SessionPoolEmpty("Timeout on acquire session")
 
@@ -83,7 +83,7 @@ def acquire(self, timeout: float) -> QuerySessionSync:
 
     def release(self, session: QuerySessionSync) -> None:
         with self._lock:
-            self._queue.put_nowait((1, session))
+            self._queue.put_nowait(session)
             logger.debug("Session returned to queue: %s", session._state.session_id)
 
     def checkout(self, timeout: float = 10) -> "SimpleQuerySessionCheckout":
@@ -145,7 +145,7 @@ def stop(self, timeout=None):
             self._should_stop.set()
             while True:
                 try:
-                    _, session = self._queue.get_nowait()
+                    session = self._queue.get_nowait()
                     session.delete()
                 except queue.Empty:
                     break