WebSocket 任务分发系统代码深度分析与应用

一 系统概述

本系统是一个基于 WebSocket 的分布式任务处理架构，实现了企业微信机器人消息接收、任务分发、文件下载和业务逻辑处理的完整流程。系统采用客户端-服务器模式，通过 WebSocket 进行高效通信，支持多种任务类型的并行处理。

1.1核心组件

组件名称	功能描述	关键技术
WebSocketServer	任务路由中心	websockets, asyncio
HTTP Server	企业微信消息接收	BaseHTTPRequestHandler
Download Client	DS文件下载处理	aiohttp, base64
Logic Client	业务逻辑处理	数据分析，报告生成
Flask API	外部下载接口	Flask, RESTful

1.2系统场景

企业微信机器人收到请求参数之后，解密将参数传给webscoket server然后分发给不同的客户端进行处理，然后再将结果返回到server server再同步给http服务器

1.3架构流程图

企业微信消息 → HTTP Server → WebSocket Server → Download Client → Logic Client → 结果返回 外部API请求 → Flask Server → WebSocket Server → Download Client → 文件下载返回

二 WebSocket 服务器核心代码分析 (webServer.py)

2.1 任务路由器实现

python

class WebSocketTaskRouter:
    def __init__(self):
        # 三类客户端连接管理
        self.clients = {
            "http": set(),      # HTTP API客户端（消息接收）
            "download": set(),  # 下载处理客户端（文件下载）
            "logic": set()      # 业务逻辑客户端（数据分析）
        }
        self.clients_lock = asyncio.Lock()  # 线程安全的客户端管理
        self.pending_tasks = {}             # 任务状态跟踪字典
        self.pending_tasks_lock = asyncio.Lock()  # 任务状态锁

2.2 客户端注册机制

python

async def register_client(self, ws, client_type):
    """注册客户端并开始消息监听循环"""
    async with self.clients_lock:
        self.clients[client_type].add(ws)  # 添加到对应客户端集合
    
    print(f"{client_type} client registered")
    try:
        # 持续监听该客户端的消息
        async for message in ws:
            await self.handle_message(ws, client_type, message)
    finally:
        # 客户端断开时清理
        async with self.clients_lock:
            self.clients[client_type].discard(ws)

2.3 消息路由逻辑

python

async def handle_message(self, ws, client_type, message):
    """根据客户端类型分发消息处理"""
    try:
        data = json.loads(message)
        task_type = data.get("task_type")
        task_id = data.get("task_id")
        
        # 根据客户端类型选择处理方式
        if client_type == "http":
            await self.handle_http_message(ws, task_type, task_id, message)
        elif client_type == "download":
            await self.handle_download_message(task_type, task_id, message)
        elif client_type == "logic":
            await self.handle_logic_message(task_type, task_id, message)
            
    except Exception as e:
        print(f"消息处理失败: {e}")

2.4 任务转发机制

python

async def forward_task(self, from_ws, target_type, task_id, message, update_type=None):
    """将任务转发到目标客户端"""
    async with self.clients_lock:
        targets = list(self.clients[target_type])  # 获取目标客户端列表
    
    if not targets:
        error_msg = f"没有可用的 {target_type} 客户端"
        if from_ws:
            await from_ws.send(json.dumps({"error": f"no {target_type} client"}))
        return
    
    target_ws = targets[0]  # 选择第一个可用客户端
    try:
        await target_ws.send(message)  # 发送任务消息
        
        # 更新任务状态
        async with self.pending_tasks_lock:
            if task_id not in self.pending_tasks:
                self.pending_tasks[task_id] = {
                    "from": from_ws,
                    "type": update_type or target_type,
                    "timestamp": time.time()
                }
    except Exception as e:
        print(f"发送到 {target_type} 客户端失败: {e}")

三 HTTP 服务器代码分析 (httpserver.py)

3.1 WebSocket 客户端封装

python

class WSClient:
    """全局WebSocket客户端，支持同步等待响应"""
    _instance = None
    _lock = threading.Lock()  # 线程安全的单例模式

    @classmethod
    def get_instance(cls, ws_url=None):
        with cls._lock:
            if cls._instance is None:
                cls._instance = WSClient(ws_url)
                cls._instance.start()  # 启动后台线程
            return cls._instance

    def start(self):
        """启动WebSocket客户端线程"""
        self.loop = asyncio.new_event_loop()
        threading.Thread(target=self._run_loop, daemon=True).start()
        self.connected_event.wait()  # 等待连接成功

3.2 同步等待响应机制

python

