301 lines
10 KiB
Python
301 lines
10 KiB
Python
import json
|
||
import socket
|
||
import struct
|
||
import sys
|
||
import threading
|
||
import os
|
||
import atexit
|
||
from typing import Callable, Any, Union, Tuple, Optional, List
|
||
|
||
|
||
class StatusServer:
|
||
def __init__(self, get_status_func: Callable[[bool], Any], server_address: Union[str, Tuple[str, int]]):
|
||
"""
|
||
初始化服务端
|
||
:param get_status_func: 获取状态的函数，返回当前进程状态字典, 支持一个参数 init，
|
||
当init为True时，表示获取初始化状态，否则为更新状态
|
||
:param server_address: 本地套接字文件路径（Unix域）或 (host, port)（TCP）
|
||
"""
|
||
self.get_status_func = get_status_func
|
||
self.server_address = server_address
|
||
self.clients: List[socket.socket] = []
|
||
self.lock = threading.Lock() # 线程锁
|
||
self.running = False
|
||
self.server_socket = None
|
||
|
||
def handle_client(self, client_socket):
|
||
"""处理客户端连接"""
|
||
try:
|
||
# 发送初始状态
|
||
new_status = self.get_status_func(True)
|
||
status_bytes = json.dumps(new_status).encode() # 使用 JSON 更安全
|
||
packed_data = len(status_bytes).to_bytes(4, 'little') + status_bytes # 添加长度头
|
||
|
||
# 添加到客户端列表
|
||
try:
|
||
# 分块发送
|
||
client_socket.sendall(packed_data) # 发送结束标志
|
||
except Exception as e:
|
||
print(f"Failed to send update to client: {e}")
|
||
client_socket.close()
|
||
return
|
||
|
||
with self.lock:
|
||
self.clients.append(client_socket)
|
||
|
||
# 保持连接以支持后续更新
|
||
while self.running:
|
||
try:
|
||
# 可选：接收客户端心跳或命令
|
||
data = client_socket.recv(1024)
|
||
if not data:
|
||
break
|
||
except:
|
||
break
|
||
|
||
finally:
|
||
# 关闭连接并从列表中移除
|
||
client_socket.close()
|
||
with self.lock:
|
||
if client_socket in self.clients:
|
||
self.clients.remove(client_socket)
|
||
|
||
def start_server(self):
|
||
"""启动本地套接字服务端"""
|
||
self.running = True
|
||
|
||
if isinstance(self.server_address, str):
|
||
# Unix 域套接字
|
||
self.server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||
try:
|
||
os.unlink(self.server_address)
|
||
except OSError:
|
||
if os.path.exists(self.server_address):
|
||
raise
|
||
self.server_socket.bind(self.server_address)
|
||
else:
|
||
# TCP 套接字
|
||
self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||
self.server_socket.bind(self.server_address)
|
||
|
||
self.server_socket.listen(5)
|
||
print(f"Server is listening on {self.server_address}...")
|
||
|
||
try:
|
||
self.running = True
|
||
while self.running:
|
||
client_socket, _ = self.server_socket.accept()
|
||
print("Client connected")
|
||
|
||
# 启动新线程处理客户端
|
||
client_thread = threading.Thread(target=self.handle_client, args=(client_socket,))
|
||
client_thread.start()
|
||
except KeyboardInterrupt:
|
||
print("Shutting down server...")
|
||
finally:
|
||
self.stop()
|
||
|
||
def stop(self):
|
||
"""停止服务端并清理资源"""
|
||
if not self.running:
|
||
return
|
||
self.running = False
|
||
|
||
with self.lock:
|
||
for client in self.clients:
|
||
client.close()
|
||
self.clients.clear()
|
||
|
||
if self.server_socket:
|
||
self.server_socket.close()
|
||
self.server_socket = None
|
||
|
||
# 清理 Unix 套接字文件
|
||
if isinstance(self.server_address, str) and os.path.exists(self.server_address):
|
||
os.remove(self.server_address)
|
||
print(f"Socket file removed: {self.server_address}")
|
||
|
||
def update_status(self, update_data: Optional[dict]=None):
|
||
"""获取最新的状态并推送给所有客户端"""
|
||
if not update_data:
|
||
new_status = self.get_status_func(False)
|
||
else:
|
||
new_status = update_data
|
||
status_bytes = json.dumps(new_status).encode() # 使用 JSON 更安全
|
||
packed_data = len(status_bytes).to_bytes(4, 'little') + status_bytes # 添加长度头
|
||
|
||
with self.lock:
|
||
for client in self.clients:
|
||
print("Sending update to client...")
|
||
print(len(status_bytes), status_bytes, packed_data)
|
||
try:
|
||
client.sendall(packed_data) # 直接发送完整数据
|
||
except Exception as e:
|
||
print(f"Failed to send update to client: {e}")
|
||
client.close()
|
||
if client in self.clients:
|
||
self.clients.remove(client)
|
||
|
||
|
||
class StatusClient:
|
||
def __init__(self, server_address, callback=None):
|
||
"""
|
||
初始化客户端
|
||
:param server_address: Unix 域路径（字符串） 或 TCP 地址元组 (host, port)
|
||
:param callback: 接收到状态更新时的回调函数，接受一个 dict 参数
|
||
"""
|
||
self.server_address = server_address
|
||
self.callback = callback
|
||
self.sock: Optional[socket.socket] = None
|
||
self.running = False
|
||
self.receive_thread = None
|
||
|
||
def connect(self):
|
||
"""连接到服务端"""
|
||
if isinstance(self.server_address, str):
|
||
print("Connecting to Unix socket...", self.server_address)
|
||
# Unix 域套接字
|
||
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||
self.sock.connect(self.server_address)
|
||
else:
|
||
# TCP 套接字
|
||
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||
self.sock.connect(self.server_address)
|
||
|
||
print("Connected to server.")
|
||
|
||
# 启动接收线程
|
||
self.running = True
|
||
self.receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
|
||
self.receive_thread.start()
|
||
|
||
def _receive_loop(self):
|
||
buffer = b''
|
||
while self.running:
|
||
try:
|
||
# 读取长度头
|
||
while len(buffer) < 4:
|
||
data = self.sock.recv(4)
|
||
if not data:
|
||
raise ConnectionResetError("Server closed the connection")
|
||
buffer += data
|
||
length = int.from_bytes(buffer[:4], 'little')
|
||
buffer = buffer[4:]
|
||
|
||
# 读取完整数据
|
||
while len(buffer) < length:
|
||
data = self.sock.recv(length - len(buffer))
|
||
if not data:
|
||
raise ConnectionResetError("Server closed the connection")
|
||
buffer += data
|
||
message = buffer[:length]
|
||
buffer = buffer[length:]
|
||
|
||
# 解析JSON
|
||
status = json.loads(message.decode())
|
||
if self.callback:
|
||
self.callback(status)
|
||
except ConnectionResetError as e:
|
||
print("Connection interrupted:", e)
|
||
self.disconnect()
|
||
break
|
||
except json.JSONDecodeError as e:
|
||
print("JSON parsing failed:", e)
|
||
continue
|
||
except Exception as e:
|
||
print("reception error:", e)
|
||
self.disconnect()
|
||
break
|
||
|
||
def disconnect(self):
|
||
"""断开连接"""
|
||
self.running = False
|
||
if self.sock:
|
||
self.sock.close()
|
||
self.sock = None
|
||
print("Disconnected from server.")
|
||
|
||
def stop(self):
|
||
"""停止客户端"""
|
||
self.disconnect()
|
||
if self.receive_thread and self.receive_thread.is_alive():
|
||
self.receive_thread.join()
|
||
print("Client stopped.")
|
||
|
||
def wait_receive(self):
|
||
if self.receive_thread and self.receive_thread.is_alive():
|
||
self.receive_thread.join()
|
||
|
||
|
||
# 注册退出清理钩子
|
||
def register_cleanup(server_instance):
|
||
def cleanup():
|
||
server_instance.stop()
|
||
|
||
atexit.register(cleanup)
|
||
|
||
# # 示例使用
|
||
# if __name__ == '__main__' and "server" in sys.argv:
|
||
#
|
||
# import time
|
||
#
|
||
# # 模拟的状态存储
|
||
# process_status = {
|
||
# 'process1': 'running',
|
||
# 'process2': 'stopped',
|
||
# "big_data": "<AAAAAAAAAAAAAAAAAFFFFFFFFFFFFFFFFFFAAAAFFFFFFFFFFFFFFFAAAAAAAAAAAAAAAAAAAAAAAAA>"
|
||
# }
|
||
#
|
||
# def get_status():
|
||
# return process_status
|
||
#
|
||
# # 设置 Unix 域套接字地址
|
||
# server_address = './socket_filetransfer.sock'
|
||
#
|
||
# # 创建服务端实例
|
||
# server = StatusServer(get_status, server_address)
|
||
# register_cleanup(server) # 注册退出时清理
|
||
#
|
||
# # 启动服务端线程
|
||
# server_thread = threading.Thread(target=server.start_server)
|
||
# server_thread.daemon = True
|
||
# server_thread.start()
|
||
#
|
||
# # 模拟状态更新
|
||
# try:
|
||
# while True:
|
||
# print(">>>>>>>change<<<<<<<<<<<<<<<")
|
||
# time.sleep(5)
|
||
# process_status['process1'] = 'stopped' if process_status['process1'] == 'running' else 'running'
|
||
# server.update_status()
|
||
#
|
||
# time.sleep(5)
|
||
# process_status['process2'] = 'running' if process_status['process2'] == 'stopped' else 'stopped'
|
||
# server.update_status()
|
||
# except KeyboardInterrupt:
|
||
# pass
|
||
#
|
||
# # 示例使用
|
||
# if __name__ == '__main__' and "client" in sys.argv:
|
||
# # Unix 域套接字示例：
|
||
# server_address = './socket_filetransfer.sock'
|
||
#
|
||
# # 示例回调函数
|
||
# def on_status_update(status_dict):
|
||
# print("[Callback] New status received:")
|
||
# for k, v in status_dict.items():
|
||
# print(f" - {k}: {v}")
|
||
#
|
||
#
|
||
# client = StatusClient(server_address, callback=on_status_update)
|
||
# try:
|
||
# client.connect()
|
||
#
|
||
# # 主线程保持运行，防止程序退出
|
||
# while client.running:
|
||
# pass
|
||
# except KeyboardInterrupt:
|
||
# print("Client shutting down...")
|
||
# client.stop()
|