qing/blackbird
3
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

重构db

Browse Source
This commit is contained in:
qing committed 2026-07-30 20:15:02 +08:00
1 parent 74610c32c8
commit 32d3a5a19f
8 files changed
+427 -88
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

captcha/captcha_service.py
+5 -3
View File
@@ -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:
db/dbv2.py
+67 -34
View File
@@ -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:
pymsg/README.md
+14 -1
View File
@@ -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，调用方需自行串行化。
pymsg/README.robot.md
+58 -25
View File
@@ -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` 单元测试 |
---
pymsg/pymsg/__init__.py
+108 -20
View File
@@ -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()
pymsg/pyproject.toml
+1 -1
View File
@@ -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"
pymsg/tests/test_threadpool.py
+164
View File
@@ -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()
smtp/smtp_v2.py
+10 -4
View File
@@ -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: