Python 基于队列实现 tcp socket 毗连池

毗连池实现

socket_pool.py

1. # -*- coding:utf-8 -*-
2. import socket
3. import time
4. import threading
5. import os
6. import logging
7. import traceback
8. from queue import Queue, Empty
10. _logger = logging.getLogger('mylogger')
12. class SocketPool:
13.     def __init__(self, host, port, min_connections=10, max_connections=10):
14.         '''
15.         初始化Socket连接池
16.         :param host: 目标主机地址
17.         :param port: 目标端口号
18.         :param min_connections: 最小连接数
19.         :param max_connections: 最大连接数
20.         '''
21.         self.host = host
22.         self.port = port
23.         self.min_connections = min_connections
24.         self.max_connections = max_connections
25.         self.busy_sockets_dict = {} # 存放从连接池取出的socket的id
26.         self._sock_lock = threading.Lock()  # 线程锁保证计数正确
27.         self._pool = Queue(max_connections)  # 基于线程安全的队列存储连接
28.         self._lock = threading.Lock()        # 线程锁保证资源安全:
29.         self._init_pool()                    # 预创建连接
30.         self._start_health_check()           # 启动连接健康检查线程
32.     def _init_pool(self):
33.         '''预创建连接并填充到池中'''
34.         
35.         for _ in range(self.min_connections):
36.             sock = self._create_socket()
37.             self._pool.put(sock)
39.     def _create_socket(self):
40.         '''创建新的Socket连接'''
41.         
42.         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
43.         try:
44.             sock.connect((self.host, self.port))
45.             return sock
46.         except socket.error as e:
47.             raise ConnectionError(f'Failed to connect: {e}')  # 连接失败抛出异常
49.     def _start_health_check(self):
50.         '''启动后台线程定期检查连接有效性'''
51.         
52.         def check():
53.             while True:
54.                 with self._lock:
55.                     for _ in range(self._pool.qsize()):
56.                         sock = self._pool.get()
57.                         self.busy_sockets_dict[sock] = 1
58.                         try:
59.                             sock.send(b'PING<END>')  # 发送心跳包验证连接状态
60.                             # 以下 11 为服务端返回数据字节长度，不能乱写，否则会导致获取非健康检查响应报文数据存在多余内容，不符合格式，从而导致数据解析问题
61.                             sock.recv(11)
62.                             self._pool.put(sock)
63.                             self.busy_sockets_dict.pop(sock)
64.                         except (socket.error, ConnectionResetError):
65.                             _logger.error('socket连接健康检查出错：%s, 关闭失效连接并创建新连接替换' % traceback.format_exc())
66.                             sock.close()  # 关闭失效连接并创建新连接替换
67.                             self.busy_sockets_dict.pop(sock)
69.                             new_sock = self._create_socket()
70.                             self._pool.put(new_sock)
71.                     
72.                     # 如果sock数量小于最小数量，则补充
73.                     for _ in range(0, self.min_connections - self._pool.qsize()):
74.                         new_sock = self._create_socket()
75.                         self._pool.put(new_sock)
76.                 time.sleep(60)  # 每60秒检查一次
77.         threading.Thread(target=check, daemon=True).start()
79.     def get_connection(self):
80.         '''
81.         从池中获取一个可用连接
82.         :return: socket对象
83.         '''
84.         
85.         with self._sock_lock:
86.             if self._pool.empty():
87.                 if len(self.busy_sockets_dict.keys()) < self.max_connections:
88.                     new_sock = self._create_socket()
89.                     self.busy_sockets_dict[new_sock] = 1
90.                     return new_sock
91.                 else:
92.                     raise Empty('No available connections in pool')
93.             else:
94.                 try:
95.                     sock = self._pool.get(block=False)
96.                     self.busy_sockets_dict[sock] = 1
97.                     return sock
98.                 except Exception:
99.                     _logger.error('获取socket连接出错：%s' % traceback.format_exc())
100.                     raise
101.                
103.     def release_connection(self, sock):
104.         '''
105.         将连接归还到池中
106.         :param sock: 待归还的socket对象
107.         '''
108.         if not sock._closed:
109.             self._pool.put(sock)
110.         if sock in self.busy_sockets_dict:
111.             self.busy_sockets_dict.pop(sock)
114.     def close_all(self):
115.         '''关闭池中所有连接'''
116.         
117.         while not self._pool.empty():
118.             sock = self._pool.get()
119.             sock.close()
120.             self.busy_sockets_dict.pop(sock.id)
121.         self.busy_sockets_dict = {} # 兜底
123. host = os.environ.get('MODBUS_TCP_SERVER_HOST', '127.0.0.1')
124. port = int(os.environ.get('MODBUS_TCP_SERVER_PORT', '9000'))
125. min_connections = int(os.environ.get('DJANGO_SOCKET_POOL_MAX_CONNECTIONS', '10'))
126. max_connections = int(os.environ.get('DJANGO_SOCKET_POOL_MAX_CONNECTIONS', '100'))
127. socketPool = SocketPool(host, port, min_connections, max_connections)

复制代码

使用毗连池

1. from socket_pool import socketPool
3. def send_socket_msg(data):
4.     global socketPool
5.    
6.     try:
7.         sock = None
8.         # 获取连接（支持超时控制）
9.         sock = socketPool.get_connection()
10.         # 发送数据
11.         sock.sendall(data.encode('utf-8'))
12.     except Exception:
13.         error_msg = '发送消息出错：%s' % traceback.format_exc()
14.         _logger.error(error_msg)
15.         
16.         if sock is not None:
17.             sock.close()
18.             socketPool.release_connection(sock)
19.         return send_socket_msg(data)
20.    
21.     response = ''
22.     try:
23.         while True:
24.             chunk = sock.recv(4096)
25.             chunk = chunk.decode('utf-8')
26.             response += chunk
27.             if response.endswith('<END>'):
28.                 response = response.rstrip('<END>')
29.                 return {'success':True, 'message':response}
30.     except Exception:
31.         error_msg = '获取消息出错：%s' % traceback.format_exc()
32.         _logger.error(error_msg)
33.         return {'success':False, 'message': error_msg}
34.     finally:
35.         # 必须归还连接！
36.         socketPool.release_connection(sock)

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。