Fix draining plugin db pool

This commit is contained in:
Tulir Asokan 2022-03-27 19:35:54 +03:00
parent 18fa1b29a7
commit 1aaa8f2fb6

View File

@ -32,6 +32,7 @@ class ProxyPostgresDatabase(Database):
_quoted_schema: str _quoted_schema: str
_default_search_path: str _default_search_path: str
_conn_sema: asyncio.Semaphore _conn_sema: asyncio.Semaphore
_max_conns: int
def __init__( def __init__(
self, self,
@ -49,6 +50,7 @@ class ProxyPostgresDatabase(Database):
self._quoted_schema = f'"{self.schema_name}"' self._quoted_schema = f'"{self.schema_name}"'
self._default_search_path = '"$user", public' self._default_search_path = '"$user", public'
self._conn_sema = asyncio.BoundedSemaphore(max_conns) self._conn_sema = asyncio.BoundedSemaphore(max_conns)
self._max_conns = max_conns
async def start(self) -> None: async def start(self) -> None:
async with self._underlying_pool.acquire() as conn: async with self._underlying_pool.acquire() as conn:
@ -58,7 +60,7 @@ class ProxyPostgresDatabase(Database):
await super().start() await super().start()
async def stop(self) -> None: async def stop(self) -> None:
while not self._conn_sema.locked(): for _ in range(self._max_conns):
try: try:
await asyncio.wait_for(self._conn_sema.acquire(), timeout=3) await asyncio.wait_for(self._conn_sema.acquire(), timeout=3)
except asyncio.TimeoutError: except asyncio.TimeoutError: