perf(pymsg): send/recv 零拷贝优化
send: sendmsg 散射-聚集 I/O 替代 header+payload 拼接,内核直接 从两块内存读取,省掉大消息的用户态拷贝。 recv: recv_into 写入预分配 bytearray 替代 bytearray 逐步 extend, 省掉 realloc 和 bytes() 转换。
This commit is contained in:
1 parent
267d572906
commit
ae92a6edde
1 file changed
+36
-11
+36
-11
@@ -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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in new issue
Block a user