From ae92a6edde300aae19a8bd9806016c06f75e86c6 Mon Sep 17 00:00:00 2001
From: qing <1776177662@qq.com>
Date: Sat, 1 Aug 2026 02:08:03 +0800
Subject: [PATCH] =?UTF-8?q?perf(pymsg):=20send/recv=20=E9=9B=B6=E6=8B=B7?=
 =?UTF-8?q?=E8=B4=9D=E4=BC=98=E5=8C=96?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

send: sendmsg 散射-聚集 I/O 替代 header+payload 拼接，内核直接
从两块内存读取，省掉大消息的用户态拷贝。

recv: recv_into 写入预分配 bytearray 替代 bytearray 逐步 extend，
省掉 realloc 和 bytes() 转换。
---
 pymsg/pymsg/__init__.py | 47 +++++++++++++++++++++++++++++++----------
 1 file changed, 36 insertions(+), 11 deletions(-)

diff --git a/pymsg/pymsg/__init__.py b/pymsg/pymsg/__init__.py
index 76e5573..72dd332 100644
--- a/pymsg/pymsg/__init__.py
+++ b/pymsg/pymsg/__init__.py
@@ -73,7 +73,7 @@ def send(
     header = _HEADER.pack(len(payload))
     with _timeout(sock, timeout):
         try:
-            sock.sendall(header + payload)
+            _send_all(sock, header, payload)
         except _socket.timeout:
             raise TimeoutError(
                 f"send timed out after {timeout}s"
@@ -115,17 +115,42 @@ def recv(
             ) from None
 
 
-def _recv_exact(sock: "_socket.socket", n: int) -> bytes:
-    """Read exactly n bytes from sock, raising on premature close."""
-    buf = bytearray()
-    while len(buf) < n:
-        chunk = sock.recv(n - len(buf))
-        if not chunk:
-            raise ConnectionError(
-                f"connection closed: expected {n} bytes, got {len(buf)}"
+def _send_all(sock: "_socket.socket", header: bytes, payload: bytes) -> None:
+    """通过 sendmsg 散射-聚集 I/O 发送 header + payload，避免拼接拷贝。
+
+    内核直接从两块独立内存读取数据，无需在用户态拼接 header + payload。
+    处理 partial send：sendmsg 可能只发出部分数据，需循环补发。
+    """
+    header_len = len(header)
+    total = header_len + len(payload)
+    sent = sock.sendmsg([memoryview(header), memoryview(payload)])
+    while sent < total:
+        if sent < header_len:
+            sent += sock.sendmsg(
+                [memoryview(header)[sent:], memoryview(payload)]
             )
-        buf.extend(chunk)
-    return bytes(buf)
+        else:
+            offset = sent - header_len
+            sent += sock.sendmsg([memoryview(payload)[offset:]])
+
+
+def _recv_exact(sock: "_socket.socket", n: int) -> bytearray:
+    """Read exactly n bytes from sock, raising on premature close.
+
+    使用 recv_into 直接写入预分配 buffer，避免 bytearray 逐步增长和
+    最终 bytes() 转换的额外拷贝。
+    """
+    buf = bytearray(n)
+    view = memoryview(buf)
+    got = 0
+    while got < n:
+        chunk = sock.recv_into(view[got:], n - got)
+        if chunk == 0:
+            raise ConnectionError(
+                f"connection closed: expected {n} bytes, got {got}"
+            )
+        got += chunk
+    return buf
 
 
 # ---------------------------------------------------------------------------