Skip to content

Commit 4bb2566

Browse files
committed
review fixes
1 parent eb98ede commit 4bb2566

File tree

2 files changed

+8
-3
lines changed

2 files changed

+8
-3
lines changed

ydb/aio/query/pool.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import asyncio
22
import logging
3-
import functools
43
from typing import (
54
Callable,
65
Optional,
@@ -61,7 +60,8 @@ async def acquire(self) -> QuerySession:
6160
done, _ = await asyncio.wait((queue_get, task_stop), return_when=asyncio.FIRST_COMPLETED)
6261
if task_stop in done:
6362
queue_get.cancel()
64-
return await self._create_new_session()
63+
raise RuntimeError("An attempt to take session from closed session pool.")
64+
6565
task_stop.cancel()
6666
session = queue_get.result()
6767

@@ -163,7 +163,7 @@ def __del__(self):
163163
if self._should_stop.is_set() or self._loop.is_closed():
164164
return
165165

166-
self._loop.call_soon(functools.partial(self.stop))
166+
self._loop.call_soon(self.stop)
167167

168168

169169
class SimpleQuerySessionCheckoutAsync:

ydb/query/pool.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ def _create_new_session(self, timeout: Optional[float]):
4747
return session
4848

4949
def acquire(self, timeout: Optional[float] = None) -> QuerySession:
50+
start = time.monotonic()
51+
5052
lock_acquire_timeout = timeout if timeout is not None else -1
5153
acquired = self._lock.acquire(timeout=lock_acquire_timeout)
5254
try:
@@ -60,6 +62,9 @@ def acquire(self, timeout: Optional[float] = None) -> QuerySession:
6062
except queue.Empty:
6163
pass
6264

65+
finish = time.monotonic()
66+
timeout = timeout - (finish - start) if timeout is not None else None
67+
6368
start = time.monotonic()
6469
if session is None and self._current_size == self._size:
6570
try:

0 commit comments

Comments
 (0)