review fixes

Author: vgvoleg

ydb/aio/query/pool.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import functools
 from typing import (
     Callable,
     Optional,
@@ -35,6 +36,7 @@ def __init__(self, driver: common_utils.SupportedDriverType, size: int = 100):
         self._queue = asyncio.Queue()
         self._current_size = 0
         self._waiters = 0
+        self._loop = asyncio.get_running_loop()
 
     async def _create_new_session(self):
         session = QuerySessionAsync(self._driver)
@@ -157,6 +159,12 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.stop()
 
+    def __del__(self):
+        if self._should_stop.is_set() or self._loop.is_closed():
+            return
+
+        self._loop.call_soon(functools.partial(self.stop))
+
 
 class SimpleQuerySessionCheckoutAsync:
     def __init__(self, pool: QuerySessionPoolAsync):

ydb/query/pool.py

@@ -47,7 +47,8 @@ def _create_new_session(self, timeout: float):
         return session
 
     def acquire(self, timeout: float) -> QuerySessionSync:
-        with self._lock:
+        acquired = self._lock.acquire(timeout=timeout)
+        try:
             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.")
@@ -80,11 +81,13 @@ def acquire(self, timeout: float) -> QuerySessionSync:
 
             self._current_size += 1
             return session
+        finally:
+            if acquired:
+                self._lock.release()
 
     def release(self, session: QuerySessionSync) -> None:
-        with self._lock:
-            self._queue.put_nowait(session)
-            logger.debug("Session returned to queue: %s", session._state.session_id)
+        self._queue.put_nowait(session)
+        logger.debug("Session returned to queue: %s", session._state.session_id)
 
     def checkout(self, timeout: float = 10) -> "SimpleQuerySessionCheckout":
         """WARNING: This API is experimental and could be changed.
@@ -140,8 +143,9 @@ def wrapped_callee():
 
         return retry_operation_sync(wrapped_callee, retry_settings)
 
-    def stop(self, timeout=None):
-        with self._lock:
+    def stop(self, timeout=-1):
+        acquired = self._lock.acquire(timeout=timeout)
+        try:
             self._should_stop.set()
             while True:
                 try:
@@ -151,13 +155,19 @@ def stop(self, timeout=None):
                     break
 
             logger.debug("All session were deleted.")
+        finally:
+            if acquired:
+                self._lock.release()
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.stop()
 
+    def __del__(self):
+        self.stop()
+
 
 class SimpleQuerySessionCheckout:
     def __init__(self, pool: QuerySessionPool, timeout: float):