You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
68 lines
2.2 KiB
68 lines
2.2 KiB
import logging
|
|
import threading
|
|
from queue import Queue
|
|
|
|
from pymilvus import connections
|
|
|
|
# 配置日志
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 1. 手动实现 Milvus 连接池
|
|
class MilvusConnectionPool:
|
|
def __init__(self, host, port, max_connections=5):
|
|
"""
|
|
初始化 Milvus 连接池
|
|
:param host: Milvus 主机地址
|
|
:param port: Milvus 端口
|
|
:param max_connections: 最大连接数
|
|
"""
|
|
self.host = host
|
|
self.port = port
|
|
self.max_connections = max_connections
|
|
self._pool = Queue(max_connections)
|
|
self._lock = threading.Lock()
|
|
|
|
# 初始化连接池
|
|
for _ in range(max_connections):
|
|
self._pool.put(self._create_connection())
|
|
|
|
def _create_connection(self):
|
|
"""
|
|
创建一个新的 Milvus 连接
|
|
:return: Milvus 连接对象
|
|
"""
|
|
return connections.connect(host=self.host, port=self.port, alias="default")
|
|
|
|
def get_connection(self):
|
|
logger.info(f"获取连接，当前可用连接数: {self._pool.qsize()}")
|
|
"""
|
|
从连接池中获取一个连接
|
|
:return: Milvus 连接对象
|
|
"""
|
|
with self._lock:
|
|
if not self._pool.empty():
|
|
return self._pool.get()
|
|
else:
|
|
raise Exception("连接池已满，无法获取连接")
|
|
|
|
def release_connection(self, connection):
|
|
"""
|
|
释放连接，将其放回连接池
|
|
:param connection: Milvus 连接对象
|
|
"""
|
|
with self._lock:
|
|
if self._pool.qsize() < self.max_connections:
|
|
self._pool.put(connection)
|
|
else:
|
|
connections.disconnect("default")
|
|
logger.info(f"释放连接，当前可用连接数: {self._pool.qsize()}")
|
|
|
|
def close(self):
|
|
"""
|
|
关闭连接池，释放所有连接
|
|
"""
|
|
with self._lock:
|
|
while not self._pool.empty():
|
|
connection = self._pool.get()
|
|
connections.disconnect("default") |