Skip to content

Commit 97f3412

Browse files
vinaysingh8866vinaysingh8866
authored
fix(kanon):updated connection cleanup to share 1 thread and added logging to detect connection leakage (#3963)
* fix(kanon):updated connection cleanup to share 1 thread and added logs for connection leakage Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):refactor wallet and connection pool to use 1 thread Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):removed unecessary callback, added parellel opening to askar and kanon db Signed-off-by: Vinay Singh <vinay@verid.id> * style: applied ruff formatting Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):moved wallet db session to reuse profile session Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):added check to see the correct handler and changed storage to close soon Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):added explicit close to scan Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):added explicit close to scan Signed-off-by: Vinay Singh <vinay@verid.id> * style:applied ruff formating Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):added explicit release Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):test async context manager Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):added detailed logs Signed-off-by: Vinay Singh <vinay@verid.id> * style:applied ruff formating Signed-off-by: Vinay Singh <vinay@verid.id> * fix(kanon):corrected docker params Signed-off-by: Vinay Singh <vinay@verid.id> --------- Signed-off-by: Vinay Singh <vinay@verid.id>
1 parent 84006f6 commit 97f3412

10 files changed

Lines changed: 831 additions & 380 deletions

File tree

‎acapy_agent/database_manager/databases/postgresql_normalized/connection_pool.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -78,12 +78,13 @@ async def initialize(self):
7878
async def getconn(self):
7979
"""Get a connection from the pool."""
8080
try:
81-
async with asyncio.timeout(60.0):
82-
conn = await self.pool.getconn()
81+
conn = await asyncio.wait_for(self.pool.getconn(), timeout=60.0)
8382
# Rollback any existing transaction to ensure clean state
83+
# This ensures the connection is in IDLE state before returning
8484
await conn.rollback()
85-
# Ensure client encoding is set to UTF-8
86-
await conn.execute("SET client_encoding = 'UTF8'")
85+
# Note: UTF-8 encoding is already set at pool creation via kwargs
86+
# (options: -c client_encoding=UTF8), so no need to SET it here.
87+
# Executing SET here would start an implicit transaction and add latency.
8788
conn_id = self.connection_count
8889
self.connection_ids[id(conn)] = conn_id
8990
self.connection_count += 1
@@ -115,14 +116,13 @@ async def putconn(self, conn):
115116
# Roll back any open transactions to ensure clean state
116117
await conn.rollback()
117118
await self.pool.putconn(conn)
118-
conn_id = self.connection_ids.get(id(conn), -1)
119+
conn_id = self.connection_ids.pop(id(conn), -1)
119120
LOGGER.debug(
120121
"Connection ID=%d returned to pool. Pool size: %d/%d",
121122
conn_id,
122123
self.pool.get_stats().get("pool_available", 0),
123124
self.max_size,
124125
)
125-
del self.connection_ids[id(conn)]
126126
except Exception as e:
127127
LOGGER.error("Failed to return connection to pool: %s", str(e))
128128
raise DatabaseError(

‎acapy_agent/database_manager/databases/postgresql_normalized/database.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,6 @@ async def _get_profile_id(self, profile_name: str) -> int:
146146
actual_error=str(e),
147147
)
148148
finally:
149-
await conn.rollback()
150149
await self.pool.putconn(conn)
151150

152151
async def create_profile(self, name: str = None) -> str:
@@ -329,8 +328,12 @@ async def session(self, profile: str = None):
329328
code=DatabaseErrorCode.CONNECTION_POOL_EXHAUSTED,
330329
message="Maximum number of active sessions reached",
331330
)
331+
effective_profile = profile or self.default_profile
332+
cached_profile_id = (
333+
self.default_profile_id if effective_profile == self.default_profile else None
334+
)
332335
sess = PostgresSession(
333-
self, profile or self.default_profile, False, self.release_number
336+
self, effective_profile, False, self.release_number, cached_profile_id
334337
)
335338
with self.lock:
336339
self.active_sessions.append(sess)
@@ -355,8 +358,12 @@ async def transaction(self, profile: str = None):
355358
code=DatabaseErrorCode.CONNECTION_POOL_EXHAUSTED,
356359
message="Maximum number of active sessions reached",
357360
)
361+
effective_profile = profile or self.default_profile
362+
cached_profile_id = (
363+
self.default_profile_id if effective_profile == self.default_profile else None
364+
)
358365
sess = PostgresSession(
359-
self, profile or self.default_profile, True, self.release_number
366+
self, effective_profile, True, self.release_number, cached_profile_id
360367
)
361368
with self.lock:
362369
self.active_sessions.append(sess)