def send_and_wait(self, data, task_id, timeout=30):
    """发送数据并同步等待响应"""
    if not self.connected_event.is_set():
        raise Exception("WebSocket未连接")
    
    # 跨线程调用异步方法
    fut = asyncio.run_coroutine_threadsafe(
        self._send_and_wait(data, task_id, timeout), self.loop
    )
    return fut.result(timeout=timeout)

async def _send_and_wait(self, data, task_id, timeout):
    """异步发送并等待响应"""
    future = self.loop.create_future()
    self.pending_tasks[task_id] = future  # 注册等待任务
    
    await self.ws.send(data)  # 发送数据
    try:
        result = await asyncio.wait_for(future, timeout)  # 等待结果
        return result
    finally:
        self.pending_tasks.pop(task_id, None)  # 清理任务

3.3 企业微信消息处理

python

class MessageProcessor:
    def decrypt_message(self, post_data, args):
        """解密企业微信消息"""
        ret, msg = self.wxcpt.DecryptMsg(
            post_data,
            args['msg_signature'][0],
            args['timestamp'][0], 
            args['nonce'][0]
        )
        if ret != 0:
            raise ValueError("消息解密失败")
        return ET.fromstring(msg)  # 返回XML解析结果

    def is_duplicate(self, post_data: bytes) -> bool:
        """SHA256哈希检测重复请求"""
        request_hash = hashlib.sha256(post_data).hexdigest()
        return request_hash in self.processed_requests

四 下载客户端代码分析 (WebSoketDownload.py)

4.1 参数解析器

python

def parse_args_items(args_items: list) -> Tuple[Optional[str], str, str]:
    """智能参数解析，兼容多种格式"""
    if len(args_items) >= 3:
        return args_items[0], args_items[1], args_items[2]  # dsid, ds_number, server
    elif len(args_items) == 2:
        if not any("服" in x for x in args_items):
            return args_items[0], args_items[1], "测试22服"  # 默认服务器
        else:
            return None, args_items[0], args_items[1]  # 只有ds_number和server
    elif len(args_items) == 1:
        return None, args_items[0], "测试22服"  # 只有ds_number

4.2 文件下载实现

python

async def download_ds(download_url: str, auth: aiohttp.BasicAuth, timeout_sec: int = 30):
    """异步下载DS文件"""
    timeout = aiohttp.ClientTimeout(total=timeout_sec)
    try:
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(download_url, auth=auth) as resp:
                if resp.status == 200:
                    return await resp.read(), None  # 返回文件内容
                else:
                    return None, f"HTTP状态码: {resp.status}"
    except Exception as e:
        return None, f"下载异常: {e}"

4.3 消息发送封装

python

async def send_download_result(websocket, data: Dict[str, Any], box_id: Optional[str], 
                              ds_number: str, status: str, error: Optional[str],
                              send_lock: asyncio.Lock, b64_data: Optional[str] = None):
    """统一的结果发送格式"""
    result = {
        'task_type': 'download_done',
        'task_id': data.get('task_id'),
        'ds_number': ds_number,
        'status': status,
    }
    
    if box_id is not None:
        result['box_id'] = box_id
    if status == 'success':
        result['data'] = b64_data  # Base64编码的文件数据
    else:
        result['error'] = error
    
    await safe_send(websocket, json.dumps(result), send_lock)

五 逻辑客户端代码分析 (WebSoketLogic.py)

5.1 任务处理器

python

async def handle_logic_task(task: dict, websocket):
    """任务分发处理器"""
    task_type = task.get('task_type')
    task_id = task.get('task_id')
    
    try:
        if task_type == 'logic':
            # 执行扫图逻辑
            response = run(task['args_items'], task)
            await send_task_result(websocket, task_id, "completed", response)
            
        elif task_type == 'download_done':
            # 处理下载完成的任务
            await handle_download_done_task(task, websocket)
            
    except Exception as e:
        await send_task_result(websocket, task_id, "failed", str(e))

5.2 DS文件分析流程

python

async def handle_ds_analysis(temp_dir: Path, filename: Path, box_id: str):
    """DS文件分析和报告生成"""
    if box_id != "获取ds":
        # 使用宝箱配置进行分析
        analyzer = DropConfigAnalyzer('./servitor/config/游戏内掉落表.xls')
        return analyzer.full_analysis(
            drop_config_id=int(box_id),
            log_path=filename,
            output_file=temp_dir / f"drop_config_{box_id}_analysis.xlsx"
        )
    else:
        # 标准DS报告生成
        return generate_ds_reports(temp_dir)

六. 配置管理系统 (config_ini.py, PipeApi.py)

6.1 统一配置管理

python

# config_ini.py - 集中式配置管理
WECHAT_NOTIFY_KEY = "123456
WX_TOKEN = "123456"
WX_ENCODING_AES_KEY = "123456"
HTTP_SERVER_IP = "9.114.1.111"
HTTP_SERVER_PORT = 18810
WS_SERVER_URL = "ws://9.114.1.111:18819"

