sync query session pool

Author: vgvoleg

tests/query/test_query_session_pool.py

@@ -9,8 +9,6 @@ def test_checkout_provides_created_session(self, pool: QuerySessionPool):
         with pool.checkout() as session:
             assert session._state._state == QuerySessionStateEnum.CREATED
 
-        assert session._state._state == QuerySessionStateEnum.CLOSED
-
     def test_oneshot_query_normal(self, pool: QuerySessionPool):
         res = pool.execute_with_retries("select 1;")
         assert len(res) == 1
@@ -47,3 +45,58 @@ def callee(session: QuerySessionSync):
 
         with pytest.raises(CustomException):
             pool.retry_operation_sync(callee)
+
+    def test_pool_size_limit_logic(self, pool: QuerySessionPool):
+        target_size = 5
+        pool._size = target_size
+        ids = set()
+
+        for i in range(1, target_size + 1):
+            session = pool.acquire(timeout=0.5)
+            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.release(session)
+
+        session = pool.acquire(timeout=0.5)
+        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
+        for _ in range(10):
+            with pool.checkout() as session:
+                if session_id is None:
+                    session_id = session._state.session_id
+                assert pool._current_size == 1
+                assert session_id == session._state.session_id
+
+    def test_pool_recreates_bad_sessions(self, pool: QuerySessionPool):
+        with pool.checkout() as session:
+            session_id = session._state.session_id
+            session.delete()
+
+        with pool.checkout() as session:
+            assert session_id != session._state.session_id
+            assert pool._current_size == 1
+
+    def test_acquire_from_closed_pool_raises(self, pool: QuerySessionPool):
+        pool.stop()
+        with pytest.raises(RuntimeError):
+            pool.acquire(1)
+
+    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)
+        except ydb.Error:
+            pass
+        assert pool._current_size == 0
+
+        docker_project.start()
+        pool.stop()

ydb/query/pool.py

@@ -4,6 +4,8 @@
     Optional,
     List,
 )
+import threading
+import queue
 
 from .session import (
     QuerySessionSync,
@@ -12,6 +14,7 @@
     RetrySettings,
     retry_operation_sync,
 )
+from .. import issues
 from .. import convert
 from .._grpc.grpcwrapper import common_utils
 
@@ -22,20 +25,61 @@
 class QuerySessionPool:
     """QuerySessionPool is an object to simplify operations with sessions of Query Service."""
 
-    def __init__(self, driver: common_utils.SupportedDriverType):
+    def __init__(self, driver: common_utils.SupportedDriverType, size: int = 10):
         """
         :param driver: A driver instance
         """
 
         logger.warning("QuerySessionPool is an experimental API, which could be changed.")
         self._driver = driver
-
-    def checkout(self) -> "SimpleQuerySessionCheckout":
+        self._queue = queue.PriorityQueue()
+        self._current_size = 0
+        self._size = size
+        self._should_stop = threading.Event()
+        self._lock = threading.RLock()
+
+    def _create_new_session(self):
+        session = QuerySessionSync(self._driver)
+        session.create()
+        logger.debug(f"New session was created for pool. Session id: {session._state.session_id}")
+        return session
+
+    def acquire(self, timeout: float) -> QuerySessionSync:
+        with self._lock:
+            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:
+                _, session = self._queue.get_nowait()
+                logger.debug(f"Acquired active session from queue: {session._state.session_id}")
+                return session if session._state.attached else self._create_new_session()
+            except queue.Empty:
+                pass
+
+            if self._current_size < self._size:
+                logger.debug(f"Session pool is not large enough: {self._current_size} < {self._size}, will create new one.")
+                session = self._create_new_session()
+                self._current_size += 1
+                return session
+
+            try:
+                _, session = self._queue.get(block=True, timeout=timeout)
+                return session if session._state.attached else self._create_new_session()
+            except queue.Empty:
+                raise issues.SessionPoolEmpty("Timeout on acquire session")
+
+    def release(self, session: QuerySessionSync) -> None:
+        with self._lock:
+            self._queue.put_nowait((1, 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.
         Return a Session context manager, that opens session on enter and closes session on exit.
         """
 
-        return SimpleQuerySessionCheckout(self)
+        return SimpleQuerySessionCheckout(self, timeout)
 
     def retry_operation_sync(self, callee: Callable, retry_settings: Optional[RetrySettings] = None, *args, **kwargs):
         """WARNING: This API is experimental and could be changed.
@@ -85,7 +129,17 @@ def wrapped_callee():
         return retry_operation_sync(wrapped_callee, retry_settings)
 
     def stop(self, timeout=None):
-        pass  # TODO: implement
+        with self._lock:
+            self._should_stop.set()
+            while True:
+                try:
+                    _, session = self._queue.get_nowait()
+                    session.delete()
+                except queue.Empty:
+                    break
+
+            logger.debug("All session were deleted.")
+
 
     def __enter__(self):
         return self
@@ -95,13 +149,14 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
 
 class SimpleQuerySessionCheckout:
-    def __init__(self, pool: QuerySessionPool):
+    def __init__(self, pool: QuerySessionPool, timeout: float):
         self._pool = pool
-        self._session = QuerySessionSync(pool._driver)
+        self._timeout = timeout
+        self._session = None
 
     def __enter__(self) -> QuerySessionSync:
-        self._session.create()
+        self._session = self._pool.acquire(self._timeout)
         return self._session
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self._session.delete()
+        self._pool.release(self._session)