‎acapy_agent/database_manager/databases/postgresql_normalized/handlers/normalized_handler.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,10 @@ def __init__(
123123

124124
self.EXPIRY_CLAUSE = "(i.expiry IS NULL OR i.expiry > CURRENT_TIMESTAMP)"
125125

126-
async def _ensure_utf8(self, cursor: AsyncCursor) -> None:
127-
await cursor.execute(SQL_SET_UTF8)
126+
async def _ensure_utf8(self, _cursor: AsyncCursor) -> None:
127+
# UTF8 encoding is set via connection pool options (-c client_encoding=UTF8)
128+
# No need to execute SET here - would add unnecessary latency
129+
pass
128130

129131
def _validate_order_by(self, order_by: Optional[str]) -> None:
130132
if order_by and order_by not in self.ALLOWED_ORDER_BY_COLUMNS:

‎acapy_agent/database_manager/databases/postgresql_normalized/session.py‎

Lines changed: 59 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -27,16 +27,26 @@ def __init__(
2727
profile: str,
2828
is_txn: bool,
2929
release_number: str = "release_0",
30+
profile_id: int = None,
3031
):
31-
"""Initialize PostgreSQL session."""
32+
"""Initialize PostgreSQL session.
33+
34+
Args:
35+
database: The PostgresDatabase instance
36+
profile: Profile name
37+
is_txn: Whether this is a transaction
38+
release_number: Schema release number
39+
profile_id: Optional cached profile ID (avoids DB lookup)
40+
41+
"""
3242
self.lock = threading.RLock()
3343
self.database = database
3444
self.pool = database.pool
3545
self.profile = profile
3646
self.is_txn = is_txn
3747
self.release_number = release_number
3848
self.conn = None
39-
self.profile_id = None
49+
self.profile_id = profile_id
4050
self.schema_context = database.schema_context
4151

4252
def _process_value(
@@ -173,19 +183,12 @@ async def __aenter__(self):
173183
self._handle_session_failure(max_retries, e)
174184

175185
async def _acquire_and_validate_connection(self):
176-
"""Acquire and validate database connection."""
186+
"""Acquire database connection from pool.
187+
188+
Note: Connection validation is handled by the pool and getconn's rollback.
189+
The retry logic in __aenter__ handles any stale connection issues.
190+
"""
177191
self.conn = await self.pool.getconn()
178-
try:
179-
async with self.conn.cursor() as cursor:
180-
await cursor.execute("SELECT 1")
181-
except Exception as e:
182-
await self._cleanup_connection()
183-
LOGGER.error("Invalid connection retrieved: %s", str(e))
184-
raise DatabaseError(
185-
code=DatabaseErrorCode.CONNECTION_ERROR,
186-
message="Invalid connection retrieved from pool",
187-
actual_error=str(e),
188-
)
189192

190193
async def _setup_session(self):
191194
"""Setup session with profile and transaction state."""
@@ -204,11 +207,22 @@ def _log_session_start(self):
204207
)
205208

206209
async def _cleanup_connection(self):
207-
"""Clean up database connection."""
210+
"""Clean up database connection during session setup failure."""
208211
if self.conn:
209-
await self.conn.rollback()
210-
await self.pool.putconn(self.conn)
211-
self.conn = None
212+
try:
213+
await self.conn.rollback()
214+
except Exception as e:
215+
LOGGER.warning("[cleanup_connection] Rollback failed: %s", str(e))
216+
217+
try:
218+
await self.pool.putconn(self.conn)
219+
except Exception as e:
220+
LOGGER.error(
221+
"[cleanup_connection] CRITICAL: Failed to return connection: %s",
222+
str(e),
223+
)
224+
finally:
225+
self.conn = None
212226

213227
def _handle_session_failure(self, max_retries: int, error: Exception):
214228
"""Handle session setup failure after retries."""
@@ -260,15 +274,27 @@ async def _handle_non_transaction_mode(self):
260274

261275
async def _cleanup_session(self):
262276
"""Clean up session resources."""
277+
conn_returned = False
263278
try:
264-
await self.conn.rollback()
265279
await self.pool.putconn(self.conn)
280+
conn_returned = True
266281
self.conn = None
282+
except Exception as e:
283+
LOGGER.error(
284+
"[close_session] CRITICAL: Failed to return connection to pool: %s",
285+
str(e),
286+
)
287+
self.conn = None
288+
289+
try:
267290
if self in self.database.active_sessions:
268291
self.database.active_sessions.remove(self)
269-
LOGGER.debug("[close_session] Completed")
270-
except Exception:
271-
pass
292+
except Exception as e:
293+
LOGGER.warning(
294+
"[close_session] Failed to remove from active_sessions: %s", str(e)
295+
)
296+
297+
LOGGER.debug("[close_session] Completed (connection_returned=%s)", conn_returned)
272298

273299
async def count(self, category: str, tag_filter: str | dict = None) -> int:
274300
"""Count entries in a category."""
@@ -583,16 +609,18 @@ async def close(self):
583609
"""Close session."""
584610
if self.conn:
585611
try:
586-
async with self.conn.cursor() as cursor:
587-
await cursor.execute("SELECT 1")
588-
except Exception:
589-
pass
590-
try:
591-
await self.conn.rollback()
592612
await self.pool.putconn(self.conn)
613+
except Exception as e:
614+
LOGGER.error("[close] CRITICAL: Failed to return connection: %s", str(e))
615+
finally:
593616
self.conn = None
617+
618+
try:
594619
if self in self.database.active_sessions:
595620
self.database.active_sessions.remove(self)
596-
LOGGER.debug("[close_session] Completed")
597-
except Exception:
598-
pass
621+
except Exception as e:
622+
LOGGER.warning(
623+
"[close] Failed to remove from active_sessions: %s", str(e)
624+
)
625+
626+
LOGGER.debug("[close] Completed")

‎acapy_agent/database_manager/dbstore.py‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ def create_generator() -> AsyncIterator[Entry]:
100100
return await anext(self._generator) # noqa: F821
101101
except StopAsyncIteration:
102102
LOGGER.error("StopAsyncIteration in __anext__")
103+
await self.aclose()
103104
raise
104105
else:
105106
# Handle sync generators using the executor
@@ -112,6 +113,7 @@ def get_next() -> Entry | None:
112113
loop = asyncio.get_running_loop()
113114
result = await loop.run_in_executor(self._executor, get_next)
114115
if result is None:
116+
await self.aclose()
115117
raise StopAsyncIteration
116118
return result
117119

@@ -120,6 +122,23 @@ def __del__(self) -> None:
120122
# Shut down the executor to clean up resources
121123
self._executor.shutdown(wait=False)
122124

125+
async def aclose(self) -> None:
126+
"""Close the underlying generator and release resources."""
127+
try:
128+
if self._generator:
129+
if self._is_async:
130+
agen_aclose = getattr(self._generator, "aclose", None)
131+
if agen_aclose:
132+
await agen_aclose()
133+
else:
134+
loop = asyncio.get_running_loop()
135+
await loop.run_in_executor(
136+
self._executor,
137+
lambda: getattr(self._generator, "close", lambda: None)(),
138+
)
139+
finally:
140+
self._executor.shutdown(wait=False)
141+
123142

124143
class ScanKeyset(AsyncIterator):
125144
"""Keyset-based scan iterator."""
@@ -196,6 +215,7 @@ def create_generator() -> AsyncGenerator[Entry, None]:
196215
return await anext(self._generator) # noqa: F821
197216
except StopAsyncIteration:
198217
LOGGER.error("StopAsyncIteration in __anext__")
218+
await self.aclose()
199219
raise
200220
else:
201221
# Handle sync generators using the executor
@@ -208,13 +228,31 @@ def get_next() -> Entry | None:
208228
loop = asyncio.get_running_loop()
209229
result = await loop.run_in_executor(self._executor, get_next)
210230
if result is None:
231+
await self.aclose()
211232
raise StopAsyncIteration
212233
return result
213234

214235
def __del__(self) -> None:
215236
"""Clean up resources."""
216237
self._executor.shutdown(wait=False)
217238

239+
async def aclose(self) -> None:
240+
"""Close the underlying generator and release resources."""
241+
try:
242+
if self._generator:
243+
if self._is_async:
244+
agen_aclose = getattr(self._generator, "aclose", None)
245+
if agen_aclose:
246+
await agen_aclose()
247+
else:
248+
loop = asyncio.get_running_loop()
249+
await loop.run_in_executor(
250+
self._executor,
251+
lambda: getattr(self._generator, "close", lambda: None)(),
252+
)
253+
finally:
254+
self._executor.shutdown(wait=False)
255+
218256
async def fetch_all(self) -> Sequence[Entry]:
219257
"""Perform the action."""
220258
rows = []
@@ -776,17 +814,31 @@ def is_transaction(self) -> bool:
776814

777815
async def _open(self) -> DBStoreSession:
778816
"""Perform the action."""
779-
LOGGER.debug("_open called")
817+
import time
818+
819+
start = time.perf_counter()
820+
LOGGER.debug(
821+
"DBOpenSession._open starting for profile=%s, is_txn=%s",
822+
self._profile,
823+
self._is_txn,
824+
)
780825
if self._session:
781826
raise DBStoreError(DBStoreErrorCode.WRAPPER, "Session already opened")
782827
method = self._db.transaction if self._is_txn else self._db.session
828+
LOGGER.debug("Calling db.%s...", "transaction" if self._is_txn else "session")
783829
self._db_session = (
784830
await method(self._profile)
785831
if inspect.iscoroutinefunction(method)
786832
else method(self._profile)
787833
)
834+
LOGGER.debug("Got db_session, calling __aenter__...")
788835
await self._db_session.__aenter__()
789836
self._session = DBStoreSession(self._db_session, self._is_txn)
837+
LOGGER.debug(
838+
"DBOpenSession._open completed in %.3fs for profile=%s",
839+
time.perf_counter() - start,
840+
self._profile,
841+
)
790842
return self._session
791843

792844
def __await__(self) -> DBStoreSession:

0 commit comments

Comments
 (0)