6.2 管道信息获取

python

# PipeApi.py - 构建信息获取
def get_latest_build_number(pipeline_id_dict):
    """获取最新构建编号"""
    for pipeline_custom_name, pipeline_id in pipeline_id_dict.items():
        pipeline_info = PipelineInfo(project_id, pipeline_id, pipeline_custom_name, 
                                   bk_access_token, bk_username)
        pipeline_info.fetch_pipeline_build_list()
        
        # 查找包含特定APK文件的构建
        for record in pipeline_info.pipeline_build_list['data']['records']:
            if record.get("artifactList"):
                for artifact in record["artifactList"]:
                    if artifact["name"].endswith(".shell.signed.apk"):
                        return record["buildNum"]  # 返回构建编号

七. 安全加密模块 (WXBizMsgCrypt3.py)

7.1 消息加解密核心

python

class Prpcrypt(object):
    """企业微信消息加解密"""
    def encrypt(self, text, receiveid):
        """加密明文"""
        # 添加16位随机字符串和长度信息
        text = self.get_random_str() + struct.pack("I", socket.htonl(len(text))) + text + receiveid.encode()
        text = pkcs7.encode(text)  # PKCS7填充
        
        cryptor = AES.new(self.key, self.mode, self.key[:16])  # AES-CBC加密
        ciphertext = cryptor.encrypt(text)
        return base64.b64encode(ciphertext)  # Base64编码

    def decrypt(self, text, receiveid):
        """解密密文"""
        cryptor = AES.new(self.key, self.mode, self.key[:16])
        plain_text = cryptor.decrypt(base64.b64decode(text))  # Base64解码+AES解密
        
        pad = plain_text[-1]  # 去除填充
        content = plain_text[16:-pad]  # 去除随机字符串
        
        # 解析长度信息和内容
        xml_len = socket.ntohl(struct.unpack("I", content[:4])[0])
        xml_content = content[4: xml_len + 4]
        from_receiveid = content[xml_len + 4:]
        
        return xml_content  # 返回原始消息

八 Websocket和 Http应用与深度解析

8.1. 协议基础概念

8.1.1 HTTP (HyperText Transfer Protocol)

定义：基于请求-响应模型的应用层协议，主要用于Web浏览器和服务器之间的通信。

1.2 WebSocket

定义：全双工通信协议，在单个TCP连接上提供双向通信通道。

8.2. 核心技术对比

8.2.1 连接建立方式对比

HTTP 连接建立

python

# 典型的HTTP请求响应模式
import requests

# 每次请求都需要建立新连接
response = requests.get('http://example.com/api/data')  # 请求1
response2 = requests.post('http://example.com/api/update', json=data)  # 请求2

WebSocket 连接建立

python

# WebSocket连接建立代码示例
import websockets
import asyncio

async def websocket_client():
    # 一次连接，持续通信
    async with websockets.connect('ws://example.com/ws') as websocket:
        await websocket.send("Hello Server!")  # 发送消息
        response = await websocket.recv()      # 接收消息
        # 可以持续通信...
        await websocket.send("Another message")

8.2.2 协议头对比

HTTP 请求头示例

text

GET /api/data HTTP/1.1
Host: example.com
Connection: keep-alive
User-Agent: Mozilla/5.0
Accept: application/json

WebSocket 升级请求头

text

GET /ws HTTP/1.1
Host: example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: x3JJHMbDL1EzLkh9GBhXDw==
Sec-WebSocket-Version: 13

8.3. 工作模式详细对比

8.3.1 HTTP 工作模式

8.3.2 WebSocket 工作模式

8.4. 代码层面的区别

8.4.1 HTTP 服务器示例 (Python)

python

from http.server import BaseHTTPRequestHandler, HTTPServer
import json

class HTTPHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        # 每个请求创建新的处理器实例
        self.send_response(200)
        self.send_header('Content-type', 'application/json')
        self.end_headers()
        
        response = {'data': 'some data'}
        self.wfile.write(json.dumps(response).encode())
    
    def do_POST(self):
        # 处理POST请求
        content_length = int(self.headers['Content-Length'])
        post_data = self.rfile.read(content_length)
        
        # 处理数据并响应
        self.send_response(200)
        self.end_headers()
        response = {'status': 'success'}
        self.wfile.write(json.dumps(response).encode())

# 启动HTTP服务器
server = HTTPServer(('localhost', 8080), HTTPHandler)
server.serve_forever()

8.4.2 WebSocket 服务器示例 (Python)

python

import asyncio
import websockets
import json

class WebSocketServer:
    def __init__(self):
        self.connected_clients = set()
    
    async def register(self, websocket):
        """注册新客户端"""
        self.connected_clients.add(websocket)
        try:
            async for message in websocket:
                # 处理接收到的消息
                await self.handle_message(websocket, message)
        finally:
            self.connected_clients.remove(websocket)
    
    async def handle_message(self, websocket, message):
        """处理消息并广播"""
        data = json.loads(message)
        # 处理逻辑...
        response = {'type': 'response', 'data': 'processed'}
        
        # 发送回原客户端
        await websocket.send(json.dumps(response))
        
        # 也可以广播给所有客户端
        # await self.broadcast(json.dumps(response))
    
    async def broadcast(self, message):
        """广播消息给所有客户端"""
        if self.connected_clients:
            await asyncio.wait([
                client.send(message) for client in self.connected_clients
            ])

# 启动WebSocket服务器
async def main():
    server = WebSocketServer()
    async with websockets.serve(server.register, "localhost", 8765):
        await asyncio.Future()  # 永久运行

asyncio.run(main())

8.5. 性能特征对比

8.5.1 连接开销对比

特性	HTTP	WebSocket
连接建立	每次请求都需要	一次建立，多次使用
头部开销	每个请求都有完整头部	建立后头部很小
TCP连接	可能频繁开关	持久连接

8.5.2 数据传输效率

python

# HTTP 的大量头部开销
http_overhead = """
Headers: 200-800 bytes per request
+ Cookies: 100-4000 bytes
+ Other metadata
"""

# WebSocket 的轻量级通信
websocket_efficiency = """
Initial handshake: ~600 bytes
Subsequent frames: 2-14 bytes header
+ Payload data only
"""

8.6. 适用场景分析

8.6.1 HTTP 最佳适用场景

python

# 适合HTTP的场景
scenarios_http = [
    "传统的网页浏览",
    "RESTful API接口",
    "文件下载上传",
    "偶尔的数据请求",
    "无状态的服务交互"
]

8.6.2 WebSocket 最佳适用场景

python

# 适合WebSocket的场景
scenarios_websocket = [
    "实时聊天应用",
    "多人协作编辑",
    "实时游戏",
    "股票行情推送",
    "实时监控仪表盘",
    "在线通知系统"
]

8.7. 在给定代码中的应用

8.7.1 HTTP 的使用（httpserver.py）

python

# 企业微信消息接收 - 使用HTTP
class RobotRequestHandler(BaseHTTPRequestHandler):
    def do_POST(self):
        # 处理企业微信的POST请求
        length = int(self.headers['Content-Length'])
        post_data = self.rfile.read(length)
        
        # 解密和处理消息
        self.processor.process(post_data, args)
        self.send_response(200)

8.7.2 WebSocket 的使用（webServer.py）

python

# 任务分发中心 - 使用WebSocket
class WebSocketTaskRouter:
    async def handler(self, ws, path=None):
        client_type = await ws.recv()  # 接收客户端类型
        
        # 持续监听消息
        async for message in ws:
            data = json.loads(message)
            # 根据任务类型路由到不同处理器
            await self.route_task(data, ws)

8.8. 混合架构的优势

8.8.1 为什么选择混合模式

python

# 在给定系统中同时使用两种协议的原因
architecture_reasons = {
    "http_for_compatibility": "企业微信回调使用标准HTTP",
    "websocket_for_efficiency": "内部任务分发需要高效实时通信",
    "right_tool_right_job": "每个协议做最擅长的事情",
    "scalability": "WebSocket减少连接开销，支持更多并发"
}

8.8.2 协议转换桥梁

python

# HTTP到WebSocket的桥梁代码（Flask_server.py）
@app.route("/download", methods=["POST"])
def download():
    # 接收HTTP请求
    req = request.get_json()
    ds_number = req.get("ds_number")
    
    # 转换为WebSocket任务
    task_id = str(uuid.uuid4())
    ws_data = {
        "task_type": "download",
        "task_id": task_id,
        "args_items": ["download", ds_number]
    }
    
    # 通过WebSocket发送并等待结果
    result = ws_client.send_and_wait(json.dumps(ws_data), task_id)
    
    # 将结果返回HTTP响应
    return send_file(BytesIO(file_bytes), as_attachment=True)

8.9. 总结对比表格

特性	HTTP	WebSocket
通信模式	请求-响应	全双工双向
连接持久性	无状态，短连接	有状态，长连接
服务器推送	不支持（需要轮询）	原生支持
头部开销	每个请求都有完整头部	初始握手后开销极小
适用场景	传统Web应用，API	实时应用，聊天，游戏
协议升级	不需要	需要HTTP升级握手
浏览器支持	所有浏览器	现代浏览器
消息顺序	不保证	保证消息顺序