From 282063cb25ce5ace0fb52cbf2caa5652c6a84cd0 Mon Sep 17 00:00:00 2001 From: He Wang Date: Mon, 16 Mar 2026 11:23:58 +0800 Subject: [PATCH 1/3] fix: release pooled DB connections after each operation to prevent MaxConnectionsExceeded Peewee's PooledMySQLDatabase does not automatically return connections to the pool after execute_sql(). Under high concurrency this causes pool exhaustion and MaxConnectionsExceeded errors. - Add `manual_close()` to RetryingPooledMySQLDatabase to return connection to pool - Add `@release_connection` decorator to all OceanBaseRedisDb public methods - Split set_object/set_obj/deleteIfExpired into internal (_-prefixed) and public versions to avoid releasing connections inside atomic() transactions - Add `_release_connection()` to MysqlDistributedLock for doAcquire/delete_if_equal Made-with: Cursor --- api/db/db_models.py | 12 +++++ rag/utils/ob_redis_conn.py | 103 ++++++++++++++++++++++++++++++++----- 2 files changed, 102 insertions(+), 13 deletions(-) diff --git a/api/db/db_models.py b/api/db/db_models.py index 50a50628e..c774e8cfb 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -245,6 +245,18 @@ def __init__(self, *args, **kwargs): self.retry_delay = kwargs.pop("retry_delay", 1) super().__init__(*args, **kwargs) + def manual_close(self): + """ + Close the current thread's connection and return it to the pool. + + In peewee's PooledMySQLDatabase, connections are not automatically returned + to the pool after execute_sql(). They only get returned when close() is called. + This method should be called after database operations to prevent connection + pool exhaustion, especially in scenarios with many concurrent operations. + """ + if not self.is_closed(): + self.close() + def execute_sql(self, sql, params=None, commit=True): for attempt in range(self.max_retries + 1): try: diff --git a/rag/utils/ob_redis_conn.py b/rag/utils/ob_redis_conn.py index f56458d0d..281921ee0 100644 --- a/rag/utils/ob_redis_conn.py +++ b/rag/utils/ob_redis_conn.py @@ -3,8 +3,10 @@ import logging import time import uuid +from contextlib import contextmanager from datetime import datetime, timedelta from decimal import Decimal +from functools import wraps import trio from peewee import IntegrityError, ProgrammingError @@ -27,6 +29,40 @@ def get_db(): return DATABASE +def release_connection(func): + """ + Decorator to ensure database connection is returned to pool after method execution. + + Peewee's PooledMySQLDatabase does not automatically return connections to the pool + after execute_sql(). This decorator ensures connections are properly released + to prevent connection pool exhaustion under high concurrency. + """ + @wraps(func) + def wrapper(self, *args, **kwargs): + try: + return func(self, *args, **kwargs) + finally: + if isinstance(self.db, RetryingPooledMySQLDatabase): + self.db.manual_close() + return wrapper + + +@contextmanager +def db_connection(db): + """ + Context manager to ensure database connection is returned to pool after use. + + Peewee's PooledMySQLDatabase does not automatically return connections to the pool + after execute_sql(). This context manager ensures connections are properly released + to prevent connection pool exhaustion under high concurrency. + """ + try: + yield db + finally: + if isinstance(db, RetryingPooledMySQLDatabase): + db.manual_close() + + # 由于这里数据库 message 返回值是str,而 redis_conn stream 接口返回的是 dict,所以 RedisMsg 接口初始化略有不同，因此单独声明一个 RedisMsg class RedisMsg: def __init__(self, consumer, queue_name, group_name, msg_id, message): @@ -76,6 +112,7 @@ def __init__(self, db=None): def register_scripts(self) -> None: raise NotImplementedError("Not implemented") + @release_connection def health(self): try: self.db.execute_sql("select 1 from dual") @@ -86,13 +123,12 @@ def health(self): def is_alive(self): return self.health() + @release_connection def exist(self, k): if not self.db: return try: - cursor = self.db.execute_sql('select count(1) from cache where cache_key = %s and expire_time > now()', (k)) - ret = cursor.fetchone() return ret[0] == 1 except Exception as e: @@ -101,6 +137,7 @@ def exist(self, k): else: logging.warning("RedisDB.exist " + str(k) + " got exception: " + str(e)) + @release_connection def delete_if_equal(self, key: str, expected_value: str) -> bool: try: cursor = self.db.execute_sql('delete from cache where cache_key = %s and cache_value = %s and expire_time ' @@ -114,6 +151,7 @@ def delete_if_equal(self, key: str, expected_value: str) -> bool: "RedisDB.delete_if_equal " + str(key) + ":" + str(expected_value) + " got exception: " + str(e)) return False + @release_connection def delete(self, key) -> bool: try: self.db.execute_sql('delete from cache where cache_key = %s ', (key)) @@ -125,7 +163,7 @@ def delete(self, key) -> bool: logging.warning("RedisDB.delete " + str(key) + " got exception: " + str(e)) return False - def deleteIfExpired(self, key) -> bool: + def _deleteIfExpired(self, key) -> bool: try: self.db.execute_sql('delete from cache where cache_key = %s and expire_time < now() ', key) return True @@ -136,6 +174,11 @@ def deleteIfExpired(self, key) -> bool: logging.warning("RedisDB.delete " + str(key) + " got exception: " + str(e)) return False + @release_connection + def deleteIfExpired(self, key) -> bool: + return self._deleteIfExpired(key) + + @release_connection def get(self, k): if not self.db: return None @@ -150,9 +193,9 @@ def get(self, k): else: logging.warning("RedisDB.get " + str(k) + " got exception: " + str(e)) - def set_obj(self, k, obj, exp=3600): + def _set_obj(self, k, obj, exp=3600): try: - self.set_object(k, obj, exp) + self._set_object(k, obj, exp) return True except Exception as e: if is_table_missing_exception(e): @@ -161,12 +204,21 @@ def set_obj(self, k, obj, exp=3600): logging.warning("RedisDB.set_obj " + str(k) + " got exception: " + str(e)) return False - def set_object(self, k, obj, exp=3600): + @release_connection + def set_obj(self, k, obj, exp=3600): + return self._set_obj(k, obj, exp) + + def _set_object(self, k, obj, exp=3600): expire_time = datetime.now() + timedelta(seconds=exp) self.db.execute_sql('replace into cache (cache_key, cache_value, expire_time) values (%s, %s, %s)', (k, json.dumps(obj, ensure_ascii=False), expire_time)) return True + @release_connection + def set_object(self, k, obj, exp=3600): + return self._set_object(k, obj, exp) + + @release_connection def set(self, k, v, exp=3600): try: expire_time = datetime.now() + timedelta(seconds=exp) @@ -180,10 +232,11 @@ def set(self, k, v, exp=3600): logging.warning("RedisDB.set " + str(k) + " got exception: " + str(e)) return False + @release_connection def setNx(self, k, v, exp=3600): try: # 删除过期的kv - self.deleteIfExpired(k) + self._deleteIfExpired(k) expire_time = datetime.now() + timedelta(seconds=exp) self.db.execute_sql('insert into cache (cache_key, cache_value, expire_time) values (%s, %s, %s)', (k, v, expire_time)) @@ -202,6 +255,7 @@ def transaction(self, key, value, exp=3600): return self.setNx(key, value, exp) # zset + @release_connection def zadd(self, key: str, member: str, score: float): try: with self.db.atomic(): @@ -210,7 +264,7 @@ def zadd(self, key: str, member: str, score: float): ret = cursor.fetchone() if ret is None: mp = {member: score} - return self.set_object(key, mp) + return self._set_object(key, mp) else: id = ret[0] cursor = self.db.execute_sql( @@ -228,6 +282,7 @@ def zadd(self, key: str, member: str, score: float): logging.warning("RedisDB.zadd " + str(key) + " got exception: " + str(e)) return False + @release_connection def zcount(self, key: str, min, max: float): try: cursor = self.db.execute_sql( @@ -245,6 +300,7 @@ def zcount(self, key: str, min, max: float): logging.warning("RedisDB.zcount " + str(key) + " got exception: " + str(e)) return 0 + @release_connection def zpopmin(self, key: str, count: int): try: with self.db.atomic() as trx: @@ -264,7 +320,7 @@ def zpopmin(self, key: str, count: int): break for k, v in ret.items(): del mp[k] - self.set_object(key, mp) + self._set_object(key, mp) return ret except Exception as e: if is_table_missing_exception(e): @@ -273,6 +329,7 @@ def zpopmin(self, key: str, count: int): logging.warning("RedisDB.zpopmin " + str(key) + " got exception: " + str(e)) return None + @release_connection def sadd(self, key: str, member: str): try: with self.db.atomic(): @@ -281,11 +338,11 @@ def sadd(self, key: str, member: str): ret = cursor.fetchone() if ret is None: st = {member} - return self.set_object(key, list(st)) + return self._set_object(key, list(st)) else: st = set(json.loads(ret[0])) st.add(member) - return self.set_obj(key, list(st)) + return self._set_obj(key, list(st)) except Exception as e: if is_table_missing_exception(e): pass @@ -293,6 +350,7 @@ def sadd(self, key: str, member: str): logging.warning("RedisDB.sadd " + str(key) + " got exception: " + str(e)) return False + @release_connection def srem(self, key: str, member: str): try: with self.db.atomic(): @@ -304,7 +362,7 @@ def srem(self, key: str, member: str): else: st = set(json.loads(ret[0])) st.discard(member) - return self.set_object(key, list(st)) + return self._set_object(key, list(st)) except Exception as e: if is_table_missing_exception(e): pass @@ -312,6 +370,7 @@ def srem(self, key: str, member: str): logging.warning("RedisDB.srem " + str(key) + " got exception: " + str(e)) return False + @release_connection def smembers(self, key: str): try: cursor = self.db.execute_sql("select cache_value from cache where cache_key = %s and expire_time > " @@ -328,6 +387,7 @@ def smembers(self, key: str): logging.warning("RedisDB.smembers " + str(key) + " got exception: " + str(e)) return [] + @release_connection def zrangebyscore(self, key: str, min: float, max: float): try: cursor = self.db.execute_sql("select cache_value from cache where cache_key = %s and expire_time > now() ", @@ -351,6 +411,7 @@ def zrangebyscore(self, key: str, min: float, max: float): return None # 以下是redis stream + @release_connection def queue_product(self, queue, message) -> bool: """ 向消息队列推送消息，如果消息队列不存在，则创建消息队列 @@ -370,6 +431,7 @@ def queue_product(self, queue, message) -> bool: ) return False + @release_connection def queue_consumer(self, queue_name, group_name, consumer_name, msg_id=b">"): """ 消费者拉取消息： @@ -463,6 +525,7 @@ def queue_consumer(self, queue_name, group_name, consumer_name, msg_id=b">"): ) return None + @release_connection def get_pending_msg(self, queue, group_name): """ 获取消费者组 {group_name} 对消息队列 {queue} 已经读取，但是没有 ACK 的消息。 @@ -498,6 +561,7 @@ def get_pending_msg(self, queue, group_name): return [] # 获取本消费组内没有 ack的消息 + @release_connection def get_unacked_iterator(self, queue_names: list[str], group_name, consumer_name): """ 获取消费者组未 ack 的消息，返回一个迭代器 @@ -544,7 +608,7 @@ def get_unacked_iterator(self, queue_names: list[str], group_name, consumer_name ''' 消息重新入队 ''' - + @release_connection def requeue_msg(self, queue: str, group_name: str, msg_id: object): """ 将未 ack 的消息重新入队列 @@ -570,6 +634,7 @@ def requeue_msg(self, queue: str, group_name: str, msg_id: object): "RedisDB.requeue_msg " + str(queue) + " got exception: " + str(e) ) + @release_connection def xack(self, queue: str, group_name: str, msg_id: object): """ 提交消息 ack @@ -578,6 +643,7 @@ def xack(self, queue: str, group_name: str, msg_id: object): "update message_consumption set ack = true where stream = %s and group_name = %s and " "message_id = %s", (queue, group_name, msg_id)) + @release_connection def queue_info(self, queue: str, group_name: str) -> dict | None: """ 获取消息队列，某个消费者组的消费情况。本项目用到的属性有： @@ -635,6 +701,13 @@ def __init__(self, lock_key, lock_value=None, timeout=10, blocking_timeout=1): # blocking_timeout 没用到，预留 self.blocking_timeout = blocking_timeout + def _release_connection(self): + """ + Release the current thread's database connection back to the pool. + """ + if isinstance(self.db, RetryingPooledMySQLDatabase): + self.db.manual_close() + def acquire(self): """ 获取锁 @@ -677,6 +750,8 @@ def doAcquire(self): else: logging.info(f"lock acquire failed:{self.lock_key}-{self.lock_value}-{e}") return False + finally: + self._release_connection() def release(self): """ @@ -705,6 +780,8 @@ def delete_if_equal(self): else: logging.warning(f"release lock failed:{self.lock_key}-{self.lock_value}-{e}") return False + finally: + self._release_connection() if __name__ == '__main__': From e90a562212bdb2ab4289b0bba2534d283c33fc52 Mon Sep 17 00:00:00 2001 From: He Wang Date: Mon, 16 Mar 2026 13:08:16 +0800 Subject: [PATCH 2/3] fix: address review comments on connection pool release - Remove @release_connection from get_unacked_iterator since it is a generator; the decorator's finally block runs immediately on generator creation rather than after iteration, causing premature connection release - Remove unused db_connection context manager - Wrap manual_close() in try/except to avoid masking original exceptions when called from finally blocks Made-with: Cursor --- api/db/db_models.py | 7 +++++-- rag/utils/ob_redis_conn.py | 17 ----------------- 2 files changed, 5 insertions(+), 19 deletions(-) diff --git a/api/db/db_models.py b/api/db/db_models.py index c774e8cfb..ffc36dea1 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -254,8 +254,11 @@ def manual_close(self): This method should be called after database operations to prevent connection pool exhaustion, especially in scenarios with many concurrent operations. """ - if not self.is_closed(): - self.close() + try: + if not self.is_closed(): + self.close() + except Exception: + pass def execute_sql(self, sql, params=None, commit=True): for attempt in range(self.max_retries + 1): diff --git a/rag/utils/ob_redis_conn.py b/rag/utils/ob_redis_conn.py index 281921ee0..a4ceefe71 100644 --- a/rag/utils/ob_redis_conn.py +++ b/rag/utils/ob_redis_conn.py @@ -3,7 +3,6 @@ import logging import time import uuid -from contextlib import contextmanager from datetime import datetime, timedelta from decimal import Decimal from functools import wraps @@ -47,21 +46,6 @@ def wrapper(self, *args, **kwargs): return wrapper -@contextmanager -def db_connection(db): - """ - Context manager to ensure database connection is returned to pool after use. - - Peewee's PooledMySQLDatabase does not automatically return connections to the pool - after execute_sql(). This context manager ensures connections are properly released - to prevent connection pool exhaustion under high concurrency. - """ - try: - yield db - finally: - if isinstance(db, RetryingPooledMySQLDatabase): - db.manual_close() - # 由于这里数据库 message 返回值是str,而 redis_conn stream 接口返回的是 dict,所以 RedisMsg 接口初始化略有不同，因此单独声明一个 RedisMsg class RedisMsg: @@ -561,7 +545,6 @@ def get_pending_msg(self, queue, group_name): return [] # 获取本消费组内没有 ack的消息 - @release_connection def get_unacked_iterator(self, queue_names: list[str], group_name, consumer_name): """ 获取消费者组未 ack 的消息，返回一个迭代器 From 4a784049f99bca7210321fd264586d57e27a2ffc Mon Sep 17 00:00:00 2001 From: He Wang Date: Mon, 16 Mar 2026 13:19:24 +0800 Subject: [PATCH 3/3] fix: address second round of review comments - Log at debug level in manual_close() instead of silently swallowing exceptions, to aid troubleshooting connection pool issues - Fix queue_info() comparing fetchone() result directly with 0; fetchone() returns a tuple like (count,), so use ret[0] == 0 instead Made-with: Cursor --- api/db/db_models.py | 4 ++-- rag/utils/ob_redis_conn.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/api/db/db_models.py b/api/db/db_models.py index ffc36dea1..99ee2923d 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -257,8 +257,8 @@ def manual_close(self): try: if not self.is_closed(): self.close() - except Exception: - pass + except Exception as e: + logging.debug(f"manual_close failed: {e}") def execute_sql(self, sql, params=None, commit=True): for attempt in range(self.max_retries + 1): diff --git a/rag/utils/ob_redis_conn.py b/rag/utils/ob_redis_conn.py index a4ceefe71..467f6f5f7 100644 --- a/rag/utils/ob_redis_conn.py +++ b/rag/utils/ob_redis_conn.py @@ -640,7 +640,7 @@ def queue_info(self, queue: str, group_name: str) -> dict | None: cursor = self.db.execute_sql( "select count(1) from message_subscribe where stream = %s and group_name = %s", (queue, group_name)) ret = cursor.fetchone() - if ret == 0: + if ret is None or ret[0] == 0: return None else: cursor = self.db.execute_sql(