diff --git a/captcha/captcha_service.py b/captcha/captcha_service.py index 6eb7953..91e1236 100644 --- a/captcha/captcha_service.py +++ b/captcha/captcha_service.py @@ -97,10 +97,10 @@ def generate_image(text): class CaptchaServer: """基于 pymsg 的验证码服务，通过 TCP 长连接收发 JSON。""" - def __init__(self, addr="127.0.0.1", port=8009): + def __init__(self, addr="127.0.0.1", port=8009, workers=4): self.addr = addr self.port = port - self._server = MsgServer(addr, port) + self._server = MsgServer(addr, port, workers=workers) # 路由表 self._routes = { @@ -209,9 +209,11 @@ if __name__ == '__main__': parser = argparse.ArgumentParser(description="captcha — pymsg 验证码服务") parser.add_argument("--addr", default="127.0.0.1", help="监听地址") parser.add_argument("--port", type=int, default=8009, help="监听端口") + parser.add_argument("--workers", type=int, default=4, + help="回调并发线程数（0=同步串行模式）") args = parser.parse_args() - server = CaptchaServer(args.addr, args.port) + server = CaptchaServer(args.addr, args.port, workers=args.workers) try: server.start() except KeyboardInterrupt: diff --git a/db/dbv2.py b/db/dbv2.py index e940015..23d25cd 100644 --- a/db/dbv2.py +++ b/db/dbv2.py @@ -29,6 +29,7 @@ dbv2 — 基于 pymsg 协议（MsgServer）的数据库服务。 import json import os +import threading import time import traceback from datetime import datetime, date, time as dt_time, timezone, timedelta @@ -76,16 +77,27 @@ def _json_default(o): class DBServer: """基于 pymsg 的数据库服务，通过 TCP 长连接收发 JSON。""" - def __init__(self, addr="127.0.0.1", port=8005): + def __init__(self, addr="127.0.0.1", port=8005, workers=4): self.addr = addr self.port = port self.engine = create_engine(DATABASE_URL) self.metadata = MetaData() - self._server = MsgServer(addr, port) + self._server = MsgServer(addr, port, workers=workers) # 日志表缓存（避免重复反射） self._logs_table_cache = {} + # 业务表缓存：仅缓存"已完整反射"的表。 + # 注意不能用 metadata.tables 作完成标记 —— Table(autoload_with=...) + # 会先注册进 metadata.tables（此时无列）再填充列，锁外快路径可能 + # 读到反射半途的空列表。 + self._tables_cache = {} + + # 表反射锁：workers>0 时多个 worker 线程可能对同一新表并发做 + # check-then-act 反射（Table(..., autoload_with=...) 对已定义表名会 + # 抛 InvalidRequestError），反射操作需串行化 + self._reflect_lock = threading.Lock() + # 路由表 self._routes = { "execute": self._handle_execute, @@ -136,9 +148,20 @@ class DBServer: # ── 表反射 --------------------------------------------------------- def _get_table(self, table_name): - if table_name not in self.metadata.tables: - Table(table_name, self.metadata, autoload_with=self.engine) - return self.metadata.tables[table_name] + # double-check：_tables_cache 只在反射完整完成后写入（dict 读写 GIL + # 下安全），锁外快路径读到的要么是完整表要么是 None + table = self._tables_cache.get(table_name) + if table is None: + with self._reflect_lock: + table = self._tables_cache.get(table_name) + if table is None: + # 持锁期间无其他反射在进行，metadata 中的表必然完整 + table = self.metadata.tables.get(table_name) + if table is None: + table = Table(table_name, self.metadata, + autoload_with=self.engine) + self._tables_cache[table_name] = table + return table # ── 日期/时间类型适配 ------------------------------------------------ def _coerce_value(self, column, value): @@ -316,35 +339,43 @@ class DBServer: if "Logs" in self._logs_table_cache: return self._logs_table_cache["Logs"] - with self.engine.connect() as conn: - conn.exec_driver_sql(""" - CREATE TABLE IF NOT EXISTS Logs ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - ts TEXT NOT NULL, - service TEXT NOT NULL, - level TEXT NOT NULL, - message TEXT NOT NULL, - trace_id TEXT, - extra TEXT - ) - """) - conn.exec_driver_sql( - "CREATE INDEX IF NOT EXISTS idx_logs_service_ts ON Logs(service, ts)" - ) - conn.exec_driver_sql( - "CREATE INDEX IF NOT EXISTS idx_logs_level_ts ON Logs(level, ts)" - ) - conn.exec_driver_sql( - "CREATE INDEX IF NOT EXISTS idx_logs_ts ON Logs(ts)" - ) - conn.exec_driver_sql( - "CREATE INDEX IF NOT EXISTS idx_logs_trace ON Logs(trace_id)" - ) - conn.commit() + # 与 _get_table 共用反射锁，double-check 保证并发下只建/反射一次； + # Logs 可能已被 _get_table("Logs") 反射进 metadata，需幂等处理 + with self._reflect_lock: + if "Logs" in self._logs_table_cache: + return self._logs_table_cache["Logs"] - table = Table("Logs", self.metadata, autoload_with=self.engine) - self._logs_table_cache["Logs"] = table - return table + table = self.metadata.tables.get("Logs") + if table is None: + with self.engine.connect() as conn: + conn.exec_driver_sql(""" + CREATE TABLE IF NOT EXISTS Logs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + ts TEXT NOT NULL, + service TEXT NOT NULL, + level TEXT NOT NULL, + message TEXT NOT NULL, + trace_id TEXT, + extra TEXT + ) + """) + conn.exec_driver_sql( + "CREATE INDEX IF NOT EXISTS idx_logs_service_ts ON Logs(service, ts)" + ) + conn.exec_driver_sql( + "CREATE INDEX IF NOT EXISTS idx_logs_level_ts ON Logs(level, ts)" + ) + conn.exec_driver_sql( + "CREATE INDEX IF NOT EXISTS idx_logs_ts ON Logs(ts)" + ) + conn.exec_driver_sql( + "CREATE INDEX IF NOT EXISTS idx_logs_trace ON Logs(trace_id)" + ) + conn.commit() + table = Table("Logs", self.metadata, autoload_with=self.engine) + + self._logs_table_cache["Logs"] = table + return table # ── 内部日志（直接写库，无 HTTP）────────────────────────────────── def _log_internal(self, service, level, message, @@ -432,9 +463,11 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="dbv2 — pymsg 数据库服务") parser.add_argument("--addr", default="127.0.0.1", help="监听地址") parser.add_argument("--port", type=int, default=8005, help="监听端口") + parser.add_argument("--workers", type=int, default=4, + help="回调并发线程数（0=同步串行模式）") args = parser.parse_args() - server = DBServer(args.addr, args.port) + server = DBServer(args.addr, args.port, workers=args.workers) try: server.start() except KeyboardInterrupt: diff --git a/pymsg/README.md b/pymsg/README.md index b6cd64a..3c77548 100644 --- a/pymsg/README.md +++ b/pymsg/README.md @@ -141,6 +141,7 @@ server.start(on_message) # 阻塞当前线程 |------|------| | `addr` | 监听地址 | | `port` | 监听端口 | +| `workers` | 工作线程数（关键字参数，默认 `0`）；`>0` 时回调在线程池中并发执行 | 回调签名为 `on_message(sock, addr, text)`： - `sock` — 客户端 socket，可用 `reply()` 回复 @@ -159,10 +160,22 @@ server.start(on_message) # 阻塞当前线程 服务端内部管理客户端列表，透明处理断连和管道破裂，使用 `select` 实现非阻塞 I/O，轮询间隔 100ms。 +**并发模式（v0.2.0）**： + +```python +server = MsgServer("0.0.0.0", 8005, workers=4) +``` + +默认 `workers=0` 为同步模式：回调在 select 事件循环线程内串行执行，慢回调会阻塞所有连接。`workers>0` 时回调提交到 `ThreadPoolExecutor` 并发执行，I/O 密集回调（如 SQL 查询）可获得真实并行度。保证： + +- **同一连接消息严格串行**：回调完成前该连接暂停 select 监听（恢复粒度 ≤100ms，对短连接 RPC 无影响）； +- **发送不交错**：`reply()` / `broadcast()` 经每连接锁串行化； +- **故障隔离**：回调抛异常仅移除该连接并打印 traceback，服务整体存活（同步模式下异常会直接使 `start()` 崩溃，两种模式语义不同，请注意）。 + ## 线程安全 - `MsgServer.stop()` 可从任意线程安全调用：设置原子标志并关闭监听 socket，`select` 循环会干净退出。 -- `MsgServer.broadcast()` 和 `reply()` 在 `on_message` 回调内（同一线程）调用是安全的。服务端运行期间从其他线程并发调用不支持。 +- `MsgServer.broadcast()` 和 `reply()` 在 `on_message` 回调内调用是安全的（线程池模式下经每连接发送锁串行化）。`broadcast()` 会遍历客户端列表，仅限回调线程调用，不要从无关的第三方线程并发调用。 - `MsgClient.stop()` 可从另一线程调用以中断 `listen()`。 - `send()` 和 `recv()` 内部不做同步 —— 如果多线程共享同一 socket，调用方需自行串行化。 diff --git a/pymsg/README.robot.md b/pymsg/README.robot.md index d94d743..8a8d79b 100644 --- a/pymsg/README.robot.md +++ b/pymsg/README.robot.md @@ -1,7 +1,7 @@ # pymsg 模块技术手册 > 文档版本：2026-07-30 -> 代码版本：0.1.0 +> 代码版本：0.2.0 > 语言：Python 3.9+（核心）、C++（gateway 集成）、Go（waf 集成） > 定位：跨语言 TCP 消息帧协议库 @@ -144,7 +144,9 @@ log_client.push(entry) ``` pymsg/ ├── pymsg/ -│ └── __init__.py # 核心实现（~370 行） +│ └── __init__.py # 核心实现（~460 行） +├── tests/ +│ └── test_threadpool.py # 线程池/同步模式单元测试（纯 stdlib） ├── README.md # 用户文档 ├── pyproject.toml # 项目配置 └── .gitignore @@ -165,23 +167,27 @@ pymsg/ | 函数/类 | 行号 | 功能 | |---------|------|------| -| `send(sock, data, *, timeout=None)` | L48 | 底层发送 | -| `recv(sock, *, timeout=None)` | L72 | 底层接收 | -| `_recv_exact(sock, n)` | L103 | 精确读 n 字节 | -| `_timeout(sock, seconds)` | L34 | 超时 context manager | -| `TimeoutError` | L30 | 超时异常 | -| `MsgClient` | L120 | TCP 客户端 | -| `MsgClient.__init__` | L136 | 构造器 | -| `MsgClient.connect()` | L144 | 建立长连接 | -| `MsgClient.send(msg)` | L150 | 短连接发送（带 name 前缀） | -| `MsgClient.call(msg, *, timeout=None)` | L163 | 短连接 RPC | -| `MsgClient.listen(on_message)` | L193 | 长连接监听 | -| `MsgClient.stop()` | L220 | 停止监听 | -| `MsgServer` | L234 | TCP 服务端 | -| `MsgServer.start(on_message)` | L259 | 启动服务 | -| `MsgServer.broadcast(msg, except_sock)` | L296 | 广播 | -| `MsgServer.reply(sock, msg)` | L321 | 回复 | -| `MsgServer.stop()` | L329 | 停止服务 | +| `send(sock, data, *, timeout=None)` | L51 | 底层发送 | +| `recv(sock, *, timeout=None)` | L75 | 底层接收 | +| `_recv_exact(sock, n)` | L106 | 精确读 n 字节 | +| `_timeout(sock, seconds)` | L38 | 超时 context manager | +| `TimeoutError` | L33 | 超时异常 | +| `MsgClient` | L123 | TCP 客户端 | +| `MsgClient.__init__` | L139 | 构造器 | +| `MsgClient.connect()` | L147 | 建立长连接 | +| `MsgClient.send(msg)` | L153 | 短连接发送（带 name 前缀） | +| `MsgClient.call(msg, *, timeout=None)` | L166 | 短连接 RPC | +| `MsgClient.listen(on_message)` | L196 | 长连接监听 | +| `MsgClient.stop()` | L223 | 停止监听 | +| `MsgServer` | L237 | TCP 服务端 | +| `MsgServer.__init__(addr, port, *, workers=0)` | L262 | 构造器（workers>0 启用线程池） | +| `MsgServer.start(on_message)` | L276 | 启动服务 | +| `MsgServer.broadcast(msg, except_sock)` | L329 | 广播（每连接锁串行化） | +| `MsgServer.reply(sock, msg)` | L355 | 回复（每连接锁串行化） | +| `MsgServer._locked_send(sock, msg)` | L363 | 带每连接锁的发送 | +| `MsgServer.stop()` | L373 | 停止服务 | +| `MsgServer._run_task(...)` | L413 | 工作线程：执行回调、失败打标记 | +| `MsgServer._reap_failed(socketlist)` | L434 | 主循环统一移除失败连接 | #### C++ 实现（Msg.hpp） @@ -319,12 +325,38 @@ client.listen(on_message) # 阻塞，直到 stop() #### 构造 ```python -server = MsgServer(addr: str, port: int) +server = MsgServer(addr: str, port: int, *, workers: int = 0) ``` **参数**： - `addr`：监听地址（`0.0.0.0` 表示所有网卡） - `port`：监听端口 +- `workers`：工作线程数（v0.2.0+）。默认 `0` 为同步模式（回调在 select 循环线程内串行执行）；`>0` 时回调提交到 `ThreadPoolExecutor` 并发执行 + +#### 并发模型（v0.2.0） + +``` +客户端 A ──┐ +客户端 B ──┼──► select 主循环（单线程） +客户端 C ──┘ │ + ├── accept / recv（主线程，快速 I/O） + │ + ├── workers=0：on_message 就地同步执行（串行） + │ + └── workers>0：标记 busy → submit 到 ThreadPoolExecutor + │ + ▼ + worker 线程执行 on_message + ├── reply() 经该连接发送锁串行化 + ├── 异常 → 标记 _failed（主循环下轮统一移除） + └── 完成 → 清除 busy，连接恢复 select 监听 +``` + +保证： +- **同连接串行**：busy 标记使回调完成前该连接暂停 select 监听，同一连接的消息不会并发执行（恢复粒度 ≤100ms，短连接 RPC 无感） +- **发送原子性**：`reply()`/`broadcast()` 经每连接 `Lock` 串行化，多 worker 向同一 socket 写入不会帧交错 +- **故障隔离**：回调抛非连接异常时打印 traceback 并移除该连接，服务整体存活（同步模式下异常会使 `start()` 崩溃，两模式语义不同） +- **fd 安全**：worker 线程不直接 close socket，连接移除统一由主循环执行，避免与 accept/select 产生 fd 复用竞态 #### 方法 @@ -530,11 +562,11 @@ func pymsgCall(payload []byte) ([]byte, error) { | 场景 | 安全性 | 说明 | |------|--------|------| | `MsgServer.stop()` 跨线程调用 | ✅ 安全 | 设置原子标志并关闭监听 socket | -| `MsgServer.reply()` 在回调内调用 | ✅ 安全 | 同一线程内调用 | -| `MsgServer.broadcast()` 在回调内调用 | ✅ 安全 | 同一线程内调用 | +| `MsgServer.reply()` 在回调内调用 | ✅ 安全 | 经每连接发送锁串行化，worker 线程内同样安全 | +| `MsgServer.broadcast()` 在回调内调用 | ✅ 安全 | 发送经每连接锁；但会遍历客户端列表，仅限回调线程，勿从无关第三方线程并发调用 | | `MsgClient.stop()` 跨线程调用 | ✅ 安全 | 可中断 `listen()` | -| `send()`/`recv()` 多线程共享 socket | ⚠️ 不安全 | 调用方需自行串行化 | -| `MsgServer` 并发调用 broadcast/reply | ⚠️ 不支持 | 仅在 `on_message` 回调内（同一线程）安全 | +| `send()`/`recv()` 多线程共享 socket | ⚠️ 不安全 | 底层函数无锁，调用方需自行串行化（MsgServer 内部已用每连接锁保护） | +| 同一连接消息并发执行回调 | ✅ 不会 | busy 标记暂停 select 监听，保证同连接串行 | --- @@ -702,7 +734,8 @@ python test_runner.py | 版本 | 日期 | 主要变更 | |------|------|---------| -| 0.1.0 | 当前 | 初始版本：核心协议、MsgClient、MsgServer、跨语言实现 | +| 0.1.0 | — | 初始版本：核心协议、MsgClient、MsgServer、跨语言实现 | +| 0.2.0 | 2026-07-30 | `MsgServer` 新增 `workers` 参数：回调可提交线程池并发执行；同连接串行（busy 标记）、每连接发送锁防帧交错、回调异常故障隔离；`tests/test_threadpool.py` 单元测试 | --- diff --git a/pymsg/pymsg/__init__.py b/pymsg/pymsg/__init__.py index 4a604ef..4f0a706 100644 --- a/pymsg/pymsg/__init__.py +++ b/pymsg/pymsg/__init__.py @@ -21,6 +21,9 @@ import contextlib import select import socket as _socket import struct +import threading +import traceback +from concurrent.futures import ThreadPoolExecutor __all__ = ["send", "recv", "TimeoutError", "MsgClient", "MsgServer"] @@ -246,14 +249,28 @@ class MsgServer: server.start(on_msg) # 阻塞当前线程 # 在另一线程中调用 server.stop() 可优雅退出 + + 并发模式:: + + server = MsgServer("0.0.0.0", 8003, workers=4) + + ``workers > 0`` 时回调提交到线程池并发执行（I/O 密集回调可获得真实 + 并行度）。同一连接上的消息仍严格串行处理；``reply()`` / ``broadcast()`` + 经每连接发送锁串行化，不会发生帧交错。 """ - def __init__(self, addr: str, port: int) -> None: + def __init__(self, addr: str, port: int, *, workers: int = 0) -> None: self.addr = addr self.port = port + self._workers = workers self._running = False self._server_sock: "_socket.socket | None" = None self._clients: "dict" = {} # socket -> (ip, port) + self._send_locks: "dict" = {} # socket -> Lock（防并发发送帧交错） + self._pool: "ThreadPoolExecutor | None" = None + self._lock = threading.Lock() # 保护 _busy / _failed + self._busy: "set" = set() # 回调执行中的 socket（暂停 select 监听） + self._failed: "set" = set() # 回调失败的 socket（主循环统一移除） # ---- 启动（阻塞）---------------------------------------------------- def start(self, on_message: "callable") -> None: @@ -264,6 +281,10 @@ class MsgServer: - sock: 发送消息的客户端套接字 - addr: (ip, port) 元组 - text: 解码后的字符串消息 + + ``workers > 0`` 时回调提交到线程池并发执行；同一连接的消息保持串行 + （回调完成前该连接暂停 select 监听，恢复粒度为一个 select 周期， + 最坏约 100ms——对短连接 RPC 模式无影响）。 """ self._running = True servsock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) @@ -271,26 +292,38 @@ class MsgServer: servsock.bind((self.addr, self.port)) servsock.listen(10) self._server_sock = servsock + if self._workers > 0: + self._pool = ThreadPoolExecutor( + max_workers=self._workers, + thread_name_prefix=f"pymsg-{self.port}", + ) socketlist = [servsock] - while self._running: - try: - readable, _, exception = select.select( - socketlist, [], socketlist, 0.1 - ) - except (ValueError, OSError): - break + try: + while self._running: + self._reap_failed(socketlist) + with self._lock: + watched = [s for s in socketlist if s not in self._busy] + try: + readable, _, exception = select.select( + watched, [], watched, 0.1 + ) + except (ValueError, OSError): + break - for s in exception: - self._remove_client(s, socketlist) + for s in exception: + self._remove_client(s, socketlist) - for s in readable: - if s is servsock: - self._accept_client(servsock, socketlist) - else: - self._handle_client(s, socketlist, on_message) - - self._shutdown(servsock, socketlist) + for s in readable: + if s is servsock: + self._accept_client(servsock, socketlist) + else: + self._handle_client(s, socketlist, on_message) + finally: + if self._pool is not None: + self._pool.shutdown(wait=True) + self._pool = None + self._shutdown(servsock, socketlist) # ---- 广播 ----------------------------------------------------------- def broadcast( @@ -307,11 +340,12 @@ class MsgServer: if client is except_sock: continue try: - send(client, msg) + self._locked_send(client, msg) except (BrokenPipeError, ConnectionResetError, OSError): disconnected.append(client) for client in disconnected: self._clients.pop(client, None) + self._send_locks.pop(client, None) try: client.close() except OSError: @@ -322,8 +356,18 @@ class MsgServer: """向指定客户端套接字发送一条应答消息（pymsg 协议）。 用于请求-应答模式：服务器收到消息后，通过同一 socket 回复处理结果。 + 经每连接发送锁串行化，``workers > 0`` 时可在回调线程安全调用。 """ - send(sock, msg) + self._locked_send(sock, msg) + + def _locked_send(self, sock: "_socket.socket", msg: str) -> None: + """带每连接锁的发送：防止多线程向同一 socket 并发写入导致帧交错。""" + lock = self._send_locks.get(sock) + if lock is None: + send(sock, msg) + return + with lock: + send(sock, msg) # ---- 停止 ----------------------------------------------------------- def stop(self) -> None: @@ -342,18 +386,61 @@ class MsgServer: if client_sock.fileno() != -1: socketlist.append(client_sock) self._clients[client_sock] = client_addr + self._send_locks[client_sock] = threading.Lock() except OSError: pass def _handle_client(self, s, socketlist, on_message): try: text = recv(s) - on_message(s, self._clients.get(s, ("unknown", 0)), text) except (ConnectionError, ConnectionResetError, OSError): self._remove_client(s, socketlist) + return + addr = self._clients.get(s, ("unknown", 0)) + if self._pool is None: + # 同步模式（默认）：回调在主循环线程内执行，行为与旧版一致 + try: + on_message(s, addr, text) + except (ConnectionError, ConnectionResetError, OSError): + self._remove_client(s, socketlist) + return + # 线程池模式：先标记 busy（暂停 select 监听该连接，保证同连接串行）， + # 再把回调提交到线程池执行 + with self._lock: + self._busy.add(s) + self._pool.submit(self._run_task, s, addr, text, on_message) + + def _run_task(self, s, addr, text, on_message): + """工作线程：执行用户回调。 + + 出错时只打标记（``_failed``），实际的连接关闭与移除由主循环统一 + 执行，避免 worker 线程与 accept/select 并发操作 fd。 + """ + try: + on_message(s, addr, text) + except (ConnectionError, ConnectionResetError, OSError): + with self._lock: + self._failed.add(s) + except Exception: + # 同步模式下此类异常会直接使 start() 崩溃；线程池模式无法传播， + # 打印 traceback 并移除该连接，保证服务整体可用 + traceback.print_exc() + with self._lock: + self._failed.add(s) + finally: + with self._lock: + self._busy.discard(s) + + def _reap_failed(self, socketlist): + """主循环每轮统一移除回调失败的连接。""" + with self._lock: + failed, self._failed = self._failed, set() + for s in failed: + self._remove_client(s, socketlist) def _remove_client(self, s, socketlist): self._clients.pop(s, None) + self._send_locks.pop(s, None) if s in socketlist: socketlist.remove(s) try: @@ -368,3 +455,4 @@ class MsgServer: except OSError: pass self._clients.clear() + self._send_locks.clear() diff --git a/pymsg/pyproject.toml b/pymsg/pyproject.toml index 4479156..3bbfeb0 100644 --- a/pymsg/pyproject.toml +++ b/pymsg/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pymsg" -version = "0.1.0" +version = "0.2.0" description = "Length-prefixed variable-length socket message framing" requires-python = ">=3.9" license = "MIT" diff --git a/pymsg/tests/test_threadpool.py b/pymsg/tests/test_threadpool.py new file mode 100644 index 0000000..a96e813 --- /dev/null +++ b/pymsg/tests/test_threadpool.py @@ -0,0 +1,164 @@ +"""MsgServer 线程池模式（workers > 0）与默认同步模式的单元测试。 + +纯标准库实现，直接运行： + + cd pymsg + python -m unittest discover -s tests -v +""" + +import socket +import threading +import time +import unittest +import unittest.mock + +from pymsg import MsgClient, MsgServer, recv, send + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +class _ServerFixture: + """在后台线程运行 MsgServer；handler 签名 handler(server, sock, addr, text)。""" + + def __init__(self, handler, workers=0): + self.port = _free_port() + self.server = MsgServer("127.0.0.1", self.port, workers=workers) + self.thread = threading.Thread( + target=self.server.start, + args=(lambda sock, addr, text: handler(self.server, sock, addr, text),), + daemon=True, + ) + + def __enter__(self): + self.thread.start() + deadline = time.time() + 5 + while time.time() < deadline: + try: + with socket.create_connection(("127.0.0.1", self.port), timeout=0.2): + return self + except OSError: + time.sleep(0.02) + raise RuntimeError("server did not start in time") + + def __exit__(self, *exc): + self.server.stop() + self.thread.join(timeout=5) + return False + + +def _parallel_calls(port, messages, timeout=10.0): + """多线程并发 call，按 messages 顺序返回响应列表。""" + results = [None] * len(messages) + + def worker(i, msg): + results[i] = MsgClient("127.0.0.1", port, f"t{i}").call(msg, timeout=timeout) + + threads = [threading.Thread(target=worker, args=(i, m)) for i, m in enumerate(messages)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=timeout + 5) + return results + + +class TestSyncMode(unittest.TestCase): + """workers=0（默认）：行为与旧版一致。""" + + def test_rpc_echo(self): + def handler(server, sock, addr, text): + server.reply(sock, f"echo:{text}") + + with _ServerFixture(handler, workers=0) as fx: + client = MsgClient("127.0.0.1", fx.port, "t") + self.assertEqual(client.call("hello", timeout=5), "echo:hello") + + def test_serial_baseline(self): + """同步模式处理慢回调是串行的：6 × 0.4s 应超过 2.0s。""" + def handler(server, sock, addr, text): + time.sleep(0.4) + server.reply(sock, "ok") + + with _ServerFixture(handler, workers=0) as fx: + t0 = time.monotonic() + results = _parallel_calls(fx.port, ["m"] * 6) + elapsed = time.monotonic() - t0 + self.assertEqual(results, ["ok"] * 6) + self.assertGreater(elapsed, 2.0, f"同步模式应为串行，实际 {elapsed:.2f}s") + + +class TestThreadPoolMode(unittest.TestCase): + """workers>0：回调并发执行，同一连接保持串行。""" + + def test_concurrent_calls(self): + """6 个并发 call、回调各 sleep 0.4s、4 个 worker：约 2 批 ≈ 0.8s。""" + def handler(server, sock, addr, text): + time.sleep(0.4) + server.reply(sock, f"ok:{text}") + + with _ServerFixture(handler, workers=4) as fx: + t0 = time.monotonic() + results = _parallel_calls(fx.port, [f"m{i}" for i in range(6)]) + elapsed = time.monotonic() - t0 + self.assertEqual(results, [f"ok:m{i}" for i in range(6)]) + self.assertLess(elapsed, 1.8, f"线程池模式应并发，实际 {elapsed:.2f}s") + + def test_same_connection_serialized(self): + """同一连接连发两条消息，回调执行区间不得重叠。""" + intervals = [] + lock = threading.Lock() + + def handler(server, sock, addr, text): + start = time.monotonic() + time.sleep(0.2) + end = time.monotonic() + with lock: + intervals.append((start, end)) + server.reply(sock, f"r:{text}") + + with _ServerFixture(handler, workers=4) as fx: + with socket.create_connection(("127.0.0.1", fx.port), timeout=5) as s: + send(s, "first") + send(s, "second") + r1 = recv(s, timeout=5) + r2 = recv(s, timeout=5) + + self.assertEqual((r1, r2), ("r:first", "r:second")) + self.assertEqual(len(intervals), 2) + (s1, e1), (s2, e2) = sorted(intervals) + self.assertLessEqual(e1, s2, "同一连接的两次回调发生了并发执行") + + def test_large_frames_not_interleaved(self): + """并发回复大载荷（100KB）时帧不得交错损坏。""" + payload = "x" * 100_000 + + def handler(server, sock, addr, text): + time.sleep(0.01) # 放大交错窗口 + server.reply(sock, f"{text}|{payload}") + + with _ServerFixture(handler, workers=4) as fx: + messages = [f"client{i}" for i in range(8)] + results = _parallel_calls(fx.port, messages) + self.assertEqual(results, [f"client{i}|{payload}" for i in range(8)]) + + def test_callback_exception_keeps_server_alive(self): + """回调抛非连接异常：打印 traceback、移除该连接，服务整体存活。""" + def handler(server, sock, addr, text): + if text == "boom": + raise ValueError("bad message") + server.reply(sock, "alive") + + with _ServerFixture(handler, workers=2) as fx: + with unittest.mock.patch("traceback.print_exc"): + with self.assertRaises(ConnectionError): + MsgClient("127.0.0.1", fx.port, "t").call("boom", timeout=5) + # 后续请求不受影响 + resp = MsgClient("127.0.0.1", fx.port, "t").call("ping", timeout=5) + self.assertEqual(resp, "alive") + + +if __name__ == "__main__": + unittest.main() diff --git a/smtp/smtp_v2.py b/smtp/smtp_v2.py index de9edde..463d84a 100644 --- a/smtp/smtp_v2.py +++ b/smtp/smtp_v2.py @@ -93,10 +93,10 @@ STORE_LOCK = threading.Lock() class SmtpServer: """基于 pymsg 的邮箱验证码服务，通过 TCP 长连接收发 JSON。""" - def __init__(self, addr="127.0.0.1", port=8008): + def __init__(self, addr="127.0.0.1", port=8008, workers=4): self.addr = addr self.port = port - self._server = MsgServer(addr, port) + self._server = MsgServer(addr, port, workers=workers) # 路由表 self._routes = { @@ -182,6 +182,9 @@ class SmtpServer: logs.info("smtp", f"{email}: 发送冷却中，剩余 {wait}s，已拒绝") return {"sent": False, "retry_after": wait, "error": f"发送过于频繁，请 {wait} 秒后重试"} + # 立即占位：防止并发请求同时通过冷却检查导致同邮箱重复发送； + # 发送失败时在下方回滚，"失败不计入冷却"语义不变 + send_history[email] = time.time() code = "".join(str(random.randint(0, 9)) for _ in range(CODE_LEN)) subject = str(req.get("subject") or DEFAULT_SUBJECT) @@ -211,10 +214,11 @@ class SmtpServer: # 发送失败：作废刚写入的验证码，避免"码已存但邮件没到"的不可用状态 with STORE_LOCK: code_store.pop(email, None) + send_history.pop(email, None) # 回滚冷却占位，允许立即重试 logs.error("smtp", f"{email}: 验证码发送失败({err})") return {"sent": False, "error": f"验证码发送失败：{err}"} - # 发送成功：记录冷却起点，并惰性清理已失效的冷却记录 + # 发送成功：刷新冷却起点（检查阶段已占位），并惰性清理已失效的冷却记录 with STORE_LOCK: send_history[email] = time.time() stale = [k for k, v in send_history.items() @@ -280,9 +284,11 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="smtp — pymsg 邮箱验证码服务") parser.add_argument("--addr", default="127.0.0.1", help="监听地址") parser.add_argument("--port", type=int, default=8008, help="监听端口") + parser.add_argument("--workers", type=int, default=4, + help="回调并发线程数（0=同步串行模式）") args = parser.parse_args() - server = SmtpServer(args.addr, args.port) + server = SmtpServer(args.addr, args.port, workers=args.workers) try: server.start() except KeyboardInterrupt: