重构db

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()